| """ |
| 2026.6.7 |
| 2026.6.9 |
| 5.5.0 |
| 1.7.0 |
| __UNSLOTH_VERSIONING__ |
| """ |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| from torch import Tensor |
| import torch |
| import torch.nn as nn |
| from torch.nn import functional as F |
| from unsloth_zoo.temporary_patches.common import torch_compile |
| from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable |
| from trl.trainer.rloo_trainer import (Any, AutoModelForSequenceClassification, AutoProcessor, AutoTokenizer, Dataset, DistributedBackend, GenerationConfig, IterableDataset, LoraConfig, Path, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, RLOOConfig, RLOOTrainer, RepeatSampler, RewardFunc, Sampler, SyncRefModelCallback, TrainerCallback, VLLMGeneration, Version, _BaseTrainer, apply_chat_template, asyncio, atexit, copy, create_model_from_path, defaultdict, deque, disable_dropout_in_model, disable_gradient_checkpointing, entropy_from_logits, gather, gather_object, get_config_model_id, get_peft_model, identity, inspect, is_conversational, is_peft_available, is_peft_model, is_rich_available, logger, math, nanmax, nanmin, nanstd, nn, np, pad, pd, peft, prepare_deepspeed, prepare_fsdp, prepare_multimodal_messages, print_prompt_completions_sample, profiling_context, profiling_decorator, selective_log_softmax, set_seed, shuffle_sequence_dict, shutdown_event_loop_in_daemon, split_pixel_values_by_grid, split_tensor_dict, start_event_loop_in_daemon, textwrap, time, torch, transformers, unsplit_pixel_values_by_grid, unwrap_model_for_generation, use_adapter, AutoModelForSequenceClassification, AutoProcessor, AutoTokenizer, Dataset, DistributedBackend, GenerationConfig, IterableDataset, LoraConfig, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, RLOOConfig, RLOOTrainer, RewardFunc, SyncRefModelCallback, TrainerCallback, VLLMGeneration, Version, atexit, copy, create_model_from_path, defaultdict, deque, disable_dropout_in_model, gather, get_config_model_id, get_peft_model, identity, inspect, is_peft_available, is_peft_model, logger, nn, np, pad, pd, peft, prepare_deepspeed, prepare_fsdp, set_seed, shutdown_event_loop_in_daemon, start_event_loop_in_daemon, time, torch, transformers, Any, np, profiling_decorator, shuffle_sequence_dict, split_pixel_values_by_grid, split_tensor_dict, torch, unsplit_pixel_values_by_grid, PeftModel, PreTrainedModel, is_peft_available, logger, peft, torch) |
|
|
|
|
| import os |
| import math |
| import logging |
| 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 |
| import inspect |
| from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling |
| from transformers.training_args import ParallelMode |
| from unsloth_zoo.device_type import DEVICE_TYPE, device_synchronize |
|
|
| |
| import functools |
| from types import MethodType |
| try: |
| from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers |
| except: |
| def reset_unsloth_gradient_checkpointing_buffers(): pass |
| |
| |
| try: |
| from unsloth.models._utils import _unsloth_reset_stray_compile_cache |
| except Exception: |
| def _unsloth_reset_stray_compile_cache(self): pass |
| def prepare_for_training_mode(f): |
| @functools.wraps(f) |
| def wrapper(self, *args, **kwargs): |
| |
| try: |
| _unsloth_reset_stray_compile_cache(self) |
| except Exception: |
| pass |
| |
| |
| |
| |
| |
| if getattr(self, '_unsloth_training_completed', False): |
| try: |
| import wandb |
| if wandb.run is not None: |
| wandb.finish() |
| |
| for cb in self.callback_handler.callbacks: |
| if type(cb).__name__ == 'WandbCallback': |
| cb._initialized = False |
| break |
| except: |
| pass |
| |
| _was_training = None |
| |
| use_gc = getattr(self.args, 'gradient_checkpointing', True) |
| if hasattr(self, 'model') and hasattr(self.model, "training"): |
| _was_training = self.model.training |
| if hasattr(self, 'model') and hasattr(self.model, "for_training"): |
| self.model.for_training(use_gradient_checkpointing=use_gc) |
| output = f(self, *args, **kwargs) |
| |
| if hasattr(self, 'model') and hasattr(self.model, "for_inference"): |
| if _was_training is False: |
| self.model.for_inference() |
| elif _was_training is True and hasattr(self.model, "for_training"): |
| self.model.for_training(use_gradient_checkpointing=use_gc) |
| |
| try: |
| reset_unsloth_gradient_checkpointing_buffers() |
| except: |
| pass |
| |
| |
| self._unsloth_training_completed = True |
| 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_hidden_states_selective_log_softmax( |
| hidden_states: torch.Tensor, |
| lm_head: torch.Tensor, |
| index: torch.Tensor, |
| chunks: int = 4, |
| logit_scale_multiply: float = 0.0, |
| logit_scale_divide: float = 0.0, |
| logit_softcapping: float = 0.0, |
| temperature: float = 1.0, |
| ) -> torch.Tensor: |
| |
| flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) |
| flat_index = index.reshape(-1) |
|
|
| chunked_hidden_states = torch.chunk(flat_hidden_states, chunks=chunks, dim=0) |
| chunked_index = torch.chunk(flat_index, chunks=chunks, dim=0) |
|
|
| all_per_token_logps = [] |
|
|
| for chunk_hidden_states, chunk_index in zip(chunked_hidden_states, chunked_index): |
| chunk_logits = chunk_hidden_states.to(lm_head.dtype) @ lm_head.t() |
|
|
| if logit_scale_multiply != 0.0: |
| chunk_logits = chunk_logits * logit_scale_multiply |
| if logit_scale_divide != 0.0: |
| chunk_logits = chunk_logits / logit_scale_divide |
| if logit_softcapping != 0.0: |
| chunk_logits = logit_softcapping * torch.tanh(chunk_logits / logit_softcapping) |
|
|
| chunk_logits = chunk_logits.to(torch.float32) |
|
|
| if temperature != 1.0: |
| chunk_logits = chunk_logits / temperature |
|
|
| 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) |
|
|
| all_per_token_logps = torch.concat(all_per_token_logps) |
|
|
| all_per_token_logps = all_per_token_logps.reshape((hidden_states.shape[0], hidden_states.shape[1])) |
| return all_per_token_logps |
|
|
| @torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options,) |
| def chunked_selective_log_softmax( |
| logits, |
| index, |
| temperature: float = 1.0, |
| chunks: int = 4, |
| ): |
| chunked_logits = torch.chunk(logits.reshape(-1, logits.shape[-1]), chunks = chunks, dim = 0) |
| chunked_index = torch.chunk(index.reshape(-1), chunks = chunks, dim = 0) |
| all_per_token_logps = [] |
| |
| for chunk_logits, chunk_index in zip(chunked_logits, chunked_index): |
| chunk_logits = chunk_logits.to(torch.float32) |
| if temperature != 1.0: |
| chunk_logits = chunk_logits / temperature |
| 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: |
| """Count left-padded tokens per sequence, e.g. [pad, pad, pad, cat] -> 3.""" |
| 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: |
| """Build a completion mask that zeros leading prompt and trailing pad tokens. |
| |
| For [p,p,p,c,c,c,pad,pad,pad] (p=sliced prompt, c=completion, pad=padding) |
| this returns [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: |
| """Move all padding tokens in each sequence 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 |
|
|
| def align_logprobs_with_mask( |
| logprob_tensor: torch.Tensor, |
| attention_mask: torch.Tensor, |
| pad_value: float = 0.0 |
| ) -> torch.Tensor: |
| """Align a log probability tensor with a given attention mask.""" |
|
|
| device = logprob_tensor.device |
| batch_size, logprob_seq_len = logprob_tensor.shape |
| mask_seq_len = attention_mask.shape[1] |
|
|
| padded_logprobs = torch.full( |
| attention_mask.shape, |
| fill_value=pad_value, |
| dtype=logprob_tensor.dtype, |
| device=device |
| ) |
|
|
| left_pad_counts = torch.argmax(attention_mask, dim=1) |
|
|
| cols = torch.arange(logprob_seq_len, device=device) |
| dest_indices = left_pad_counts.unsqueeze(1) + cols |
|
|
| |
| row_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand_as(dest_indices) |
|
|
| |
| valid_mask = dest_indices < mask_seq_len |
| valid_rows = row_indices[valid_mask] |
| valid_cols = dest_indices[valid_mask] |
| valid_vals = logprob_tensor[valid_mask] |
| padded_logprobs[valid_rows, valid_cols] = valid_vals |
|
|
| return padded_logprobs |
|
|
| def align_completion_tool_mask( |
| tool_mask: torch.Tensor, |
| completion_mask: torch.Tensor, |
| ) -> torch.Tensor: |
| """Align a raw completion-length tool/env mask with Unsloth's repacked loss mask.""" |
| if tool_mask is None: |
| return completion_mask |
| if tool_mask.shape[0] != completion_mask.shape[0]: |
| raise ValueError("tool_mask batch size must match completion_mask batch size.") |
|
|
| tool_mask = tool_mask.to(device=completion_mask.device) |
| if tool_mask.shape == completion_mask.shape: |
| aligned_tool_mask = tool_mask |
| else: |
| aligned_tool_mask = align_logprobs_with_mask( |
| tool_mask, |
| completion_mask, |
| pad_value=0, |
| ) |
| return completion_mask * aligned_tool_mask.to(dtype=completion_mask.dtype) |
|
|
| def autotune_batch_and_chunks( |
| total_input_rows, |
| seq_len, |
| hidden_size, |
| vocab_size, |
| dtype_bytes=16, |
| multiplier=None |
| ): |
| if multiplier is None: |
| final_m = max(4, seq_len // 4096) |
| else: |
| final_m = multiplier |
|
|
| if torch.cuda.is_available(): |
| free_bytes, _ = torch.cuda.mem_get_info() |
| limit_gb = (free_bytes / (1024**3))*.80 |
| elif hasattr(torch, "xpu") and torch.xpu.is_available(): |
| |
| total_mem = torch.xpu.get_device_properties(0).total_memory |
| reserved_mem = torch.xpu.memory_reserved() |
| free_bytes = total_mem - reserved_mem |
| limit_gb = (free_bytes / (1024**3)) * 0.80 |
| else: |
| |
| limit_gb = 8.0 |
|
|
| bytes_to_gb = 1024**3 |
|
|
| b_vals = torch.arange(total_input_rows, 0, -1, device='cpu', dtype=torch.float32) |
|
|
| hidden_gb = (b_vals * seq_len * hidden_size * dtype_bytes) / bytes_to_gb |
|
|
| base_logits = ((b_vals/total_input_rows) * b_vals * seq_len * vocab_size * dtype_bytes) / bytes_to_gb |
| logits_gb = base_logits / final_m |
|
|
| total_mem_gb = hidden_gb + logits_gb |
|
|
| valid_mask = total_mem_gb <= limit_gb |
| valid_indices = torch.nonzero(valid_mask, as_tuple=False) |
|
|
| if valid_indices.shape[0] == 0: |
| |
| return 4, final_m |
|
|
| best_idx = valid_indices[0].item() |
| final_b = int(b_vals[best_idx].item()) |
|
|
| return final_b, final_m |
|
|
| def sanitize_logprob(logprob): |
| """Local port of trl.scripts.vllm_serve.sanitize_logprob. |
| Filters NaN logprobs from vLLM outputs.""" |
| value = logprob.logprob |
| if math.isnan(value): |
| logging.getLogger(__name__).warning( |
| f"Generated NaN logprob, token logprob '{logprob}' will be ignored" |
| ) |
| return None |
| return value |
| @dataclass |
| class UnslothRLOOConfig(RLOOConfig): |
| """ |
| |
| Configuration class for the [`RLOOTrainer`]. |
| |
| This class includes only the parameters that are specific to RLOO training. For a full list of training arguments, |
| please refer to the [`~transformers.TrainingArguments`] 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: |
| > Parameters that control the model and reference model |
| |
| model_init_kwargs (`str`, `dict[str, Any]`, *optional*): |
| Keyword arguments for [`~transformers.AutoModelForCausalLM.from_pretrained`], used when the `model` |
| argument of the [`RLOOTrainer`] is provided as a string. |
| trust_remote_code (`bool`, *optional*, defaults to `False`): |
| Whether to allow loading models and tokenizers that ship custom Python code from the Hub. Forwarded to |
| [`~transformers.AutoModelForCausalLM.from_pretrained`] and |
| [`~transformers.AutoProcessor.from_pretrained`]. Also applied to reward-model and reward-tokenizer loads. |
| router_aux_loss_coef (`float`, *optional*, defaults to `0.001`): |
| Coefficient of the load-balancing auxiliary loss. Only has an effect when training a Mixture-of-Experts |
| (MoE) model; for other models it does nothing. The auxiliary loss is added to the training loss with this |
| weight. Set to `0.0` to disable it. |
| disable_dropout (`bool`, *optional*, defaults to `False`): |
| Whether to disable dropout in the model. This is useful for training with a reference model, as it prevents |
| the model from generating different logprobs for the same input. |
| |
| > Parameters that control the data preprocessing |
| |
| remove_unused_columns (`bool`, *optional*, defaults to `False`): |
| Whether to only keep the column `"prompt"` in the dataset. If you use a custom reward function that |
| requires any column other than `"prompts"` and `"completions"`, you should keep this to `False`. |
| num_generations (`int`, *optional*, defaults to `2`): |
| Number of generations per prompt to sample. The effective batch size (num_processes * per_device_batch_size |
| * gradient_accumulation_steps) must be evenly divisible by this value. |
| num_generations_eval (`int` or `None`, *optional*): |
| Number of generations to sample during evaluation. This allows using fewer generations during evaluation to |
| save computation. If `None`, uses the value of `num_generations`. |
| max_completion_length (`int` or `None`, *optional*, defaults to `256`): |
| Maximum length of the generated completion. |
| 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. Disabling this option is not compatible |
| with vLLM generation. |
| shuffle_dataset (`bool`, *optional*, defaults to `True`): |
| Whether to shuffle the training dataset. |
| pad_to_multiple_of (`int`, *optional*): |
| If set, the prompts ids and completions ids will be padded to a multiple of this value. |
| |
| > Parameters that control generation |
| |
| generation_batch_size (`int`, *optional*): |
| Batch size to use for generation. If `None`, it defaults to the effective training batch size: |
| `per_device_train_batch_size * num_processes * steps_per_generation`. In other words, there is one |
| generation batch processed per optimization step. Mutually exclusive with `steps_per_generation`. |
| steps_per_generation (`int`, *optional*): |
| Number of steps per generation. If `None`, it defaults to `gradient_accumulation_steps`. Mutually exclusive |
| with `generation_batch_size`. |
| temperature (`float`, defaults to `1.0`): |
| Temperature for sampling. The higher the temperature, the more random the completions. |
| top_p (`float`, *optional*, defaults to `1.0`): |
| Float that controls the cumulative probability of the top tokens to consider. Must be in (0, 1]. Set to |
| `1.0` to consider all tokens. |
| top_k (`int`, *optional*, defaults to `0`): |
| Number of highest probability vocabulary tokens to keep for top-k-filtering. If `0`, top-k-filtering is |
| disabled and all tokens are considered. |
| min_p (`float`, *optional*): |
| Minimum token probability, which will be scaled by the probability of the most likely token. It must be a |
| value between `0.0` and `1.0`. Typical values are in the `0.01-0.2` range. |
| generation_kwargs (`dict[str, Any]`, *optional*): |
| Additional keyword arguments to pass to [`~transformers.GenerationConfig`] (if using transformers) or |
| `SamplingParams` (if using vLLM) when sampling completions. This can be used to further customize the |
| generation behavior, such as setting `suppress_tokens`, `num_beams`, etc. If it contains keys that conflict |
| with the other generation parameters (like `min_p`, `top_p`, etc.), they will override them. |
| chat_template_kwargs (`dict[str, Any]`, *optional*): |
| Additional keyword arguments to pass to the `apply_chat_template` function when generating completions. |
| repetition_penalty (`float`, *optional*, defaults to `1.0`): |
| Float that penalizes new tokens based on whether they appear in the prompt and the generated text so far. |
| Values > `1.0` encourage the model to use new tokens, while values < `1.0` encourage the model to repeat |
| tokens. |
| cache_implementation (`str`, *optional*): |
| Implementation of the cache method for faster generation when `use_vllm` is set to `False`. |
| |
| > Parameters that control generation acceleration powered by vLLM |
| |
| use_vllm (`bool`, *optional*, defaults to `False`): |
| Whether to use vLLM for generating completions. If set to `True`, the trainer will use vLLM for generation |
| instead of the default model.generate(). Requires `vllm` to be installed. |
| vllm_mode (`str`, *optional*, defaults to `"colocate"`): |
| Mode to use for vLLM integration when `use_vllm` is set to `True`. Must be one of `"server"` or |
| `"colocate"`. |
| |
| - `"server"`: The trainer will send generation requests to a separate vLLM server. Make sure a TRL vLLM |
| server is running (start with `trl vllm-serve`). |
| - `"colocate"`: vLLM will run in the same process and share the training GPUs. This avoids the need for a |
| separate server but may cause resource contention with training. |
| vllm_model_impl (`str`, *optional*, defaults to `"vllm"`): |
| Model implementation to use for vLLM. Must be one of `"transformers"` or `"vllm"`. `"transformers"`: Use |
| the `transformers` backend for model implementation. `"vllm"`: Use the `vllm` library for model |
| implementation. |
| vllm_structured_outputs_regex (`str`, *optional*): |
| Regex for vLLM structured outputs. If `None` (default), structured outputs is disabled. |
| |
| > Parameters that control the vLLM server (only used when `vllm_mode` is `"server"`) |
| |
| vllm_server_base_url (`str`, *optional*): |
| Base URL for the vLLM server (e.g., `"http://localhost:8000"`). If provided, `vllm_server_host` and |
| `vllm_server_port` are ignored. |
| vllm_server_host (`str`, *optional*, defaults to `"0.0.0.0"`): |
| Host of the vLLM server to connect to. Ignored if `vllm_server_base_url` is provided. |
| vllm_server_port (`int`, *optional*, defaults to `8000`): |
| Port of the vLLM server to connect to. Ignored if `vllm_server_base_url` is provided. |
| vllm_server_timeout (`float`, *optional*, defaults to `240.0`): |
| Total timeout duration in seconds to wait for the vLLM server to be up. If the server is not up after the |
| timeout, a `ConnectionError` is raised. |
| vllm_group_port (`int`, *optional*, defaults to `51216`): |
| Port number for the weight update group. This is used to communicate with the vLLM server. Unless the port |
| is occupied, there is no need to change it. |
| |
| > Parameters that control colocated vLLM execution (only used when `vllm_mode` is `"colocate"`) |
| |
| vllm_gpu_memory_utilization (`float`, *optional*, defaults to `0.3`): |
| Control the GPU memory utilization for vLLM. This setting only applies when `vllm_mode` is set to |
| `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when |
| launching the vLLM server via the `--vllm_gpu_memory_utilization` flag. |
| vllm_max_model_length (`int`, *optional*): |
| Context window for vLLM. Set it to at least the maximum prompt length in the dataset plus |
| `max_completion_length`; if omitted, it is inferred from the model config. |
| vllm_tensor_parallel_size (`int`, *optional*, defaults to `1`): |
| Control the tensor parallel size for vLLM. This setting only applies when `vllm_mode` is set to |
| `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when |
| launching the vLLM server via the `--vllm_tensor_parallel_size` flag. |
| vllm_enable_sleep_mode (`bool`, *optional*, defaults to `False`): |
| Enable vLLM sleep mode to offload weights/cache during the optimizer step. Keeps GPU memory usage low, but |
| waking the engine adds host–device transfer latency. |
| |
| > Parameters that control generation acceleration powered by transformers continuous batching |
| |
| use_transformers_continuous_batching (`bool`, *optional*, defaults to `False`): |
| Whether to use transformers' continuous batching engine for generating completions. Requires |
| `transformers>=5.8.0`. |
| transformers_continuous_batching_config (`dict`, *optional*): |
| Keyword arguments for [`~transformers.generation.ContinuousBatchingConfig`]. |
| |
| > Parameters that control the training |
| |
| beta (`float`, *optional*, defaults to `0.05`): |
| KL coefficient. If `0.0`, the reference model is not loaded, reducing memory usage and improving training |
| speed. |
| num_iterations (`int`, *optional*, defaults to `1`): |
| Number of iterations per batch (denoted as μ in the algorithm). |
| epsilon (`float`, *optional*, defaults to `0.2`): |
| Epsilon value for clipping. |
| epsilon_high (`float`, *optional*): |
| Upper-bound epsilon value for clipping. If not specified, it defaults to the same value as the lower-bound |
| specified in argument `epsilon`. Paper [DAPO](https://huggingface.co/papers/2503.14476) recommends `0.28`. |
| reward_weights (`list[float]`, *optional*): |
| Weights for each reward function. Must match the number of reward functions. If `None`, all rewards are |
| weighted equally with weight `1.0`. |
| normalize_advantages (`bool`, *optional*, defaults to `False`): |
| Whether to normalize advantages. Normalization is done per generation batch to have mean `0.0` and standard |
| deviation of `1.0`. |
| reward_clip_range (`tuple[float, float]`, *optional*): |
| Clip range for rewards as (min, max). If `None`, no clipping is applied. |
| mask_truncated_completions (`bool`, *optional*, defaults to `False`): |
| When enabled, truncated completions are excluded from the loss calculation, preventing them from being |
| incorrectly penalized and introducing noise during training. According to the |
| [DAPO](https://huggingface.co/papers/2503.14476) paper, this is a good practice for training stability. |
| sync_ref_model (`bool`, *optional*, defaults to `False`): |
| Whether to synchronize the reference model with the active model every `ref_model_sync_steps` steps, using |
| the `ref_model_mixup_alpha` parameter. This synchronization originates from the |
| [TR-DPO](https://huggingface.co/papers/2404.09656) paper. |
| ref_model_mixup_alpha (`float`, *optional*, defaults to `0.6`): |
| α parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which controls the mix |
| between the current policy and the previous reference policy during updates. The reference policy is |
| updated according to the equation: `π_ref = α * π_θ + (1 - α) * π_ref_prev`. To use this parameter, you |
| must set `sync_ref_model=True`. |
| ref_model_sync_steps (`int`, *optional*, defaults to `512`): |
| τ parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which determines how |
| frequently the current policy is synchronized with the reference policy. To use this parameter, you must |
| set `sync_ref_model=True`. |
| |
| > Parameters that control the logging |
| |
| log_completions (`bool`, *optional*, defaults to `False`): |
| Whether to log a sample of (prompt, completion) pairs every `logging_steps` steps. If `rich` is installed, |
| it prints the sample. If `wandb` and/or `trackio` logging is enabled, it logs it to `wandb` and/or |
| `trackio`. |
| num_completions_to_print (`int`, *optional*): |
| Number of completions to print with `rich`. If `None`, all completions are logged. |
| log_unique_prompts (`bool`, *optional*, defaults to `False`): |
| Whether to log unique prompts. If `True`, only unique prompts are logged. If `False`, all prompts are |
| logged. |
| |
| > Deprecated parameters |
| |
| use_transformers_paged: |
| |
| <Deprecated version="1.2.0"> |
| |
| Parameter `use_transformers_paged` is deprecated and will be removed in version v2.0.0. Use |
| `use_transformers_continuous_batching` instead. |
| |
| </Deprecated> |
| |
| > [!NOTE] |
| > These parameters have default values different from [`~transformers.TrainingArguments`]: |
| > - `logging_steps`: Defaults to `10` instead of `500`. |
| > - `gradient_checkpointing`: Defaults to `True` instead of `False`. |
| > - `bf16`: Defaults to `True` if `fp16` is not set, instead of `False`. |
| > - `learning_rate`: Defaults to `1e-6` instead of `5e-5`. |
| |
| """ |
| 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.'}, |
| ) |
| unsloth_logit_chunk_multiplier : Optional[int] = field( |
| default = None, |
| metadata = {'help': 'Multiplier for chunked logit computations.'}, |
| ) |
| unsloth_grpo_mini_batch : Optional[int] = field( |
| default = None, |
| metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'}, |
| ) |
| |
| def __init__( |
| self, |
| output_dir = None, |
| per_device_train_batch_size = 4, |
| num_train_epochs = 3.0, |
| max_steps = -1, |
| learning_rate = 5e-05, |
| lr_scheduler_type = 'linear', |
| lr_scheduler_kwargs = None, |
| warmup_steps = 0.1, |
| optim = 'adamw_8bit', |
| optim_args = None, |
| weight_decay = 0.001, |
| adam_beta1 = 0.9, |
| adam_beta2 = 0.999, |
| adam_epsilon = 1e-08, |
| optim_target_modules = None, |
| gradient_accumulation_steps = 2, |
| average_tokens_across_devices = True, |
| max_grad_norm = 1.0, |
| label_smoothing_factor = 0.0, |
| bf16 = False, |
| fp16 = False, |
| bf16_full_eval = False, |
| fp16_full_eval = False, |
| tf32 = None, |
| gradient_checkpointing = True, |
| gradient_checkpointing_kwargs = None, |
| torch_compile = False, |
| torch_compile_backend = None, |
| torch_compile_mode = None, |
| use_liger_kernel = False, |
| liger_kernel_config = None, |
| use_cache = False, |
| neftune_noise_alpha = None, |
| torch_empty_cache_steps = 250, |
| auto_find_batch_size = False, |
| logging_strategy = 'steps', |
| logging_steps = 1, |
| logging_first_step = False, |
| log_on_each_node = True, |
| logging_nan_inf_filter = False, |
| include_num_input_tokens_seen = False, |
| log_level = 'passive', |
| log_level_replica = 'warning', |
| disable_tqdm = None, |
| report_to = 'none', |
| run_name = None, |
| project = 'huggingface', |
| trackio_space_id = 'trackio', |
| eval_strategy = 'no', |
| eval_steps = None, |
| eval_delay = 0, |
| per_device_eval_batch_size = 4, |
| prediction_loss_only = False, |
| eval_on_start = False, |
| eval_do_concat_batches = True, |
| eval_use_gather_object = False, |
| eval_accumulation_steps = 2, |
| batch_eval_metrics = False, |
| save_only_model = False, |
| save_strategy = 'steps', |
| save_steps = 500, |
| save_on_each_node = False, |
| save_total_limit = None, |
| enable_jit_checkpoint = False, |
| push_to_hub = False, |
| hub_token = None, |
| hub_private_repo = None, |
| hub_model_id = None, |
| hub_strategy = 'every_save', |
| hub_always_push = False, |
| hub_revision = None, |
| load_best_model_at_end = False, |
| metric_for_best_model = None, |
| greater_is_better = None, |
| ignore_data_skip = False, |
| restore_callback_states_from_checkpoint = False, |
| full_determinism = False, |
| seed = 3407, |
| data_seed = 3407, |
| use_cpu = False, |
| accelerator_config = None, |
| parallelism_config = None, |
| dataloader_drop_last = False, |
| dataloader_num_workers = 0, |
| dataloader_pin_memory = True, |
| dataloader_persistent_workers = False, |
| dataloader_prefetch_factor = None, |
| remove_unused_columns = False, |
| label_names = None, |
| train_sampling_strategy = 'random', |
| length_column_name = 'length', |
| ddp_find_unused_parameters = None, |
| ddp_bucket_cap_mb = None, |
| ddp_broadcast_buffers = None, |
| ddp_backend = None, |
| ddp_timeout = 1800, |
| fsdp = None, |
| fsdp_config = None, |
| deepspeed = None, |
| debug = '', |
| skip_memory_metrics = True, |
| do_train = False, |
| do_eval = False, |
| do_predict = False, |
| resume_from_checkpoint = None, |
| warmup_ratio = None, |
| logging_dir = None, |
| local_rank = -1, |
| model_init_kwargs = None, |
| trust_remote_code = False, |
| router_aux_loss_coef = 0.001, |
| disable_dropout = False, |
| num_generations = 8, |
| num_generations_eval = None, |
| max_completion_length = 256, |
| ds3_gather_for_generation = True, |
| shuffle_dataset = True, |
| pad_to_multiple_of = None, |
| generation_batch_size = None, |
| steps_per_generation = None, |
| temperature = 1.0, |
| top_p = 1.0, |
| top_k = None, |
| min_p = None, |
| generation_kwargs = {}, |
| chat_template_kwargs = None, |
| repetition_penalty = 1.0, |
| cache_implementation = None, |
| use_vllm = False, |
| vllm_mode = 'colocate', |
| vllm_model_impl = 'vllm', |
| vllm_enable_sleep_mode = False, |
| vllm_structured_outputs_regex = None, |
| vllm_server_base_url = None, |
| vllm_server_host = '0.0.0.0', |
| vllm_server_port = 8000, |
| vllm_server_timeout = 240.0, |
| vllm_group_port = 51216, |
| vllm_gpu_memory_utilization = 0.3, |
| vllm_max_model_length = None, |
| vllm_tensor_parallel_size = 1, |
| beta = 0.05, |
| num_iterations = 1, |
| epsilon = 0.2, |
| epsilon_high = None, |
| reward_weights = None, |
| normalize_advantages = False, |
| reward_clip_range = None, |
| mask_truncated_completions = False, |
| sync_ref_model = False, |
| ref_model_mixup_alpha = 0.6, |
| ref_model_sync_steps = 512, |
| log_completions = False, |
| num_completions_to_print = None, |
| log_unique_prompts = False, |
| use_transformers_continuous_batching = False, |
| transformers_continuous_batching_config = None, |
| use_transformers_paged = False, |
| vllm_sampling_params = None, |
| unsloth_num_chunks = -1, |
| unsloth_logit_chunk_multiplier = None, |
| unsloth_grpo_mini_batch = None, |
| |
| **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 num_train_epochs is None: |
| num_train_epochs = 3.0 |
| if output_dir is None and save_strategy == 'steps' and save_steps == 500: |
| output_dir = 'unsloth_training_checkpoints' |
| save_strategy = 'no' |
| if os.environ.get('UNSLOTH_ENABLE_FLEX_ATTENTION', '0') == '1': |
| from unsloth_zoo.flex_attention import HAS_FLEX_ATTENTION |
| if HAS_FLEX_ATTENTION and pad_to_multiple_of is None: |
| from unsloth_zoo.flex_attention import FLEX_ATTENTION_BLOCK_SIZE |
| pad_to_multiple_of = FLEX_ATTENTION_BLOCK_SIZE |
| |
| if steps_per_generation is None and generation_batch_size is None: |
| ga = gradient_accumulation_steps |
| world_size = int(os.environ.get('WORLD_SIZE', '1')) |
| if (ga * world_size * per_device_train_batch_size) % num_generations != 0: |
| print('Unsloth: We now expect `per_device_train_batch_size` * `gradient_accumulation_steps` * `world_size` to be a multiple of `num_generations`.\nWe will change the batch size of ' + str(per_device_train_batch_size) + ' to the `num_generations` of ' + str(num_generations)) |
| per_device_train_batch_size = num_generations |
| |
| if temperature <= 0: |
| raise ValueError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.') |
| elif temperature >= 10: |
| raise ValueError('Unsloth: Please set a positive non-zero temperature less than 10, since sampling will be quite erratic.') |
| |
| |
| super().__init__( |
| output_dir = output_dir, |
| per_device_train_batch_size = per_device_train_batch_size, |
| num_train_epochs = num_train_epochs, |
| max_steps = max_steps, |
| learning_rate = learning_rate, |
| lr_scheduler_type = lr_scheduler_type, |
| lr_scheduler_kwargs = lr_scheduler_kwargs, |
| warmup_steps = warmup_steps, |
| optim = optim, |
| optim_args = optim_args, |
| weight_decay = weight_decay, |
| adam_beta1 = adam_beta1, |
| adam_beta2 = adam_beta2, |
| adam_epsilon = adam_epsilon, |
| optim_target_modules = optim_target_modules, |
| gradient_accumulation_steps = gradient_accumulation_steps, |
| average_tokens_across_devices = average_tokens_across_devices, |
| max_grad_norm = max_grad_norm, |
| label_smoothing_factor = label_smoothing_factor, |
| bf16 = bf16, |
| fp16 = fp16, |
| bf16_full_eval = bf16_full_eval, |
| fp16_full_eval = fp16_full_eval, |
| tf32 = tf32, |
| gradient_checkpointing = gradient_checkpointing, |
| gradient_checkpointing_kwargs = gradient_checkpointing_kwargs, |
| torch_compile = torch_compile, |
| torch_compile_backend = torch_compile_backend, |
| torch_compile_mode = torch_compile_mode, |
| use_liger_kernel = use_liger_kernel, |
| liger_kernel_config = liger_kernel_config, |
| use_cache = use_cache, |
| neftune_noise_alpha = neftune_noise_alpha, |
| torch_empty_cache_steps = torch_empty_cache_steps, |
| auto_find_batch_size = auto_find_batch_size, |
| logging_strategy = logging_strategy, |
| logging_steps = logging_steps, |
| logging_first_step = logging_first_step, |
| log_on_each_node = log_on_each_node, |
| logging_nan_inf_filter = logging_nan_inf_filter, |
| include_num_input_tokens_seen = include_num_input_tokens_seen, |
| log_level = log_level, |
| log_level_replica = log_level_replica, |
| disable_tqdm = disable_tqdm, |
| report_to = report_to, |
| run_name = run_name, |
| project = project, |
| trackio_space_id = trackio_space_id, |
| eval_strategy = eval_strategy, |
| eval_steps = eval_steps, |
| eval_delay = eval_delay, |
| per_device_eval_batch_size = per_device_eval_batch_size, |
| prediction_loss_only = prediction_loss_only, |
| eval_on_start = eval_on_start, |
| eval_do_concat_batches = eval_do_concat_batches, |
| eval_use_gather_object = eval_use_gather_object, |
| eval_accumulation_steps = eval_accumulation_steps, |
| batch_eval_metrics = batch_eval_metrics, |
| save_only_model = save_only_model, |
| save_strategy = save_strategy, |
| save_steps = save_steps, |
| save_on_each_node = save_on_each_node, |
| save_total_limit = save_total_limit, |
| enable_jit_checkpoint = enable_jit_checkpoint, |
| push_to_hub = push_to_hub, |
| hub_token = hub_token, |
| hub_private_repo = hub_private_repo, |
| hub_model_id = hub_model_id, |
| hub_strategy = hub_strategy, |
| hub_always_push = hub_always_push, |
| hub_revision = hub_revision, |
| 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, |
| restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint, |
| full_determinism = full_determinism, |
| seed = seed, |
| data_seed = data_seed, |
| use_cpu = use_cpu, |
| accelerator_config = accelerator_config, |
| parallelism_config = parallelism_config, |
| dataloader_drop_last = dataloader_drop_last, |
| dataloader_num_workers = dataloader_num_workers, |
| dataloader_pin_memory = dataloader_pin_memory, |
| dataloader_persistent_workers = dataloader_persistent_workers, |
| dataloader_prefetch_factor = dataloader_prefetch_factor, |
| remove_unused_columns = remove_unused_columns, |
| label_names = label_names, |
| train_sampling_strategy = train_sampling_strategy, |
| length_column_name = length_column_name, |
| ddp_find_unused_parameters = ddp_find_unused_parameters, |
| ddp_bucket_cap_mb = ddp_bucket_cap_mb, |
| ddp_broadcast_buffers = ddp_broadcast_buffers, |
| ddp_backend = ddp_backend, |
| ddp_timeout = ddp_timeout, |
| fsdp = fsdp, |
| fsdp_config = fsdp_config, |
| deepspeed = deepspeed, |
| debug = debug, |
| skip_memory_metrics = skip_memory_metrics, |
| do_train = do_train, |
| do_eval = do_eval, |
| do_predict = do_predict, |
| resume_from_checkpoint = resume_from_checkpoint, |
| warmup_ratio = warmup_ratio, |
| logging_dir = logging_dir, |
| local_rank = local_rank, |
| model_init_kwargs = model_init_kwargs, |
| trust_remote_code = trust_remote_code, |
| router_aux_loss_coef = router_aux_loss_coef, |
| disable_dropout = disable_dropout, |
| num_generations = num_generations, |
| num_generations_eval = num_generations_eval, |
| max_completion_length = max_completion_length, |
| ds3_gather_for_generation = ds3_gather_for_generation, |
| shuffle_dataset = shuffle_dataset, |
| pad_to_multiple_of = pad_to_multiple_of, |
| generation_batch_size = generation_batch_size, |
| steps_per_generation = steps_per_generation, |
| temperature = temperature, |
| top_p = top_p, |
| top_k = top_k, |
| min_p = min_p, |
| generation_kwargs = generation_kwargs, |
| chat_template_kwargs = chat_template_kwargs, |
| repetition_penalty = repetition_penalty, |
| cache_implementation = cache_implementation, |
| use_vllm = use_vllm, |
| vllm_mode = vllm_mode, |
| vllm_model_impl = vllm_model_impl, |
| vllm_enable_sleep_mode = vllm_enable_sleep_mode, |
| vllm_structured_outputs_regex = vllm_structured_outputs_regex, |
| vllm_server_base_url = vllm_server_base_url, |
| vllm_server_host = vllm_server_host, |
| vllm_server_port = vllm_server_port, |
| vllm_server_timeout = vllm_server_timeout, |
| vllm_group_port = vllm_group_port, |
| vllm_gpu_memory_utilization = vllm_gpu_memory_utilization, |
| vllm_max_model_length = vllm_max_model_length, |
| vllm_tensor_parallel_size = vllm_tensor_parallel_size, |
| beta = beta, |
| num_iterations = num_iterations, |
| epsilon = epsilon, |
| epsilon_high = epsilon_high, |
| reward_weights = reward_weights, |
| normalize_advantages = normalize_advantages, |
| reward_clip_range = reward_clip_range, |
| mask_truncated_completions = mask_truncated_completions, |
| sync_ref_model = sync_ref_model, |
| ref_model_mixup_alpha = ref_model_mixup_alpha, |
| ref_model_sync_steps = ref_model_sync_steps, |
| log_completions = log_completions, |
| num_completions_to_print = num_completions_to_print, |
| log_unique_prompts = log_unique_prompts, |
| use_transformers_continuous_batching = use_transformers_continuous_batching, |
| transformers_continuous_batching_config = transformers_continuous_batching_config, |
| use_transformers_paged = use_transformers_paged,**kwargs) |
| self.vllm_sampling_params = vllm_sampling_params |
| self.unsloth_num_chunks = unsloth_num_chunks |
| if unsloth_grpo_mini_batch is not None: |
| if self.generation_batch_size >= unsloth_grpo_mini_batch: |
| self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch |
| else: |
| raise ValueError( |
| f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, " |
| f"which is self.per_device_train_batch_size * gradient_accumulation_steps." |
| ) |
| self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier |
| |
| |
| if getattr(self, 'gradient_checkpointing_kwargs', None) is not None: |
| if 'use_reentrant' in self.gradient_checkpointing_kwargs: |
| del self.gradient_checkpointing_kwargs['use_reentrant'] |
|
|
| pass |
|
|
| class _UnslothRLOOTrainer(_BaseTrainer): |
| """""" |
|
|
| _tag_names = ["trl", "rloo"] |
| _name = "RLOO" |
| _paper = { |
| "title": "Back to Basics: Revisiting REINFORCE-Style Optimization for Learning from Human Feedback in LLMs", |
| "id": "2402.14740", |
| |
| "citation": textwrap.dedent("""\ |
| @inproceedings{ahmadian2024back, |
| title = {{Back to Basics: Revisiting REINFORCE-Style Optimization for Learning from Human Feedback in LLMs}}, |
| author = {Arash Ahmadian and Chris Cremer and Matthias Gall{\'{e}} and Marzieh Fadaee and Julia Kreutzer and Olivier Pietquin and Ahmet {\"{U}}st{\"{u}}n and Sara Hooker}, |
| year = 2024, |
| booktitle = {Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), {ACL} 2024, Bangkok, Thailand, August 11-16, 2024}, |
| pages = {12248--12267}, |
| publisher = {Association for Computational Linguistics}, |
| editor = {Lun{-}Wei Ku and Andre Martins and Vivek Srikumar}, |
| }"""), |
| } |
|
|
| def __init__( |
| self, |
| model: "str | PreTrainedModel | PeftModel", |
| reward_funcs: RewardFunc | list[RewardFunc], |
| args: RLOOConfig | None = None, |
| train_dataset: Dataset | IterableDataset | None = None, |
| eval_dataset: Dataset | IterableDataset | dict[str, Dataset | IterableDataset] | None = None, |
| processing_class: PreTrainedTokenizerBase | ProcessorMixin | None = None, |
| reward_processing_classes: PreTrainedTokenizerBase | list[PreTrainedTokenizerBase] | None = None, |
| callbacks: list[TrainerCallback] | None = None, |
| optimizers: tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None] = (None, None), |
| peft_config: "PeftConfig | None" = None, |
| ): |
|
|
| if hasattr(model, 'vllm_engine') and hasattr(args, 'use_vllm'): |
| if (getattr(args, 'use_vllm', False) == False): |
| args.use_vllm = True |
| |
| if args is None: |
| model_name = model if isinstance(model, str) else get_config_model_id(model.config) |
| model_name = model_name.split("/")[-1] |
| args = RLOOConfig(f"{model_name}-RLOO") |
|
|
| |
| if isinstance(model, str): |
| model_init_kwargs = args.model_init_kwargs or {} |
| |
| if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: |
| model_init_kwargs["device_map"] = None |
| model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) |
| model = create_model_from_path(model, **model_init_kwargs) |
| else: |
| if args.model_init_kwargs is not None: |
| logger.warning( |
| "You passed `model_init_kwargs` to the `RLOOConfig`, but your model is already instantiated. " |
| "The `model_init_kwargs` will be ignored." |
| ) |
| |
| _is_quantized_model = getattr(model, "is_loaded_in_4bit", False) or getattr(model, "is_loaded_in_8bit", False) |
|
|
| |
| |
| self.model_kwarg_keys = ( |
| inspect.signature(model.forward).parameters.keys() |
| if not hasattr(model, "get_base_model") |
| else inspect.signature(model.get_base_model().forward).parameters.keys() |
| ) |
|
|
| |
| if processing_class is None: |
| processing_class = AutoProcessor.from_pretrained( |
| get_config_model_id(model.config), |
| truncation_side="left", |
| padding_side="left", |
| trust_remote_code=args.trust_remote_code, |
| ) |
|
|
| if args.use_transformers_continuous_batching and isinstance(processing_class, ProcessorMixin): |
| raise ValueError( |
| "`use_transformers_continuous_batching` does not support multimodal models. Use `use_vllm` instead." |
| ) |
|
|
| |
| if isinstance(processing_class, ProcessorMixin): |
| self._tokenizer = processing_class.tokenizer |
| elif isinstance(processing_class, PreTrainedTokenizerBase): |
| self._tokenizer = processing_class |
| else: |
| raise TypeError("The `processing_class` must be either a `PreTrainedTokenizerBase` or a `ProcessorMixin`") |
|
|
| if self._tokenizer.pad_token is None: |
| self._tokenizer.pad_token = self._tokenizer.eos_token |
|
|
| |
| if False: |
| if not is_peft_available(): |
| raise ImportError( |
| "You passed `peft_config` but the `peft` library is not installed. " |
| "Install it with `pip install trl[peft]`." |
| ) |
| if not isinstance(peft_config, PeftConfig): |
| raise TypeError( |
| f"`peft_config` must be a `peft.PeftConfig` instance (e.g. `peft.LoraConfig`), " |
| f"got {type(peft_config).__name__}." |
| ) |
| if is_peft_model(model): |
| raise ValueError( |
| "You passed a `PeftModel` instance together with a `peft_config` to the trainer. Please first merge " |
| "and unload the existing adapter, save the resulting base model, and then pass that base model along " |
| "with the new `peft_config` to the trainer." |
| ) |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| get_peft_model_kwargs = {} |
| if ( |
| args.deepspeed_plugin is not None |
| and args.deepspeed_plugin.zero_stage == 3 |
| and not _is_quantized_model |
| and Version(peft.__version__) >= Version("0.12.0") |
| ): |
| get_peft_model_kwargs["autocast_adapter_dtype"] = False |
| model = get_peft_model(model, peft_config, **get_peft_model_kwargs) |
|
|
| elif is_peft_model(model): |
| |
| |
| |
| |
| |
| default_config = model.peft_config["default"] |
| if isinstance(default_config, LoraConfig) and default_config.target_parameters: |
| logger.warning( |
| "PEFT can't add a frozen reference adapter alongside one that uses `target_parameters` " |
| "(peft#3340], so the reference log probs are computed from the base model [adapters disabled]. " |
| "If you wrapped the model only to apply LoRA, pass a `peft_config` to the trainer instead; if you " |
| "wrapped it deliberately (pretrained adapter or custom init), note that the base model matches " |
| "your adapter only when it's freshly zero-initialized. If it is, this warning is safe to ignore." |
| ) |
| else: |
| model.add_adapter("ref", default_config) |
| for name, param in model.named_parameters(): |
| if ".default." in name: |
| ref_name = name.replace(".default.", ".ref.") |
| ref_param = model.get_parameter(ref_name) |
| ref_param.data.copy_(param.data) |
|
|
| |
| |
| if is_peft_model(model) and args.gradient_checkpointing: |
| model.enable_input_require_grads() |
|
|
| |
| |
| |
| |
| if _is_quantized_model: |
| for param in model.parameters(): |
| if param.requires_grad: |
| param.data = param.data.to(torch.bfloat16) |
|
|
| |
| if not isinstance(reward_funcs, list): |
| reward_funcs = [reward_funcs] |
| self.reward_func_names = [] |
| for i, reward_func in enumerate(reward_funcs): |
| if isinstance(reward_func, str): |
| model_init_kwargs = args.model_init_kwargs or {} |
| |
| if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: |
| model_init_kwargs["device_map"] = None |
| model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) |
| reward_funcs[i] = AutoModelForSequenceClassification.from_pretrained( |
| reward_func, num_labels=1, **model_init_kwargs |
| ) |
| if isinstance(reward_funcs[i], nn.Module): |
| self.reward_func_names.append(get_config_model_id(reward_funcs[i].config).split("/")[-1]) |
| else: |
| self.reward_func_names.append(reward_funcs[i].__name__) |
| self.reward_funcs = reward_funcs |
|
|
| self._has_async_funcs = any(inspect.iscoroutinefunction(func) for func in self.reward_funcs) |
| if self._has_async_funcs: |
| self.async_loop_thread, self.async_loop, self.async_loop_ready_event = start_event_loop_in_daemon( |
| name="RLOOTrainer-AsyncRewardLoop" |
| ) |
| |
| self.async_loop_ready_event.wait() |
| atexit.register(shutdown_event_loop_in_daemon, self.async_loop_thread, self.async_loop) |
|
|
| |
| if args.reward_weights is not None: |
| if len(args.reward_weights) != len(reward_funcs): |
| raise ValueError( |
| f"Number of reward weights ({len(args.reward_weights)}) must match number of reward " |
| f"functions ({len(reward_funcs)})" |
| ) |
| self.reward_weights = torch.tensor(args.reward_weights, dtype=torch.float32) |
| else: |
| self.reward_weights = torch.ones(len(reward_funcs), dtype=torch.float32) |
|
|
| |
| if reward_processing_classes is None: |
| reward_processing_classes = [None] * len(reward_funcs) |
| elif not isinstance(reward_processing_classes, list): |
| reward_processing_classes = [reward_processing_classes] |
| if len(reward_processing_classes) != len(reward_funcs): |
| raise ValueError( |
| f"The number of reward processing classes ({len(reward_processing_classes)}) must match the number of " |
| f"reward functions ({len(reward_funcs)})." |
| ) |
|
|
| for i, (reward_processing_class, reward_func) in enumerate( |
| zip(reward_processing_classes, reward_funcs, strict=True) |
| ): |
| if isinstance(reward_func, PreTrainedModel): |
| if reward_processing_class is None: |
| reward_processing_class = AutoTokenizer.from_pretrained( |
| get_config_model_id(reward_func.config), trust_remote_code=args.trust_remote_code |
| ) |
| if reward_processing_class.pad_token_id is None: |
| reward_processing_class.pad_token = reward_processing_class.eos_token |
| |
| |
| reward_func.config.pad_token_id = reward_processing_class.pad_token_id |
| reward_processing_classes[i] = reward_processing_class |
|
|
| self.reward_processing_classes = reward_processing_classes |
|
|
| |
| self.max_completion_length = args.max_completion_length |
| self.num_generations = args.num_generations |
| self.num_generations_eval = args.num_generations_eval or self.num_generations |
| self.chat_template_kwargs = args.chat_template_kwargs or {} |
| self.temperature = args.temperature |
| self.top_p = args.top_p |
| self.top_k = args.top_k |
| self.min_p = args.min_p |
| self.repetition_penalty = args.repetition_penalty |
| self.use_transformers_continuous_batching = args.use_transformers_continuous_batching |
| if self.use_transformers_continuous_batching: |
| if not Version(transformers.__version__) >= Version("5.8.0"): |
| raise ImportError( |
| "Using `use_transformers_continuous_batching` requires transformers>=5.8.0. " |
| "Please upgrade with `pip install --upgrade transformers`." |
| ) |
| from transformers.generation import ContinuousBatchingConfig |
|
|
| cb_kwargs = dict(args.transformers_continuous_batching_config or {}) |
| |
| |
| cb_kwargs.setdefault("max_memory_percent", 0.5) |
| self.continuous_batching_config = ContinuousBatchingConfig(**cb_kwargs) |
| else: |
| self.continuous_batching_config = None |
| self.pad_to_multiple_of = args.pad_to_multiple_of |
| self.use_vllm = args.use_vllm |
| self.vllm_mode = args.vllm_mode |
| self.vllm_gpu_memory_utilization = args.vllm_gpu_memory_utilization |
| self.vllm_tensor_parallel_size = args.vllm_tensor_parallel_size |
| self.normalize_advantages = args.normalize_advantages |
| self.mask_truncated_completions = args.mask_truncated_completions |
| self.reward_clip_range = args.reward_clip_range |
|
|
| |
| self.shuffle_dataset = args.shuffle_dataset |
|
|
| if train_dataset is None: |
| raise ValueError("`train_dataset` is required") |
| elif ( |
| isinstance(train_dataset, IterableDataset) |
| or isinstance(eval_dataset, IterableDataset) |
| or ( |
| isinstance(eval_dataset, dict) and any(isinstance(ds, IterableDataset) for ds in eval_dataset.values()) |
| ) |
| ): |
| |
| raise NotImplementedError( |
| "Iterable datasets are not yet supported in RLOOTrainer. Please use a standard dataset instead." |
| ) |
|
|
| |
| self.num_iterations = args.num_iterations |
| self.epsilon_low = args.epsilon |
| self.epsilon_high = args.epsilon_high if args.epsilon_high is not None else args.epsilon |
|
|
| |
| text_config = model.config.get_text_config() |
| is_moe = getattr(text_config, "output_router_logits", None) is not None |
| self.aux_loss_enabled = is_moe and args.router_aux_loss_coef != 0.0 |
| self.router_aux_loss_coef = args.router_aux_loss_coef |
| |
| self._step = 0 |
| |
| |
| self._buffered_inputs = None |
|
|
| |
| |
| |
| |
| if args.gradient_checkpointing and Version(transformers.__version__) < Version("5.0.0"): |
| args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} |
| args.gradient_checkpointing_kwargs.setdefault("use_reentrant", False) |
|
|
| super().__init__( |
| model=model, |
| args=args, |
| data_collator=identity, |
| train_dataset=train_dataset, |
| eval_dataset=eval_dataset, |
| processing_class=processing_class, |
| callbacks=callbacks, |
| optimizers=optimizers, |
| ) |
|
|
| |
| self.beta = args.beta |
| if self.beta == 0.0: |
| |
| self.ref_model = None |
| elif is_peft_model(model): |
| |
| |
| self.ref_model = None |
| else: |
| |
| model_init_kwargs = args.model_init_kwargs or {} |
| |
| if self.args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: |
| model_init_kwargs["device_map"] = None |
| model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) |
| self.ref_model = create_model_from_path(get_config_model_id(self.model.config), **model_init_kwargs) |
|
|
| |
| if args.disable_dropout: |
| disable_dropout_in_model(model) |
| if self.ref_model is not None: |
| disable_dropout_in_model(self.ref_model) |
|
|
| |
| self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} |
| self._total_train_tokens = 0 |
| self._current_train_step_time = 0.0 |
| self.log_completions = args.log_completions |
| self.log_unique_prompts = args.log_unique_prompts |
| self.num_completions_to_print = args.num_completions_to_print |
| |
| self._logs = { |
| "images": deque(maxlen=args.generation_batch_size), |
| "prompt": deque(maxlen=args.generation_batch_size), |
| "completion": deque(maxlen=args.generation_batch_size), |
| "rewards": defaultdict(lambda: deque(maxlen=args.generation_batch_size)), |
| "advantages": deque(maxlen=args.generation_batch_size), |
| "extra": defaultdict(lambda: deque(maxlen=args.generation_batch_size)), |
| } |
| |
| self._pending_extra_logs = defaultdict(list) |
| self._pending_metrics = defaultdict(list) |
|
|
| |
| |
| |
| set_seed(args.seed, device_specific=True) |
|
|
| if self.use_vllm: |
| self.vllm_generation = VLLMGeneration( |
| model=self.model, |
| accelerator=self.accelerator, |
| processing_class=self.processing_class, |
| mode=args.vllm_mode, |
| structured_outputs_regex=args.vllm_structured_outputs_regex, |
| server_base_url=args.vllm_server_base_url, |
| server_host=args.vllm_server_host, |
| server_port=args.vllm_server_port, |
| group_port=args.vllm_group_port, |
| server_timeout=args.vllm_server_timeout, |
| tensor_parallel_size=args.vllm_tensor_parallel_size, |
| gpu_memory_utilization=args.vllm_gpu_memory_utilization, |
| max_model_length=args.vllm_max_model_length, |
| max_num_seqs=args.per_device_train_batch_size |
| * args.vllm_tensor_parallel_size |
| * args.steps_per_generation, |
| enable_sleep_mode=args.vllm_enable_sleep_mode, |
| model_impl=args.vllm_model_impl, |
| repetition_penalty=self.repetition_penalty, |
| temperature=self.temperature, |
| top_p=self.top_p, |
| top_k=self.top_k, |
| min_p=self.min_p, |
| max_completion_length=self.max_completion_length, |
| logprobs=None, |
| generation_kwargs=args.generation_kwargs, |
| ) |
| self._last_loaded_step = -1 |
| else: |
| generation_kwargs = { |
| "max_new_tokens": self.max_completion_length, |
| "do_sample": True, |
| "pad_token_id": self._tokenizer.pad_token_id, |
| "bos_token_id": self._tokenizer.bos_token_id, |
| "eos_token_id": self._tokenizer.eos_token_id, |
| "temperature": self.temperature, |
| "top_p": self.top_p, |
| "top_k": self.top_k, |
| "min_p": self.min_p, |
| "repetition_penalty": self.repetition_penalty, |
| "cache_implementation": args.cache_implementation, |
| } |
| if args.generation_kwargs is not None: |
| generation_kwargs.update(args.generation_kwargs) |
| self.generation_config = GenerationConfig(**generation_kwargs, disable_compile=True) |
| |
| self.generation_kwargs = generation_kwargs |
|
|
| |
| |
| |
| self.model_accepts_loss_kwargs = False |
| self._dist = DistributedBackend(self.accelerator) |
|
|
| |
| self.model.add_model_tags(self._tag_names) |
|
|
| if self.ref_model is not None: |
| if self.is_deepspeed_enabled: |
| self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator) |
| elif self.is_fsdp_enabled: |
| self.ref_model = prepare_fsdp(self.ref_model, self.accelerator) |
| else: |
| self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True) |
|
|
| if args.sync_ref_model: |
| if self.beta == 0.0: |
| raise ValueError( |
| "You passed `sync_ref_model=True` while `beta=0.0`, which means the reference model is not used " |
| "during training. Consequently, RLOOTrainer does not create a `ref_model` instance, and there is " |
| "nothing to synchronize. Please set `sync_ref_model=False`, or set `beta` to a non-zero value." |
| ) |
| if is_peft_model(model): |
| raise NotImplementedError( |
| "You passed `sync_ref_model=True` while using a PEFT model, which is currently not supported. " |
| "With PEFT, RLOOTrainer does not keep a separate reference model in memory; instead, it recovers " |
| "reference behavior by temporarily disabling the adapter. As a result, there is no standalone " |
| "`ref_model` instance to synchronize. Use `sync_ref_model=False`, or opt for full fine-tuning if " |
| "you need a synced reference model. If you need `sync_ref_model` to work with PEFT, please open a " |
| "feature request at https://github.com/huggingface/trl/issues." |
| ) |
| self.add_callback(SyncRefModelCallback(ref_model=self.ref_model, accelerator=self.accelerator)) |
|
|
| for i, reward_func in enumerate(self.reward_funcs): |
| if isinstance(reward_func, PreTrainedModel): |
| if self.is_deepspeed_enabled: |
| self.reward_funcs[i] = prepare_deepspeed(reward_func, self.accelerator) |
| else: |
| |
| self.reward_funcs[i] = self.accelerator.prepare_model( |
| reward_func, evaluation_mode=True, device_placement=True |
| ) |
|
|
| def _set_signature_columns_if_needed(self): |
| |
| |
| |
| |
| if self._signature_columns is None: |
| self._signature_columns = ["prompt", "image", "images"] |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| def get_train_dataloader(self): |
| return self._get_dataloader( |
| dataset=self.train_dataset, |
| description="Training", |
| batch_size=self._train_batch_size * self.args.steps_per_generation, |
| sampler_fn=self._get_train_sampler, |
| is_training=True, |
| ) |
|
|
| def _get_train_sampler(self, dataset: Dataset | None = None) -> Sampler: |
| |
| |
| |
| |
| |
| |
| |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| if dataset is None: |
| dataset = self.train_dataset |
| return RepeatSampler( |
| data_source=dataset, |
| mini_repeat_count=self.num_generations, |
| batch_size=self.args.generation_batch_size // self.num_generations, |
| repeat_count=self.num_iterations * self.args.steps_per_generation, |
| shuffle=self.shuffle_dataset, |
| seed=self.args.seed, |
| ) |
|
|
| def _get_eval_sampler(self, eval_dataset) -> Sampler: |
| |
| return RepeatSampler( |
| data_source=eval_dataset, |
| mini_repeat_count=self.num_generations_eval, |
| seed=self.args.seed, |
| ) |
|
|
| @profiling_decorator |
| def _get_per_token_logps_and_entropies( |
| self, |
| model, |
| input_ids, |
| attention_mask, |
| logits_to_keep, |
| batch_size=None, |
| compute_entropy=False, |
| compute_aux_loss=False, |
| pixel_values=None, |
| image_grid_thw=None, |
| num_images=None, |
| pixel_attention_mask=None, |
| spatial_shapes=None, |
| num_tiles=None, |
| image_sizes=None, |
| token_type_ids=None, |
| mm_token_type_ids=None, |
| image_position_ids=None, |
| ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: |
| """Compute log-probs, (optionally) entropies, and (optionally) the MoE load-balancing aux loss.""" |
| batch_size = batch_size or input_ids.size(0) |
| all_logps = [] |
| all_entropies = [] |
| all_aux_losses = [] |
| for start in range(0, input_ids.size(0), batch_size): |
| input_ids_batch = input_ids[start : start + batch_size] |
| attention_mask_batch = attention_mask[start : start + batch_size] |
|
|
| |
| model_inputs = {"input_ids": input_ids_batch, "attention_mask": attention_mask_batch} |
| if image_grid_thw is not None and pixel_values is not None: |
| rows_per_image = image_grid_thw.prod(dim=-1) |
| rows_per_sample = torch.split(rows_per_image, num_images) |
| rows_per_sample = torch.stack([s.sum() for s in rows_per_sample]) |
| cum_rows = torch.cat([torch.tensor([0], device=rows_per_sample.device), rows_per_sample.cumsum(0)]) |
| row_start, row_end = cum_rows[start].item(), cum_rows[start + batch_size].item() |
| model_inputs["pixel_values"] = pixel_values[row_start:row_end] |
| cum_imgs = torch.tensor([0] + num_images).cumsum(0) |
| img_start, img_end = cum_imgs[start], cum_imgs[start + batch_size] |
| model_inputs["image_grid_thw"] = image_grid_thw[img_start:img_end] |
| elif image_position_ids is not None and pixel_values is not None: |
| cum_imgs = torch.tensor([0] + num_images).cumsum(0) |
| img_start, img_end = cum_imgs[start], cum_imgs[start + batch_size] |
| model_inputs["pixel_values"] = pixel_values[img_start:img_end] |
| model_inputs["image_position_ids"] = image_position_ids[img_start:img_end] |
| elif spatial_shapes is not None and pixel_values is not None: |
| |
| cum_tiles = torch.tensor([0] + num_tiles).cumsum(0) |
| tile_start, tile_end = cum_tiles[start], cum_tiles[start + batch_size] |
| model_inputs["pixel_values"] = pixel_values[tile_start:tile_end] |
| model_inputs["pixel_attention_mask"] = pixel_attention_mask[tile_start:tile_end] |
| model_inputs["spatial_shapes"] = spatial_shapes[tile_start:tile_end] |
| elif pixel_values is not None: |
| model_inputs["pixel_values"] = pixel_values[start : start + batch_size] |
| if pixel_attention_mask is not None and spatial_shapes is None: |
| model_inputs["pixel_attention_mask"] = pixel_attention_mask[start : start + batch_size] |
| if image_sizes is not None: |
| model_inputs["image_sizes"] = image_sizes[start : start + batch_size] |
| if token_type_ids is not None: |
| model_inputs["token_type_ids"] = token_type_ids[start : start + batch_size] |
| if mm_token_type_ids is not None: |
| model_inputs["mm_token_type_ids"] = mm_token_type_ids[start : start + batch_size] |
|
|
| |
| if "logits_to_keep" in self.model_kwarg_keys: |
| |
| model_inputs["logits_to_keep"] = logits_to_keep + 1 |
|
|
| model_inputs["use_cache"] = False |
|
|
| |
| |
| if compute_aux_loss: |
| model_inputs["output_router_logits"] = True |
|
|
| outputs = model(**model_inputs) |
| logits = outputs.logits |
| |
| logits = logits[:, :-1, :] |
| |
| logits = logits[:, -logits_to_keep:, :] |
| |
| |
| logits.div_(self.temperature) |
| completion_ids = input_ids_batch[:, -logits_to_keep:] |
| logps = selective_log_softmax(logits, completion_ids) |
| all_logps.append(logps) |
|
|
| if compute_entropy: |
| with torch.no_grad(): |
| entropies = entropy_from_logits(logits) |
| all_entropies.append(entropies) |
|
|
| if compute_aux_loss: |
| all_aux_losses.append(outputs.aux_loss) |
|
|
| logps = torch.cat(all_logps, dim=0) |
| entropies = torch.cat(all_entropies, dim=0) if compute_entropy else None |
| aux_loss = torch.stack(all_aux_losses).mean() if compute_aux_loss else None |
| return logps, entropies, aux_loss |
|
|
| def training_step(self, model, inputs, num_items_in_batch): |
| time_before = time.perf_counter() |
| output = super().training_step(model, inputs, num_items_in_batch) |
| self._step += 1 |
| time_after = time.perf_counter() |
| self._current_train_step_time += time_after - time_before |
| if self._step % self.current_gradient_accumulation_steps == 0: |
| self._metrics["train"]["step_time"].append(self._current_train_step_time) |
| self._current_train_step_time = 0.0 |
| return output |
|
|
| @profiling_decorator |
| def _prepare_inputs(self, generation_batch: dict[str, torch.Tensor | Any]) -> dict[str, torch.Tensor | Any]: |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| mode = "train" if self.model.training else "eval" |
| if mode == "train": |
| generate_every = self.args.steps_per_generation * self.num_iterations |
| if self._step % generate_every == 0 or self._buffered_inputs is None: |
| |
| generation_batch = self._generate_and_score_completions(generation_batch) |
| generation_batch = split_pixel_values_by_grid(generation_batch) |
|
|
| try: generation_batch = shuffle_sequence_dict(generation_batch) |
|
|
| except: pass |
| generation_batches = split_tensor_dict(generation_batch, self.args.steps_per_generation) |
| self._buffered_inputs = [unsplit_pixel_values_by_grid(batch) for batch in generation_batches] |
| inputs = self._buffered_inputs[self._step % self.args.steps_per_generation] |
| else: |
| |
| |
| inputs = self._generate_and_score_completions(generation_batch) |
| return inputs |
|
|
| def _log_completion_extra(self, column: str, values: list): |
| """ |
| Log extra columns to the completions table. Called from reward functions via the `log_extra` kwarg. |
| |
| Args: |
| column (`str`): |
| Name of the column to add. |
| values (`list`): |
| Values for the column, one per sample in the batch. |
| """ |
| self._pending_extra_logs[column].extend(values) |
|
|
| def _log_metric(self, name: str, value: float): |
| """ |
| Log a scalar metric from a reward function. Called via the `log_metric` kwarg. Values are averaged over each |
| logging step and reported alongside built-in metrics like `kl` and `entropy`. |
| |
| Args: |
| name (`str`): |
| Name of the metric. |
| value (`float`): |
| Scalar value for this batch. |
| """ |
| self._pending_metrics[name].append(value) |
|
|
| @profiling_decorator |
| def _calculate_rewards(self, inputs, prompts, completions, completion_ids_list): |
| device = self.accelerator.device |
| rewards_per_func = torch.zeros(len(prompts), len(self.reward_funcs), device=device) |
|
|
| |
| keys = [key for key in inputs[0] if key not in ["prompt", "completion", "completion_ids"]] |
| reward_kwargs = {key: [example[key] for example in inputs] for key in keys} |
|
|
| |
| reward_kwargs["trainer_state"] = self.state |
|
|
| |
| reward_kwargs["log_extra"] = self._log_completion_extra |
|
|
| |
| reward_kwargs["log_metric"] = self._log_metric |
|
|
| async_funcs_info = [] |
|
|
| for i, (reward_func, reward_processing_class, reward_func_name) in enumerate( |
| zip(self.reward_funcs, self.reward_processing_classes, self.reward_func_names, strict=True) |
| ): |
| if isinstance(reward_func, nn.Module): |
| with profiling_context(self, reward_func_name): |
| if is_conversational(inputs[0]): |
| messages = [{"messages": p + c} for p, c in zip(prompts, completions, strict=True)] |
| texts = [ |
| apply_chat_template(x, reward_processing_class, **self.chat_template_kwargs)["text"] |
| for x in messages |
| ] |
| else: |
| texts = [p + c for p, c in zip(prompts, completions, strict=True)] |
| reward_inputs = reward_processing_class( |
| text=texts, return_tensors="pt", padding=True, padding_side="right", add_special_tokens=False |
| ) |
| reward_inputs = super()._prepare_inputs(reward_inputs) |
| with torch.inference_mode(): |
| rewards_per_func[:, i] = reward_func(**reward_inputs).logits[:, 0] |
| elif inspect.iscoroutinefunction(reward_func): |
| async_funcs_info.append((i, reward_func, reward_func_name)) |
| else: |
| |
| with profiling_context(self, reward_func_name): |
| output_reward_func = reward_func( |
| prompts=prompts, completions=completions, completion_ids=completion_ids_list, **reward_kwargs |
| ) |
| |
| output_reward_func = [reward if reward is not None else torch.nan for reward in output_reward_func] |
| rewards_per_func[:, i] = torch.tensor(output_reward_func, dtype=torch.float32, device=device) |
|
|
| |
| if async_funcs_info: |
|
|
| async def _invoke_async(index, func, func_name): |
| with profiling_context(self, func_name): |
| output = await func( |
| prompts=prompts, completions=completions, completion_ids=completion_ids_list, **reward_kwargs |
| ) |
| output = [r if r is not None else torch.nan for r in output] |
| return index, output |
|
|
| async def _run_async_funcs(): |
| coros = [_invoke_async(i, func, func_name) for (i, func, func_name) in async_funcs_info] |
| return await asyncio.gather(*coros) |
|
|
| async_results = asyncio.run_coroutine_threadsafe(_run_async_funcs(), self.async_loop).result() |
| for idx, output_reward_func in async_results: |
| rewards_per_func[:, idx] = torch.tensor(output_reward_func, dtype=torch.float32, device=device) |
|
|
| |
| if torch.isnan(rewards_per_func).all(dim=1).any(): |
| nan_row_idx = torch.isnan(rewards_per_func).all(dim=1).nonzero(as_tuple=True)[0][0] |
| row_reward_kwargs = { |
| key: value[nan_row_idx] |
| for key, value in reward_kwargs.items() |
| if key not in ("trainer_state", "log_extra", "log_metric") |
| } |
| row_reward_kwargs["prompt"] = prompts[nan_row_idx] |
| row_reward_kwargs["completion"] = completions[nan_row_idx] |
| logger.warning( |
| f"All reward functions returned None for the following kwargs:\n{row_reward_kwargs}\n" |
| "Please ensure that at least one reward function returns a valid reward." |
| ) |
|
|
| |
| |
| rewards_per_func = gather(rewards_per_func) |
| return rewards_per_func |
|
|
| def _tokenize_prompts(self, prompts: list): |
| """Tokenize prompts and extract images/multimodal fields for generation.""" |
| if is_conversational({"prompt": prompts[0]}): |
| |
| images = [] |
| has_images = False |
| for prompt in prompts: |
| prompt_images = [] |
| for message in prompt: |
| if isinstance(message["content"], list): |
| for part in message["content"]: |
| if part["type"] == "image": |
| prompt_images.append(part["image"]) |
| has_images = True |
| images.append(prompt_images if prompt_images else None) |
| images = images if has_images else None |
|
|
| |
| |
| |
| needs_padding_workaround = Version("5.3.0") <= Version(transformers.__version__) < Version("5.4.0") |
| tokenized = self.processing_class.apply_chat_template( |
| conversation=prompts, |
| add_generation_prompt=True, |
| tokenize=True, |
| return_dict=True, |
| **({"padding": True} if needs_padding_workaround else {}), |
| **self.chat_template_kwargs, |
| ) |
| if needs_padding_workaround: |
| |
| prompt_ids = [ |
| [tok for tok, m in zip(ids, mask, strict=True) if m] |
| for ids, mask in zip(tokenized["input_ids"], tokenized["attention_mask"], strict=True) |
| ] |
| else: |
| prompt_ids = tokenized["input_ids"] |
| |
| multimodal_fields = {k: v for k, v in tokenized.items() if k not in ("input_ids", "attention_mask")} |
| else: |
| prompt_ids = self.processing_class(text=prompts)["input_ids"] |
| images = None |
| multimodal_fields = {} |
| return prompt_ids, images, multimodal_fields |
|
|
| def _generate_single_turn(self, prompt_ids, images, multimodal_fields): |
| device = self.accelerator.device |
| mode = "train" if self.model.training else "eval" |
|
|
| |
| if self.use_vllm: |
| |
| if self.state.global_step != self._last_loaded_step: |
| with profiling_context(self, "sync_weights"): |
| self.vllm_generation.sync_weights() |
| self._last_loaded_step = self.state.global_step |
|
|
| |
| num_generations = self.num_generations if mode == "train" else self.num_generations_eval |
| _, completion_ids, _, _ = self.vllm_generation.generate( |
| prompts=prompt_ids, |
| images=images, |
| num_generations=num_generations, |
| profiler=profiling_context(self, "vLLM.generate"), |
| ) |
|
|
| elif self.use_transformers_continuous_batching: |
| with ( |
| profiling_context(self, "transformers.generate_batch"), |
| unwrap_model_for_generation( |
| self.model_wrapped, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation |
| ) as unwrapped_model, |
| torch.no_grad(), |
| self._dist.summon_full_params(self.model_wrapped, recurse=False), |
| ): |
| |
| if self.args.bf16: |
| unwrapped_model.to(torch.bfloat16) |
| elif self.args.fp16: |
| unwrapped_model.to(torch.float16) |
| all_outputs = unwrapped_model.generate_batch( |
| prompt_ids, |
| generation_config=self.generation_config, |
| continuous_batching_config=self.continuous_batching_config, |
| progress_bar=False, |
| ) |
| unwrapped_model.train() |
| completion_ids = [output.generated_tokens for output in all_outputs.values()] |
|
|
| else: |
| |
| prompt_tensors = [torch.tensor(ids) for ids in prompt_ids] |
| padded_ids = pad(prompt_tensors, padding_value=self._tokenizer.pad_token_id, padding_side="left") |
| attention_mask = pad([torch.ones_like(t) for t in prompt_tensors], padding_value=0, padding_side="left") |
| generate_inputs = {"input_ids": padded_ids, "attention_mask": attention_mask} |
| |
| for k, v in multimodal_fields.items(): |
| if isinstance(v, torch.Tensor): |
| generate_inputs[k] = v |
| elif isinstance(v, list) and v and isinstance(v[0], list): |
| |
| generate_inputs[k] = pad([torch.tensor(x) for x in v], padding_value=0, padding_side="left") |
| else: |
| generate_inputs[k] = torch.tensor(np.array(v)) |
| generate_inputs = super()._prepare_inputs(generate_inputs) |
|
|
| with ( |
| profiling_context(self, "transformers.generate"), |
| unwrap_model_for_generation( |
| self.model_wrapped, |
| self.accelerator, |
| gather_deepspeed3_params=self.args.ds3_gather_for_generation, |
| generation_kwargs=self.generation_kwargs, |
| ) as unwrapped_model, |
| torch.no_grad(), |
| self._dist.summon_full_params(self.model_wrapped, recurse=False), |
| ): |
| prompt_completion_ids = unwrapped_model.generate( |
| **generate_inputs, generation_config=self.generation_config |
| ) |
| |
| prompt_length = generate_inputs["input_ids"].size(1) |
| completion_ids = prompt_completion_ids[:, prompt_length:] |
|
|
| |
| is_eos = completion_ids == self._tokenizer.eos_token_id |
| eos_idx = torch.full((is_eos.size(0),), is_eos.size(1), dtype=torch.long, device=device) |
| eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)] |
| sequence_indices = torch.arange(is_eos.size(1), device=device).expand(is_eos.size(0), -1) |
| completion_mask = (sequence_indices <= eos_idx.unsqueeze(1)).int() |
| completion_ids = [ |
| c[m].tolist() for c, m in zip(completion_ids.cpu(), completion_mask.bool().cpu(), strict=True) |
| ] |
|
|
| return completion_ids |
|
|
| def _generate(self, prompts: list): |
| device = self.accelerator.device |
| mode = "train" if self.model.training else "eval" |
|
|
| |
| prompts = copy.deepcopy(prompts) |
|
|
| prompt_ids, images, multimodal_fields = self._tokenize_prompts(prompts) |
| completion_ids = self._generate_single_turn(prompt_ids, images, multimodal_fields) |
|
|
| |
| if is_conversational({"prompt": prompts[0]}): |
| contents = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) |
| completions = [[{"role": "assistant", "content": content}] for content in contents] |
| else: |
| completions = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) |
|
|
| |
| prompt_lengths = torch.tensor([len(ids) for ids in prompt_ids], device=device) |
| completion_lengths = torch.tensor([len(ids) for ids in completion_ids], device=device) |
| agg_prompt_lengths = self.accelerator.gather(prompt_lengths) |
| agg_completion_lengths = self.accelerator.gather(completion_lengths) |
| total_prompt_tokens = agg_prompt_lengths.sum() |
| total_completion_tokens = agg_completion_lengths.sum() |
|
|
| |
| if mode == "train": |
| self.state.num_input_tokens_seen += (total_prompt_tokens + total_completion_tokens).item() |
| self._metrics[mode]["num_tokens"] = [self.state.num_input_tokens_seen] |
|
|
| |
| self._metrics[mode]["completions/mean_length"].append(agg_completion_lengths.float().mean().item()) |
| self._metrics[mode]["completions/min_length"].append(agg_completion_lengths.float().min().item()) |
| self._metrics[mode]["completions/max_length"].append(agg_completion_lengths.float().max().item()) |
|
|
| |
| eos_and_pad = [self._tokenizer.eos_token_id, self._tokenizer.pad_token_id] |
| is_truncated = torch.tensor([ids[-1] not in eos_and_pad for ids in completion_ids], device=device) |
| agg_is_truncated = self.accelerator.gather(is_truncated) |
| self._metrics[mode]["completions/clipped_ratio"].append(agg_is_truncated.float().mean().item()) |
| term_completion_lengths = agg_completion_lengths[~agg_is_truncated] |
| if len(term_completion_lengths) == 0: |
| term_completion_lengths = torch.zeros(1, device=device) |
| self._metrics[mode]["completions/mean_terminated_length"].append(term_completion_lengths.float().mean().item()) |
| self._metrics[mode]["completions/min_terminated_length"].append(term_completion_lengths.float().min().item()) |
| self._metrics[mode]["completions/max_terminated_length"].append(term_completion_lengths.float().max().item()) |
|
|
| return prompt_ids, completion_ids, completions |
|
|
| def _generate_and_score_completions( |
| self, inputs: list[dict[str, torch.Tensor | Any]] |
| ) -> dict[str, torch.Tensor | Any]: |
| device = self.accelerator.device |
| mode = "train" if self.model.training else "eval" |
|
|
| prompts = [x["prompt"] for x in inputs] |
|
|
| if "images" in inputs[0]: |
| images = [example.get("images") for example in inputs] |
| elif "image" in inputs[0]: |
| images = [[example.get("image")] if example.get("image") is not None else None for example in inputs] |
| else: |
| images = None |
| |
| if images is not None and all(img_list == [] for img_list in images): |
| images = None |
|
|
| |
| |
| |
| if images is not None: |
| if not is_conversational(inputs[0]): |
| raise ValueError( |
| "Multimodal training requires conversational prompts. It looks like the dataset contains " |
| "non-conversational inputs, likely because a chat template was applied before passing the dataset " |
| "to the trainer. Please provide the raw conversational prompts and let the trainer apply the chat " |
| "template internally." |
| ) |
| prompts = [ |
| prepare_multimodal_messages(prompt, images=image_list) |
| for prompt, image_list in zip(prompts, images, strict=True) |
| ] |
|
|
| prompt_ids_list, completion_ids_list, completions = self._generate(prompts) |
|
|
| |
| prompt_ids = [torch.tensor(ids) for ids in prompt_ids_list] |
| prompt_mask = [torch.ones_like(ids, dtype=torch.long) for ids in prompt_ids] |
| prompt_ids = pad( |
| prompt_ids, |
| padding_value=self._tokenizer.pad_token_id, |
| padding_side="left", |
| pad_to_multiple_of=self.pad_to_multiple_of, |
| ).to(device=device) |
| prompt_mask = pad( |
| prompt_mask, padding_value=0, padding_side="left", pad_to_multiple_of=self.pad_to_multiple_of |
| ).to(device=device) |
| completion_ids = [torch.tensor(ids) for ids in completion_ids_list] |
| completion_mask = [torch.ones_like(ids, dtype=torch.long) for ids in completion_ids] |
| completion_ids = pad( |
| completion_ids, |
| padding_value=self._tokenizer.pad_token_id, |
| padding_side="right", |
| pad_to_multiple_of=self.pad_to_multiple_of, |
| ).to(device=device) |
| completion_mask = pad( |
| completion_mask, padding_value=0, padding_side="right", pad_to_multiple_of=self.pad_to_multiple_of |
| ).to(device=device) |
|
|
| |
| if self.mask_truncated_completions: |
| eos_and_pad = [self._tokenizer.eos_token_id, self._tokenizer.pad_token_id] |
| |
| is_truncated = torch.tensor([ids[-1] not in eos_and_pad for ids in completion_ids_list], device=device) |
| completion_mask = completion_mask * (~is_truncated).unsqueeze(1).int() |
|
|
| |
| prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1) |
| attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) |
|
|
| logits_to_keep = completion_ids.size(1) |
| batch_size = self.args.per_device_train_batch_size if mode == "train" else self.args.per_device_eval_batch_size |
|
|
| num_images = [len(img_list) if img_list else 0 for img_list in images] if images is not None else None |
|
|
| |
| if images is not None: |
| prompts_text = [ |
| apply_chat_template({"prompt": prompt}, self.processing_class, **self.chat_template_kwargs)["prompt"] |
| for prompt in prompts |
| ] |
| prompt_inputs = self.processing_class(images=images, text=prompts_text, padding=True, return_tensors="pt") |
| prompt_inputs = super()._prepare_inputs(prompt_inputs) |
| forward_kwargs = {k: v for k, v in prompt_inputs.items() if k not in ["input_ids", "attention_mask"]} |
| else: |
| forward_kwargs = {} |
|
|
| |
| num_tiles = None |
| if images is not None and "spatial_shapes" in forward_kwargs: |
| image_info = self.processing_class.image_processor( |
| images=images, return_tensors="pt", return_row_col_info=True |
| ) |
| tiles_per_image = image_info["image_rows"] * image_info["image_cols"] |
| if self.processing_class.image_processor.use_thumbnail: |
| tiles_per_image = tiles_per_image + (tiles_per_image > 1).to(tiles_per_image.dtype) |
| num_tiles = [group.sum().item() for group in torch.split(tiles_per_image, num_images)] |
|
|
| |
| if "token_type_ids" in forward_kwargs: |
| token_type_ids = forward_kwargs["token_type_ids"] |
| if self.pad_to_multiple_of is not None: |
| |
| padding_size = prompt_ids.size(1) - token_type_ids.size(1) |
| if padding_size > 0: |
| token_type_ids = torch.cat( |
| [token_type_ids.new_zeros((token_type_ids.size(0), padding_size)), token_type_ids], dim=1 |
| ) |
| forward_kwargs["token_type_ids"] = torch.cat( |
| [token_type_ids, token_type_ids.new_zeros(completion_ids.shape)], dim=1 |
| ) |
| |
| if "mm_token_type_ids" in forward_kwargs: |
| mm_token_type_ids = forward_kwargs["mm_token_type_ids"] |
| if self.pad_to_multiple_of is not None: |
| |
| padding_size = prompt_ids.size(1) - mm_token_type_ids.size(1) |
| if padding_size > 0: |
| mm_token_type_ids = torch.cat( |
| [mm_token_type_ids.new_zeros((mm_token_type_ids.size(0), padding_size)), mm_token_type_ids], |
| dim=1, |
| ) |
| forward_kwargs["mm_token_type_ids"] = torch.cat( |
| [mm_token_type_ids, mm_token_type_ids.new_zeros(completion_ids.shape)], dim=1 |
| ) |
|
|
| |
| |
| |
| with torch.no_grad(), disable_gradient_checkpointing(self.model, self.args.gradient_checkpointing_kwargs): |
| |
| old_per_token_logps, _, _ = self._get_per_token_logps_and_entropies( |
| self.model, |
| prompt_completion_ids, |
| attention_mask, |
| logits_to_keep, |
| batch_size, |
| num_images=num_images, |
| num_tiles=num_tiles, |
| **forward_kwargs, |
| ) |
| old_logps = (old_per_token_logps * completion_mask).sum(1) |
|
|
| |
| if self.beta != 0.0: |
| if self.ref_model is not None: |
| ref_per_token_logps, _, _ = self._get_per_token_logps_and_entropies( |
| self.ref_model, |
| prompt_completion_ids, |
| attention_mask, |
| logits_to_keep, |
| batch_size=batch_size, |
| num_images=num_images, |
| num_tiles=num_tiles, |
| **forward_kwargs, |
| ) |
| else: |
| |
| |
| |
| model = self.accelerator.unwrap_model(self.model) |
| with use_adapter(model, adapter_name="ref" if "ref" in model.peft_config else None): |
| ref_per_token_logps, _, _ = self._get_per_token_logps_and_entropies( |
| self.model, |
| prompt_completion_ids, |
| attention_mask, |
| logits_to_keep, |
| batch_size=batch_size, |
| num_images=num_images, |
| num_tiles=num_tiles, |
| **forward_kwargs, |
| ) |
| else: |
| ref_per_token_logps = None |
|
|
| |
| prompts_text = self.processing_class.batch_decode(prompt_ids, skip_special_tokens=True) |
| completions_text = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) |
|
|
| |
| |
| |
| rewards_per_func = self._calculate_rewards(inputs, prompts, completions, completion_ids_list) |
| num_generations = self.num_generations if mode == "train" else self.num_generations_eval |
|
|
| |
| |
| |
| unscorable_mask = torch.isnan(rewards_per_func).all(dim=1) |
|
|
| |
| rewards = (rewards_per_func * self.reward_weights.to(device).unsqueeze(0)).nansum(dim=1) |
| rewards[unscorable_mask] = torch.nan |
|
|
| |
| if self.reward_clip_range: |
| rewards = rewards.clamp(min=self.reward_clip_range[0], max=self.reward_clip_range[1]) |
|
|
| |
| if self.beta != 0.0: |
| |
| |
| |
| |
| per_token_kl = old_per_token_logps - ref_per_token_logps |
| |
| kl = (per_token_kl * completion_mask).sum(-1) |
| kl = gather(kl) |
| rewards = rewards - self.beta * kl |
|
|
| grouped_rewards = rewards.view(-1, num_generations) |
| mean_grouped_rewards = torch.nanmean(grouped_rewards, dim=1) |
| if num_generations > 1: |
| std_rewards = nanstd(grouped_rewards, dim=1) |
| else: |
| std_rewards = torch.zeros_like(mean_grouped_rewards) |
|
|
| |
| |
| |
| scorable_counts = (~torch.isnan(grouped_rewards)).sum(dim=1, keepdim=True) |
| grouped_sum = torch.nansum(grouped_rewards, dim=1, keepdim=True) |
| if num_generations > 1: |
| baselines = (grouped_sum - grouped_rewards) / (scorable_counts - 1) |
| baselines = baselines.view(-1) |
| advantages = rewards - baselines |
| else: |
| advantages = torch.zeros_like(rewards) |
|
|
| |
| if self.normalize_advantages: |
| advantages = (advantages - torch.nanmean(advantages)) / (nanstd(advantages) + 1e-4) |
|
|
| |
| advantages = torch.nan_to_num(advantages, nan=0.0) |
|
|
| is_std_zero = torch.isclose(std_rewards, torch.zeros_like(std_rewards)) |
|
|
| |
| process_slice = slice( |
| self.accelerator.process_index * len(prompts), |
| (self.accelerator.process_index + 1) * len(prompts), |
| ) |
| all_process_advantages = advantages.clone() |
| advantages = advantages[process_slice] |
|
|
| |
| if self.beta != 0.0: |
| mean_kl = (per_token_kl * completion_mask).sum() / completion_mask.sum().clamp(min=1.0) |
| self._metrics[mode]["kl"].append(self.accelerator.gather(mean_kl).nanmean().item()) |
|
|
| |
| for i, reward_func_name in enumerate(self.reward_func_names): |
| mean_rewards = torch.nanmean(rewards_per_func[:, i]).item() |
| self._metrics[mode][f"rewards/{reward_func_name}/mean"].append(mean_rewards) |
| std_func_rewards = nanstd(rewards_per_func[:, i]).item() |
| self._metrics[mode][f"rewards/{reward_func_name}/std"].append(std_func_rewards) |
| rewards = (rewards_per_func * self.reward_weights.to(rewards_per_func.device).unsqueeze(0)).nansum(dim=1) |
| rewards[unscorable_mask] = torch.nan |
| self._metrics[mode]["reward"].append(torch.nanmean(rewards).item()) |
| self._metrics[mode]["reward_std"].append(nanstd(rewards).item()) |
| self._metrics[mode]["frac_reward_zero_std"].append(is_std_zero.float().mean().item()) |
|
|
| |
| self._logs["prompt"].extend(gather_object(prompts_text)) |
| self._logs["completion"].extend(gather_object(completions_text)) |
| for i, name in enumerate(self.reward_func_names): |
| self._logs["rewards"][name].extend(rewards_per_func[:, i].tolist()) |
| self._logs["advantages"].extend(all_process_advantages.tolist()) |
|
|
| |
| |
| |
| for column in sorted(self._pending_extra_logs): |
| self._logs["extra"][column].extend(gather_object(self._pending_extra_logs[column])) |
| self._pending_extra_logs.clear() |
|
|
| |
| |
| |
| for name in sorted(self._pending_metrics): |
| values = self._pending_metrics[name] |
| local_mean = sum(values) / len(values) |
| global_mean = self.accelerator.gather(torch.tensor(local_mean, device=device)).mean().item() |
| self._metrics[mode][name].append(global_mean) |
| self._pending_metrics.clear() |
|
|
| if images is not None: |
| self._logs["images"].extend(gather_object(images)) |
|
|
| output = { |
| "prompt_ids": prompt_ids, |
| "prompt_mask": prompt_mask, |
| "completion_ids": completion_ids, |
| "completion_mask": completion_mask, |
| "old_logps": old_logps, |
| "advantages": advantages, |
| } |
| if "pixel_values" in forward_kwargs: |
| output["pixel_values"] = forward_kwargs["pixel_values"] |
| if "image_grid_thw" in forward_kwargs: |
| output["image_grid_thw"] = forward_kwargs["image_grid_thw"] |
| if "pixel_attention_mask" in forward_kwargs: |
| output["pixel_attention_mask"] = forward_kwargs["pixel_attention_mask"] |
| if "spatial_shapes" in forward_kwargs: |
| output["spatial_shapes"] = forward_kwargs["spatial_shapes"] |
| if "image_sizes" in forward_kwargs: |
| output["image_sizes"] = forward_kwargs["image_sizes"] |
| if "token_type_ids" in forward_kwargs: |
| output["token_type_ids"] = forward_kwargs["token_type_ids"] |
| if "mm_token_type_ids" in forward_kwargs: |
| output["mm_token_type_ids"] = forward_kwargs["mm_token_type_ids"] |
| if "image_position_ids" in forward_kwargs: |
| output["image_position_ids"] = forward_kwargs["image_position_ids"] |
| if images is not None: |
| output["num_images"] = num_images |
| if num_tiles is not None: |
| output["num_tiles"] = num_tiles |
| return output |
|
|
| @profiling_decorator |
| def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): |
| if return_outputs: |
| raise ValueError("The RLOOTrainer does not support returning outputs") |
| return self._compute_loss(model, inputs) |
|
|
| def _compute_loss(self, model, inputs): |
| |
| prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"] |
| completion_ids, completion_mask = inputs["completion_ids"], inputs["completion_mask"] |
| input_ids = torch.cat([prompt_ids, completion_ids], dim=1) |
| attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) |
| logits_to_keep = completion_ids.size(1) |
|
|
| |
| per_token_logps, entropies, aux_loss = self._get_per_token_logps_and_entropies( |
| model, |
| input_ids, |
| attention_mask, |
| logits_to_keep, |
| compute_entropy=True, |
| compute_aux_loss=self.aux_loss_enabled, |
| pixel_values=inputs.get("pixel_values"), |
| image_grid_thw=inputs.get("image_grid_thw"), |
| num_images=inputs.get("num_images"), |
| pixel_attention_mask=inputs.get("pixel_attention_mask"), |
| spatial_shapes=inputs.get("spatial_shapes"), |
| num_tiles=inputs.get("num_tiles"), |
| image_sizes=inputs.get("image_sizes"), |
| token_type_ids=inputs.get("token_type_ids"), |
| mm_token_type_ids=inputs.get("mm_token_type_ids"), |
| image_position_ids=inputs.get("image_position_ids"), |
| ) |
|
|
| logps = (per_token_logps * completion_mask).sum(1) |
| old_logps = inputs["old_logps"] |
| log_ratio = logps - old_logps |
|
|
| |
| advantages = inputs["advantages"] |
| coef_1 = torch.exp(log_ratio) |
| coef_2 = torch.clamp(coef_1, 1 - self.epsilon_low, 1 + self.epsilon_high) |
| per_sequence_loss1 = coef_1 * advantages |
| per_sequence_loss2 = coef_2 * advantages |
| per_sequence_loss = -torch.min(per_sequence_loss1, per_sequence_loss2) |
| loss = per_sequence_loss.mean() |
|
|
| |
| mode = "train" if self.model.training else "eval" |
|
|
| |
| if self.aux_loss_enabled: |
| loss = loss + self.router_aux_loss_coef * aux_loss |
| self._metrics[mode]["aux_loss"].append(self.accelerator.gather_for_metrics(aux_loss).mean().item()) |
|
|
| |
| mean_entropy = (entropies * completion_mask).sum() / completion_mask.sum().clamp(min=1.0) |
| self._metrics[mode]["entropy"].append(self.accelerator.gather(mean_entropy).nanmean().item()) |
|
|
| |
| is_low_clipped = (coef_1 < 1 - self.epsilon_low) & (advantages < 0) |
| is_high_clipped = (coef_1 > 1 + self.epsilon_high) & (advantages > 0) |
| is_region_clipped = is_low_clipped | is_high_clipped |
| gathered_low_clip = self.accelerator.gather(is_low_clipped.float().mean()) |
| self._metrics[mode]["clip_ratio/low_mean"].append(gathered_low_clip.nanmean().item()) |
| self._metrics[mode]["clip_ratio/low_min"].append(nanmin(gathered_low_clip).item()) |
| gathered_high_clip = self.accelerator.gather(is_high_clipped.float().mean()) |
| self._metrics[mode]["clip_ratio/high_mean"].append(gathered_high_clip.nanmean().item()) |
| self._metrics[mode]["clip_ratio/high_max"].append(nanmax(gathered_high_clip).item()) |
| gathered_clip_ratio = self.accelerator.gather(is_region_clipped.float().mean()) |
| self._metrics[mode]["clip_ratio/region_mean"].append(gathered_clip_ratio.nanmean().item()) |
| return loss |
|
|
| |
| |
| def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys: list[str] | None = None): |
| inputs = self._prepare_inputs(inputs) |
| with torch.no_grad(): |
| with self.compute_loss_context_manager(): |
| loss = self.compute_loss(model, inputs) |
| loss = loss.mean().detach() |
| return loss, None, None |
|
|
| def log(self, logs: dict[str, float], start_time: float | None = None) -> None: |
| mode = "train" if self.model.training else "eval" |
| |
| metrics = {} |
| for key, val in self._metrics[mode].items(): |
| |
| |
| |
| |
| valid = [v for v in val if not math.isnan(v)] |
| metrics[key] = sum(valid) / len(valid) if valid else None |
|
|
| |
| |
| if mode == "eval": |
| metrics = {f"eval_{key}": val for key, val in metrics.items()} |
|
|
| logs.update(metrics) |
| super().log(logs, start_time) |
| self._metrics[mode].clear() |
|
|
| if self.accelerator.is_main_process and self.log_completions: |
| if is_rich_available(): |
| print_prompt_completions_sample( |
| self._logs["prompt"], |
| self._logs["completion"], |
| self._logs["rewards"], |
| self._logs["advantages"], |
| self.state.global_step, |
| self.num_completions_to_print, |
| extra=dict(self._logs["extra"]), |
| ) |
|
|
| logging_backends = [] |
| if self.args.report_to and "wandb" in self.args.report_to and wandb.run is not None: |
| logging_backends.append(wandb) |
| if self.args.report_to and "trackio" in self.args.report_to: |
| logging_backends.append(trackio) |
|
|
| table = { |
| "step": [self.state.global_step] * len(self._logs["prompt"]), |
| "prompt": self._logs["prompt"], |
| "completion": self._logs["completion"], |
| **self._logs["rewards"], |
| **self._logs["extra"], |
| "advantage": self._logs["advantages"], |
| } |
|
|
| df_base = pd.DataFrame(table) |
| images_raw = self._logs["images"] or [] |
|
|
| for logging_backend in logging_backends: |
| if images_raw: |
| images = [] |
| for image_list in self._logs["images"]: |
| images.append([logging_backend.Image(image) for image in image_list]) |
| df = pd.concat( |
| [df_base, pd.Series(images, name="image")], |
| axis=1, |
| copy=False, |
| ) |
| else: |
| df = df_base |
|
|
| if self.log_unique_prompts: |
| df = df.drop_duplicates(subset=["prompt"]) |
|
|
| logging_backend.log({"completions": logging_backend.Table(dataframe=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) |
| class UnslothRLOOTrainer(_UnslothRLOOTrainer): |
| """ |
| |
| Trainer for the Reinforce Leave One Out (RLOO) method. This algorithm was initially proposed in the paper [Back to |
| Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in |
| LLMs](https://huggingface.co/papers/2402.14740). |
| |
| Example: |
| |
| ```python |
| >>> from trl import RLOOTrainer |
| >>> from trl.rewards import accuracy_reward |
| >>> from datasets import load_dataset |
| |
| >>> dataset = load_dataset("trl-lib/DeepMath-103K", split="train") |
| |
| >>> trainer = RLOOTrainer( |
| ... model="Qwen/Qwen2.5-0.5B-Instruct", |
| ... reward_funcs=accuracy_reward, |
| ... train_dataset=dataset, |
| ... ) |
| >>> trainer.train() |
| ``` |
| |
| Args: |
| model (`str` or [`~transformers.PreTrainedModel`] or [`~peft.PeftModel`]): |
| Model to be trained. Can be either: |
| |
| - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a |
| path to a *directory* containing model weights saved using |
| [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded |
| using `<ModelArchitecture>.from_pretrained` (where `<ModelArchitecture>` is derived from the model |
| config) with the keyword arguments in `args.model_init_kwargs`. |
| - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. |
| - A [`~peft.PeftModel`] object. Only causal language models are supported. |
| reward_funcs (`RewardFunc | list[RewardFunc]`): |
| Reward functions to be used for computing the rewards. To compute the rewards, we call all the reward |
| functions with the prompts and completions and sum the rewards. Can be either: |
| |
| - A single reward function, such as: |
| - A string: The *model ID* of a pretrained model hosted inside a model repo on huggingface.co, or a |
| path to a *directory* containing model weights saved using |
| [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded |
| using [`~transformers.AutoModelForSequenceClassification.from_pretrained`] with `num_labels=1` and the |
| keyword arguments in `args.model_init_kwargs`. |
| - A [`~transformers.PreTrainedModel`] object: Only sequence classification models are supported. |
| - A custom reward function: The function is provided with the prompts and the generated completions, |
| plus any additional columns in the dataset. It should return a list of rewards. Custom reward |
| functions can be either synchronous or asynchronous and can also return `None` when the reward is |
| not applicable to those samples. This is useful for multi-task training where different reward |
| functions apply to different types of samples. When a reward function returns `None` for a sample, |
| that reward function is excluded from the reward calculation for that sample. For more details, see |
| [Using a custom reward |
| function](#using-a-custom-reward-function). |
| |
| The trainer's state is also passed to the reward function. The trainer's state is an instance of |
| [`~transformers.TrainerState`] and can be accessed by accessing the `trainer_state` argument to the |
| reward function's signature. |
| - A list of reward functions, where each item can independently be any of the above types. Mixing different |
| types within the list (e.g., a string model ID and a custom reward function) is allowed. |
| args ([`RLOOConfig`], *optional*): |
| Configuration for this trainer. If `None`, a default configuration is used. |
| train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): |
| Dataset to use for training. It must include a column `"prompt"`. Any additional columns in the dataset is |
| ignored. The format of the samples can be either: |
| |
| - [Standard](dataset_formats#standard): Each sample contains plain text. |
| - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role |
| and content). |
| eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`] or `dict[str, Dataset | IterableDataset]`): |
| Dataset to use for evaluation. It must meet the same requirements as `train_dataset`. |
| processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.ProcessorMixin`], *optional*): |
| Processing class used to process the data. The padding side must be set to "left". If `None`, the |
| processing class is loaded from the model's name with [`~transformers.AutoProcessor.from_pretrained`]. A |
| padding token, `tokenizer.pad_token`, must be set. If the processing class has not set a padding token, |
| `tokenizer.eos_token` will be used as the default. |
| reward_processing_classes ([`~transformers.PreTrainedTokenizerBase`] or `list[PreTrainedTokenizerBase]`, *optional*): |
| Processing classes corresponding to the reward functions specified in `reward_funcs`. Can be either: |
| |
| - A single processing class: Used when `reward_funcs` contains only one reward function. |
| - A list of processing classes: Must match the order and length of the reward functions in `reward_funcs`. |
| If set to `None`, or if an element of the list corresponding to a [`~transformers.PreTrainedModel`] is |
| `None`, the tokenizer for the model is automatically loaded using |
| [`~transformers.AutoTokenizer.from_pretrained`]. For elements in `reward_funcs` that are custom reward |
| functions (not [`~transformers.PreTrainedModel`]), the corresponding entries in `reward_processing_classes` |
| are ignored. |
| callbacks (list of [`~transformers.TrainerCallback`], *optional*): |
| List of callbacks to customize the training loop. Will add those to the list of default callbacks detailed |
| in [here](https://huggingface.co/docs/transformers/main_classes/callback). |
| |
| If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`] |
| method. |
| optimizers (`tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None]`, *optional*, defaults to `(None, None)`): |
| A tuple containing the optimizer and the scheduler to use. Will default to an instance of `AdamW` on your |
| model and a scheduler given by [`~transformers.get_linear_schedule_with_warmup`] controlled by `args`. |
| peft_config ([`~peft.PeftConfig`], *optional*): |
| PEFT configuration used to wrap the model. If `None`, the model is not wrapped. |
| |
| """ |
| def __init__( |
| self, |
| model, |
| reward_funcs, |
| args = None, |
| train_dataset = None, |
| eval_dataset = None, |
| processing_class = None, |
| reward_processing_classes = None, |
| callbacks = None, |
| peft_config = None, |
| **kwargs |
| ): |
| if args is None: args = UnslothRLOOConfig() |
| 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().weight.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' |
| if hasattr(args, 'mixed_precision'): args.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 hasattr(args, 'mixed_precision'): args.mixed_precision = 'fp16' if float16 else 'bf16' |
| |
| elif mixed_precision_dtype == 'bfloat16': |
| |
| args.fp16 = False |
| args.bf16 = False |
| os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' |
| if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' |
| |
| |
| 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 model is not None: |
| _warnings_issued = getattr(model, 'warnings_issued', None) |
| if _warnings_issued is None: |
| model.warnings_issued = {} |
| elif not isinstance(_warnings_issued, dict): |
| try: |
| model.warnings_issued = dict(_warnings_issued) |
| except Exception: |
| model.warnings_issued = {} |
| 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 |
| elif args_max_seq_length is not None and model_max_seq_length is not None: |
| if args_max_seq_length > model_max_seq_length: |
| print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but ' |
| 'the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.') |
| args.max_seq_length = model_max_seq_length |
| if model is not None and hasattr(model, 'for_training'): |
| model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True)) |
| 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' |
| other_metrics = [] |
| |
| from unsloth_zoo.logging_utils import PatchRLStatistics |
| PatchRLStatistics('rloo_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(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True)) |
| super().__init__( |
| model = model, |
| reward_funcs = reward_funcs, |
| args = args, |
| train_dataset = train_dataset, |
| eval_dataset = eval_dataset, |
| processing_class = processing_class, |
| reward_processing_classes = reward_processing_classes, |
| 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 |
| if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'): |
| _vllm_tok = self.llm.get_tokenizer() |
| _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None) |
| if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None: |
| _vllm_tok.chat_template = _pc.chat_template |
| pass |
| |
| pass |
|
|
|
|
| if hasattr(logger, "addFilter"): |
| import logging |
| class HideLoggingMessage(logging.Filter): |
| def __init__(self, text): self.text = text |
| def filter(self, x): return not (self.text in x.getMessage()) |
| pass |
| logger.addFilter(HideLoggingMessage("`use_cache=True`")) |
|
|
|
|
