# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. __all__ = [ "RL_EXTRA_ARGS", "RL_FUNCTIONS", "RL_PRE_ITEMS", "RL_CONFIG_CHANGES", "RL_METRICS_CHANGES", ] import os import re import torch import inspect import linecache from collections import defaultdict from unsloth_zoo.rl_replacements import ( RL_REPLACEMENTS, left_pack_padding, create_completion_attention_mask, chunked_selective_log_softmax, chunked_hidden_states_selective_log_softmax, _unsloth_get_mm_token_id, _unsloth_fix_mm_token_type_ids, ) from unsloth_zoo.utils import Version from trl import __version__ as trl_version_raw from importlib.metadata import version as importlib_version from unsloth_zoo.log import logger from unsloth_zoo.device_type import device_synchronize import importlib.util from ..device_type import ( is_hip, get_device_type, DEVICE_TYPE, DEVICE_TYPE_TORCH, DEVICE_COUNT, ALLOW_PREQUANTIZED_MODELS, ) import textwrap from ._utils import _get_inference_mode_context_manager, UNSLOTH_ENABLE_LOGGING # One-time GRPO sequence-packing gates; mirrored into the generated trainer cache via RL_PRE_ITEMS. UNSLOTH_GRPO_SEQ_PACKING_ON = os.environ.get("UNSLOTH_GRPO_SEQ_PACKING", "1").lower() not in ( "0", "false", "no", "off", ) # Packing needs zoo#840's masked-column guard in grpo_compute_loss (installed zoo is fixed per-process). try: UNSLOTH_ZOO_HAS_MASKED_COL_GUARD = "torch.where(_keep, new" in inspect.getsource( RL_REPLACEMENTS["grpo_compute_loss"] ) except Exception: UNSLOTH_ZOO_HAS_MASKED_COL_GUARD = False # One-time PrefixGrouper gate; any import failure degrades to "PrefixGrouper off". _pg_build_layout = _pg_enabled_fn = _pg_verify_on = _pg_tol_ok = _PG_TOL_KILL = None UNSLOTH_GRPO_PREFIX_GROUPER_ON = os.environ.get("UNSLOTH_GRPO_PREFIX_GROUPER", "1").lower() not in ( "0", "false", "no", "off", ) if UNSLOTH_GRPO_PREFIX_GROUPER_ON: try: from ..utils.prefix_grouper import ( build_group_layout as _pg_build_layout, prefix_grouper_enabled as _pg_enabled_fn, verify_on as _pg_verify_on, tol_ok as _pg_tol_ok, TOL_KILL as _PG_TOL_KILL, ) except Exception: UNSLOTH_GRPO_PREFIX_GROUPER_ON = False RL_EXTRA_ARGS = defaultdict(list) RL_FUNCTIONS = defaultdict(list) RL_PRE_ITEMS = defaultdict(list) def _unsloth_clear_stateful_mrope(model): modules = getattr(model, "modules", None) if modules is None: return False cleared = False for module in modules(): if hasattr(module, "compute_3d_position_ids") and hasattr(module, "rope_deltas"): module.rope_deltas = None cleared = True return cleared RL_CONFIG_CHANGES = defaultdict(list) RL_METRICS_CHANGES = defaultdict(list) RL_ADDITIONAL_FUNCTIONS = defaultdict(list) _DPO_VISION_KEYS = ( "pixel_position_ids", "image_position_ids", "mm_token_type_ids", ) torch_compile_options = { "epilogue_fusion": True, "max_autotune": False, # I saw speedups, but not sure if this has issues in collab "shape_padding": True, "trace.enabled": False, "triton.cudagraphs": False, } try: trl_version = Version(trl_version_raw) except Exception: try: trl_version = Version(importlib_version("trl")) except Exception: trl_version = Version("0.0.0") # Check untrained tokens def sft_trainer_fix_untrained_tokens(call_args, extra_args): if "model" in call_args and "train_dataset" in call_args: fix_tokenizer = ( "IGNORED_TOKENIZER_NAMES = os.environ.get('UNSLOTH_IGNORED_TOKENIZER_NAMES', '').split('\\n')\n" "from unsloth_zoo.tokenizer_utils import fix_untrained_tokens\n" "from unsloth_zoo.training_utils import fix_zero_training_loss\n" "if 'tokenizer' not in locals(): tokenizer = processing_class\n" "fix_untrained_tokens(model, tokenizer, train_dataset, IGNORED_TOKENIZER_NAMES, eps = 1e-16)\n" "fix_zero_training_loss(model, tokenizer, train_dataset)\n" ) return fix_tokenizer return "" RL_EXTRA_ARGS["sft_trainer"].append(sft_trainer_fix_untrained_tokens) # Fix top_k for GRPO vLLM. # https://github.com/huggingface/trl/pull/4695 with this change trl added top_k in GRPOConfig and defaults to 0 # We don't want that since vllm's all include top_k is -1 and 0 returns an error on SamplingParams creation. def grpo_config_fix_vllm_top_k(old_RLTrainer_source, old_RLConfig_source): return "if use_vllm and (top_k is None or top_k == 0): top_k = -1\n" RL_CONFIG_CHANGES["grpo_trainer"].append(grpo_config_fix_vllm_top_k) # Remove DPO columns which might randomnly be tokenized def dpo_trainer_fix_columns(call_args, extra_args): if "model" in call_args and "train_dataset" in call_args: fix_dpo = ( "if hasattr(train_dataset, 'column_names'):\n" " column_names = set(train_dataset.column_names)\n" " check = ['chosen', 'rejected', 'prompt', 'chosen_input_ids', 'chosen_attention_mask',\n" " 'chosen_labels', 'rejected_input_ids', 'rejected_attention_mask', 'rejected_labels',\n" " 'prompt_input_ids', 'prompt_attention_mask']\n" " if all(x in column_names for x in check):\n" " train_dataset = train_dataset.remove_columns(['chosen', 'rejected', 'prompt'])\n" " del check, column_names\n" ) return fix_dpo return "" RL_EXTRA_ARGS["dpo_trainer"].append(dpo_trainer_fix_columns) def dpo_trainer_fix_data_collator(call_args, extra_args): if ( "data_collator" in call_args and "train_dataset" in call_args and "processing_class" in call_args ): fix_collator = ( "if hasattr(train_dataset, 'column_names'):\n" " column_names = set(train_dataset.column_names)\n" " is_dpo_dataset = ({'chosen', 'rejected'}.issubset(column_names) or\n" " {'prompt_input_ids', 'chosen_input_ids', 'rejected_input_ids'}.issubset(column_names))\n" " if is_dpo_dataset and isinstance(data_collator, TransformersDataCollatorForLanguageModeling):\n" " data_collator = None\n" " del is_dpo_dataset, column_names\n" ) return fix_collator return "" RL_EXTRA_ARGS["dpo_trainer"].append(dpo_trainer_fix_data_collator) def dpo_trainer_vision_process_row( features, processing_class, max_prompt_length = None, max_completion_length = None, add_special_tokens = True, is_chat = False, ): text = features.get("prompt", "") images = features.get("images") processor, tokenizer = processing_class, processing_class.tokenizer processed_features = processor( images = images, text = text, add_special_tokens = False, ) prompt_input_ids = processed_features["input_ids"][0] chosen_input_ids = tokenizer(features["chosen"], add_special_tokens = False)["input_ids"] rejected_input_ids = tokenizer(features["rejected"], add_special_tokens = False)["input_ids"] if add_special_tokens: if tokenizer.bos_token_id is not None: prompt_input_ids = [tokenizer.bos_token_id] + prompt_input_ids if tokenizer.eos_token_id is not None: prompt_input_ids = prompt_input_ids + [tokenizer.eos_token_id] if not is_chat and tokenizer.eos_token_id is not None: chosen_input_ids = chosen_input_ids + [tokenizer.eos_token_id] rejected_input_ids = rejected_input_ids + [tokenizer.eos_token_id] if max_prompt_length is not None: prompt_input_ids = prompt_input_ids[-max_prompt_length:] if max_completion_length is not None: chosen_input_ids = chosen_input_ids[:max_completion_length] rejected_input_ids = rejected_input_ids[:max_completion_length] output = { "prompt_input_ids": prompt_input_ids, "chosen_input_ids": chosen_input_ids, "rejected_input_ids": rejected_input_ids, } if "pixel_values" in processed_features: output["pixel_values"] = processed_features["pixel_values"][0] if "pixel_attention_mask" in processed_features: output["pixel_attention_mask"] = processed_features["pixel_attention_mask"][0] if "image_sizes" in processed_features: output["image_sizes"] = processed_features["image_sizes"][0] if "token_type_ids" in processed_features: token_type_ids = processed_features["token_type_ids"][0] if max_prompt_length is not None: token_type_ids = token_type_ids[-max_prompt_length:] output["token_type_ids"] = token_type_ids if "pixel_position_ids" in processed_features: output["pixel_position_ids"] = processed_features["pixel_position_ids"][0] if "image_position_ids" in processed_features: output["image_position_ids"] = processed_features["image_position_ids"][0] if "mm_token_type_ids" in processed_features: mm_token_type_ids = processed_features["mm_token_type_ids"][0] if max_prompt_length is not None: mm_token_type_ids = mm_token_type_ids[-max_prompt_length:] output["mm_token_type_ids"] = mm_token_type_ids return output def dpo_trainer_vision_signature_columns(function_name, function): if function_name != "_set_signature_columns_if_needed": return function if all(_k in function for _k in _DPO_VISION_KEYS): return function _extra_columns = "".join(f' "{_k}",\n' for _k in _DPO_VISION_KEYS) new_function = function.replace( ' "image_sizes",\n "token_type_ids",\n', f' "image_sizes",\n' f"{_extra_columns}" f' "token_type_ids",\n', ) if new_function != function: return new_function return function.replace( ' "image_sizes",\n "ref_chosen_logps",\n', f' "image_sizes",\n' f"{_extra_columns}" f' "ref_chosen_logps",\n', ) def dpo_trainer_concatenated_inputs(function_name, function): if function_name != "concatenated_inputs": return function if all(_k in function for _k in _DPO_VISION_KEYS): return function _extra_inputs = "".join( f' if "{_k}" in batch:\n' f' output["{_k}"] = torch.cat((batch["{_k}"], batch["{_k}"]), dim=0)\n' for _k in _DPO_VISION_KEYS ) image_sizes_block = ( ' if "image_sizes" in batch:\n' ' output["image_sizes"] = torch.cat([batch["image_sizes"], batch["image_sizes"]], dim=0)\n' ) new_function = function.replace( image_sizes_block + ' if "token_type_ids" in batch:\n', image_sizes_block + _extra_inputs + ' if "token_type_ids" in batch:\n', ) if new_function != function: return new_function if image_sizes_block in function: return function.replace(image_sizes_block, image_sizes_block + _extra_inputs, 1) return function def _dpo_trainer_extend_vision_model_kwargs(function): if all(_k in function for _k in _DPO_VISION_KEYS): return function _extra_forward = "".join( f' if "{_k}" in concatenated_batch:\n' f' model_kwargs["{_k}"] = concatenated_batch["{_k}"]\n' for _k in ( "pixel_values", "pixel_attention_mask", "image_sizes", *_DPO_VISION_KEYS, ) ) return function.replace( ' if "pixel_values" in concatenated_batch:\n' ' model_kwargs["pixel_values"] = concatenated_batch["pixel_values"]\n' ' if "pixel_attention_mask" in concatenated_batch:\n' ' model_kwargs["pixel_attention_mask"] = concatenated_batch["pixel_attention_mask"]\n' ' if "image_sizes" in concatenated_batch:\n' ' model_kwargs["image_sizes"] = concatenated_batch["image_sizes"]\n', f"{_extra_forward}", ) def dpo_trainer_concatenated_forward(function_name, function): if function_name != "concatenated_forward": return function return _dpo_trainer_extend_vision_model_kwargs(function) def dpo_trainer_compute_loss_liger(function_name, function): if function_name != "_compute_loss_liger": return function return _dpo_trainer_extend_vision_model_kwargs(function) def dpo_trainer_data_collator_vision_keys(call_args, extra_args): if "data_collator" not in call_args: return "" _vision_keys = str(_DPO_VISION_KEYS) return ( "from trl.trainer.dpo_trainer import DataCollatorForPreference\n" "if not hasattr(DataCollatorForPreference, '_unsloth_vision_keys_patch'):\n" " _old_dpo_collator_torch_call = DataCollatorForPreference.torch_call\n" "\n" " def _unsloth_dpo_torch_call(self, examples):\n" " output = _old_dpo_collator_torch_call(self, examples)\n" " import torch as _unsloth_torch\n" " try:\n" " from trl.trainer.utils import pad as _unsloth_trl_pad\n" " except Exception:\n" " _unsloth_trl_pad = None\n" " for _k in " + _vision_keys + ":\n" " if not all(_k in example for example in examples):\n" " continue\n" " _is_position_key = _k.endswith('position_ids')\n" " _padding_value = -1 if _is_position_key else 0\n" " _padding_side = 'right' if _is_position_key else 'left'\n" " _values = [_unsloth_torch.as_tensor(example[_k]) for example in examples]\n" " try:\n" " if _unsloth_trl_pad is not None:\n" " output[_k] = _unsloth_trl_pad(_values, padding_value=_padding_value, padding_side=_padding_side)\n" " else:\n" " from torch.nn.utils.rnn import pad_sequence as _unsloth_pad_sequence\n" " output[_k] = _unsloth_pad_sequence(_values, batch_first=True, padding_value=_padding_value)\n" " except Exception:\n" " from torch.nn.utils.rnn import pad_sequence as _unsloth_pad_sequence\n" " output[_k] = _unsloth_pad_sequence(_values, batch_first=True, padding_value=_padding_value)\n" " return output\n" "\n" " DataCollatorForPreference.torch_call = _unsloth_dpo_torch_call\n" " DataCollatorForPreference._unsloth_vision_keys_patch = True\n" ) def dpo_trainer_prepare_dataset(function_name, function): if function_name != "_prepare_dataset": return function legacy_call = "self.tokenize_row if not self.is_vision_model else self.process_row" if legacy_call not in function: return function function = function.replace( legacy_call, "self.tokenize_row if not self.is_vision_model else dpo_trainer_vision_process_row", ) legacy_tokenize_block = ( " # Tokenize the dataset\n" " if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc`\n" ' map_kwargs["desc"] = f"Tokenizing {dataset_name} dataset"\n' "\n" " dataset = dataset.map(\n" " self.tokenize_row if not self.is_vision_model else dpo_trainer_vision_process_row,\n" ) patched_tokenize_block = ( " # Tokenize the dataset\n" " if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc`\n" ' map_kwargs["desc"] = f"Tokenizing {dataset_name} dataset"\n' " if self.is_vision_model:\n" ' map_kwargs.pop("num_proc", None)\n' "\n" " dataset = dataset.map(\n" " self.tokenize_row if not self.is_vision_model else dpo_trainer_vision_process_row,\n" ) if legacy_tokenize_block in function: function = function.replace(legacy_tokenize_block, patched_tokenize_block, 1) return function RL_FUNCTIONS["dpo_trainer"].append(dpo_trainer_prepare_dataset) RL_PRE_ITEMS["dpo_trainer"].append(inspect.getsource(dpo_trainer_vision_process_row)) RL_FUNCTIONS["dpo_trainer"].append(dpo_trainer_vision_signature_columns) RL_FUNCTIONS["dpo_trainer"].append(dpo_trainer_concatenated_inputs) RL_FUNCTIONS["dpo_trainer"].append(dpo_trainer_concatenated_forward) RL_FUNCTIONS["dpo_trainer"].append(dpo_trainer_compute_loss_liger) RL_EXTRA_ARGS["dpo_trainer"].append(dpo_trainer_data_collator_vision_keys) # Fix tokenizer double BOS def sft_trainer_prepare_dataset(function_name, function): if function_name != "_prepare_non_packed_dataloader" and function_name != "_prepare_dataset": return function fast_sft_prepare_dataset = RL_REPLACEMENTS.get("sft_prepare_dataset", None) if fast_sft_prepare_dataset is not None: params = inspect.signature(fast_sft_prepare_dataset).parameters.keys() params = ".*?".join(params) matched = re.match( r"[\s]{0,}def _prepare_dataset\(.*?" + params + r".*?\)", function, flags = re.MULTILINE | re.DOTALL, ) if matched: # Use fast version! function = inspect.getsource(fast_sft_prepare_dataset) function = function.split("\n") function = "\n".join(" " * 4 + x for x in function) function = function.replace("def sft_prepare_dataset", "def _prepare_dataset") return function check_text = ( "if 'skip_prepare_dataset' in locals() and skip_prepare_dataset:\n" " return dataset\n" "if 'tokenizer' not in locals(): tokenizer = processing_class\n" "if 'formatting_func' not in locals(): raise RuntimeError('Unsloth: Please file a bug report - `formatting_func` does not exist!')\n" "if 'dataset_text_field' not in locals() and 'args' in locals(): dataset_text_field = args.dataset_text_field\n" "if 'dataset_text_field' not in locals(): dataset_text_field = None\n" "if formatting_func is None and dataset_text_field is None and 'prompt' in dataset[0] and 'completion' in dataset[0]:\n" " test_text = (dataset[0]['prompt'] + dataset[0]['completion']) if (isinstance(dataset[0]['prompt'], str) and isinstance(dataset[0]['completion'], str)) else None\n" "elif formatting_func is None and dataset_text_field is not None:\n" " test_text = dataset[0][dataset_text_field]\n" "elif formatting_func is not None:\n" " test_text = formatting_func(dataset[0])[0]\n" "else:\n" " test_text = None\n" "chat_template = getattr(tokenizer, 'chat_template', None)\n" "chat_template = '' if chat_template is None else chat_template\n" "has_bos_token_already = ((test_text is not None and test_text.startswith(tokenizer.bos_token)) or tokenizer.bos_token in chat_template) " "if getattr(tokenizer, 'bos_token', None) is not None else False\n" "if 'add_special_tokens' not in locals() and has_bos_token_already:\n" " from functools import partial\n" " tokenizer_call = tokenizer.__call__\n" " tokenizer.__call__ = partial(tokenizer_call, add_special_tokens = False)\n" " processing_class = tokenizer\n" "else:\n" " tokenizer_call = None\n" " add_special_tokens = False if has_bos_token_already else locals().get('add_special_tokens', False)\n" ) check_text = check_text.split("\n") check_text = "\n".join(" " * 8 + x for x in check_text) check_text = check_text.rstrip() + "\n" # .*? matches first match. .+? matches final match. replacer = re.findall( r"def " + function_name + r"\(.*?\).*?\:\n", function, flags = re.MULTILINE | re.DOTALL, ) if len(replacer) != 0: replacer = replacer[0] function = function.replace(replacer, replacer + check_text) # Return tokenizer's original state return_state = "if tokenizer_call is not None: tokenizer.__call__ = tokenizer_call\n" function = re.sub( r"\n([ ]{4,})(return .*?[\s]{0,})$", rf"\1{return_state}\1\2", function, ) return function RL_FUNCTIONS["sft_trainer"].append(sft_trainer_prepare_dataset) # Ignore mean_token_accuracy since it needs logits # We override it directly with our version def sft_trainer_compute_loss(function_name, function): if function_name != "compute_loss": return function def compute_loss( self, model, inputs, return_outputs = False, num_items_in_batch = None, ): outputs = super().compute_loss( model, inputs, return_outputs = return_outputs, num_items_in_batch = num_items_in_batch, ) return outputs function = inspect.getsource(compute_loss) return function RL_FUNCTIONS["sft_trainer"].append(sft_trainer_compute_loss) # Route ORPO/CPO row tokenization through the underlying text tokenizer when the # processing class is a multimodal processor; CPO reuses this code (#4952). def orpo_trainer_text_tokenizer(function_name, function): if function_name == "build_tokenized_answer": function = re.sub( r"(?m)^([ \t]*)full_tokenized = self\.processing_class\(prompt \+ answer, add_special_tokens=False\)\n" r'\1prompt_input_ids = self\.processing_class\(prompt, add_special_tokens=False\)\["input_ids"\]\n', r'\1tokenizer = getattr(self.processing_class, "tokenizer", self.processing_class)' "\n" r"\1full_tokenized = tokenizer(prompt + answer, add_special_tokens=False)" "\n" r'\1prompt_input_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"]' "\n", function, count = 1, ) return function if function_name != "tokenize_row": return function if ( 'tokenizer = getattr(self.processing_class, "tokenizer", self.processing_class)' not in function ): new_function = re.sub( r"(?m)^([ \t]*)batch = \{\}\n", r"\1batch = {}" "\n" r'\1tokenizer = getattr(self.processing_class, "tokenizer", self.processing_class)' "\n", function, count = 1, ) if new_function == function: return function function = new_function function = function.replace("self.processing_class(", "tokenizer(") function = function.replace("self.processing_class.bos_token_id", "tokenizer.bos_token_id") function = function.replace("self.processing_class.eos_token_id", "tokenizer.eos_token_id") return function RL_FUNCTIONS["orpo_trainer"].append(orpo_trainer_text_tokenizer) RL_FUNCTIONS["cpo_trainer"].append(orpo_trainer_text_tokenizer) # Resolve `processing_class.pad_token_id` through the underlying tokenizer when # a multimodal processor is supplied (processors lack `pad_token_id`). Without # this, ORPO/CPOTrainer.__init__ raises AttributeError on # `DPODataCollatorWithPadding(pad_token_id=processing_class.pad_token_id, ...)` # and on `self.padding_value = ... else processing_class.pad_token_id`. _PAD_FALLBACK = ( "(getattr(processing_class, 'pad_token_id', None) " "if getattr(processing_class, 'pad_token_id', None) is not None " "else getattr(getattr(processing_class, 'tokenizer', None), 'pad_token_id', None))" ) def orpo_trainer_processor_pad_token(function_name, function): if function_name != "__init__": return function # Multimodal processors (e.g. Gemma3/Gemma4 Processor) expose pad_token / # eos_token on `.tokenizer`, not on the processor itself. TRL 1.x CPO/ORPO # __init__ defaults `processing_class.pad_token` from `.eos_token` before # tokenizing, which AttributeErrors on such a processor. Route the default # through the inner tokenizer. Older TRL lacks this block, so the sub is a # no-op there and only the pad_token_id fallback below applies. function = re.sub( r"(?m)^([ \t]*)if processing_class\.pad_token is None:\n" r"\1[ \t]+processing_class\.pad_token\s*=\s*processing_class\.eos_token\n", r"\1_unsloth_proc_tok = getattr(processing_class, 'tokenizer', processing_class)\n" r"\1if getattr(_unsloth_proc_tok, 'pad_token', None) is None:\n" r"\1 _unsloth_proc_tok.pad_token = getattr(_unsloth_proc_tok, 'eos_token', None)\n", function, count = 1, ) if "processing_class.pad_token_id" not in function: return function return function.replace("processing_class.pad_token_id", _PAD_FALLBACK) RL_FUNCTIONS["orpo_trainer"].append(orpo_trainer_processor_pad_token) RL_FUNCTIONS["cpo_trainer"].append(orpo_trainer_processor_pad_token) # Fix bare pop("push_to_hub_token") in compiled SFT/IterativeSFT trainer __init__ # On transformers 5.0+, to_dict() no longer includes push_to_hub_token, so bare pop KeyErrors def sft_trainer_push_to_hub_token(function_name, function): if function_name != "__init__": return function return function.replace( 'dict_args.pop("push_to_hub_token")', 'dict_args.pop("push_to_hub_token", None)' ) RL_FUNCTIONS["sft_trainer"].append(sft_trainer_push_to_hub_token) # Autocast precision for GRPO def grpo_trainer__prepare_inputs(function_name, function): if function_name != "_prepare_inputs": return function # Add mixed precision training function = function.replace( "with torch.inference_mode():", "with torch.inference_mode(), " "torch.amp.autocast(device_type = 'cuda', " "dtype = ((torch.float16 if os.environ.get('ACCELERATE_MIXED_PRECISION', 'fp16') == 'fp16' else torch.bfloat16) " "if not torch.is_autocast_enabled('cuda') else nullcontext())" "if os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '0' else torch.float16):", ) function = function.replace( "self.accelerator.unwrap_model(self.model)", "self.accelerator.unwrap_model(self.model, keep_fp32_wrapper = False)", ) return function RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__prepare_inputs) # Guard reload_weights and sync_weights - skip when fast inference LoRA shares weights with vLLM # https://github.com/huggingface/trl/commit/7856d3b1f6518601732f489883b341bb6dd36434#diff-964e6fd373aa93037604064cb2b822d7f8e2735e33f791065acf2c4c3552d393R1168-R1169 def _guard_vllm_sync_reload_for_shared_weights(function): # Guard reload_weights - only call when not sharing weights with vLLM reload_weights_pattern = re.compile( r"^(?P[ \t]*)self\.llm\.collective_rpc\(\s*(['\"])reload_weights\2\s*\)\s*$", re.MULTILINE, ) def replace_reload_weights_line(match): indent = match.group("indent") return ( f"{indent}if not getattr(self.llm, 'shared_weights', False):\n" f'{indent} self.llm.collective_rpc("reload_weights")\n' ) function = reload_weights_pattern.sub(replace_reload_weights_line, function) # Guard sync_weights - skip when sharing weights with vLLM sync_weights_block = re.compile( r"(?P[ \t]*)with profiling_context\(self,\s*(['\"])sync_weights\2\s*\):\n" r"(?P=indent)[ \t]+self\.vllm_generation\.sync_weights\(\)\n", re.MULTILINE, ) def guard_sync_weights_block(match): indent = match.group("indent") return ( f"{indent}if not getattr(getattr(self.vllm_generation, 'llm', None), 'shared_weights', False):\n" f"{indent} with profiling_context(self, 'sync_weights'):\n" f"{indent} self.vllm_generation.sync_weights()\n" ) function = sync_weights_block.sub(guard_sync_weights_block, function) return function def grpo_trainer__generate_single_turn(function_name, function): if function_name != "_generate_single_turn": return function function = _guard_vllm_sync_reload_for_shared_weights(function) # TRL 0.24.0-0.25.1 truncation regression fix # # TRL 0.22.2-0.23.1 used smart truncation via truncate_with_protected_tokens(): # - Tokenizes first without truncation # - Then truncates keeping the RIGHTMOST tokens (preserves assistant turn) # - Protects special tokens (image_token, vision_start/end) from removal # # TRL 0.24.0-0.25.1 removed this and passed kwargs directly to the tokenizer: # max_length=self.max_prompt_length, truncation=True, add_special_tokens=False # This causes issues because tokenizer truncation doesn't protect special tokens # and may not preserve the end of the prompt properly. # # TRL 0.26.2+ removed these kwargs entirely (no tokenizer-level truncation). # # Fix: Remove these kwargs so TRL 0.24.0-0.25.1 behaves like 0.26.2+ (no truncation). # This is a no-op for versions that don't have these kwargs (0.22.2-0.23.1, 0.26.2+). for pattern in [ r'["\']?max_length["\']?\s*[:=]\s*self\.max_prompt_length\s*,\s*\n?', r'["\']?truncation["\']?\s*[:=]\s*True\s*,\s*\n?', r'["\']?add_special_tokens["\']?\s*[:=]\s*False\s*,\s*\n?', ]: function = re.sub(pattern, "", function) string_to_find = " generate_inputs = super()._prepare_inputs(generate_inputs)" replacement_string = ( string_to_find + """ if "mm_token_type_ids" in generate_inputs or "image_grid_thw" in generate_inputs: mm_token_type_ids = _unsloth_fix_mm_token_type_ids( self.processing_class, generate_inputs["input_ids"], generate_inputs.get("mm_token_type_ids", None), ) if mm_token_type_ids is not None: generate_inputs["mm_token_type_ids"] = mm_token_type_ids""" ) function = function.replace(string_to_find, replacement_string) return function RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__generate_single_turn) def grpo_trainer__generate(function_name, function): if function_name != "_generate": return function return _guard_vllm_sync_reload_for_shared_weights(function) RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__generate) # Fix incorrect special tokens handling and truncation in older TRL versions def grpo_trainer__generate_and_score_completions(function_name, function): if function_name != "_generate_and_score_completions": return function # TRL 0.19.0 did skip_special_tokens = True which should be False function = function.replace( "prompt_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False", "prompt_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False", ) # Left pad prompt before calculation old and ref hidden states line_to_replace = 'batch_size = self.args.per_device_train_batch_size if mode == "train" else self.args.per_device_eval_batch_size' # The new multi-line string that will replace the line above replacement_lines = """ max_left_pad = None batch_size = self.args.per_device_train_batch_size if mode == "train" else self.args.per_device_eval_batch_size try: # TRL 0.23.1 and below path if not has_images: # Left pad prompt before calculation old and ref hidden states left_pad_tokens_per_prompt = calculate_pad_tokens_in_prompt(prompt_completion_ids, logits_to_keep, self.processing_class.pad_token_id) max_left_pad = torch.max(left_pad_tokens_per_prompt).item() except: # TRL 0.24.0 and below path if images is None: # Left pad prompt before calculation old and ref hidden states left_pad_tokens_per_prompt = calculate_pad_tokens_in_prompt(prompt_completion_ids, logits_to_keep, self.processing_class.pad_token_id) max_left_pad = torch.max(left_pad_tokens_per_prompt).item() _use_gc = self.model._unsloth_gradient_checkpointing if hasattr(self.model, '_unsloth_gradient_checkpointing') else getattr(self.args, 'gradient_checkpointing', True) self.model.for_training(use_gradient_checkpointing=_use_gc)""" function = function.replace(line_to_replace, replacement_lines) pattern_to_find = re.compile( r"^\s*if self\.args\.gradient_accumulation_steps % generate_every != 0 or \(\s*" r"self\.use_vllm and self\.vllm_importance_sampling_correction\s*" r"\):", re.MULTILINE, ) replacement_text = """ if self.args.gradient_accumulation_steps % generate_every != 0 or ( self.use_vllm ):""" # Use re.sub() to perform the replacement function, num_replacements = pattern_to_find.subn(replacement_text, function) pattern_to_find = re.compile( r"(^\s*)all_logprobs = \[" # Capture indentation (group 1) r".*?" # Match everything inside non-greedily r"for output in outputs\.outputs\s*" r"\]", re.DOTALL | re.MULTILINE, ) # sanitize_logprob is injected as a module-level function via RLTrainer_replacement # template in rl.py (from RL_REPLACEMENTS), so just reference it directly here. replacement_text = ( r"\1all_logprobs = [\n" r"\1 [sanitize_logprob(next(iter(logprob.values()))) for logprob in output.logprobs]\n" r"\1 for outputs in all_outputs\n" r"\1 for output in outputs.outputs\n" r"\1]" ) function, num_replacements = pattern_to_find.subn(replacement_text, function) # Always between max_prompt_length and use_vllm found = re.findall( r"\n(([ ]{8,})if self\.max_prompt_length is not None:.*?\2if self\.use_vllm:)", function, flags = re.DOTALL | re.MULTILINE, ) if len(found) != 0: replace_part, spacing = found[0] removed_comments = re.sub(r"\#[^\n]{1,}", "", replace_part) splits = removed_comments.split("\n") if ( sum(re.match(rf"{spacing}[^\s]", x) is not None for x in splits) == 2 and len(spacing) >= 8 ): new_replacement = f"""\n{spacing}if self.max_prompt_length is not None: # If max_prompt_length is set, we trim the prompt to keep only the last `max_prompt_length` tokens. # Then we decode those tokens back into text. We manually remove leading pad tokens from the decoded text, # because we can't use `skip_special_tokens=True` (some special tokens are still needed for generation). protected = [self.image_token_id, self.vision_start_token_id, self.vision_end_token_id] protected = [token for token in protected if token is not None] prompt_ids, prompt_mask = truncate_with_protected_tokens( prompt_ids, prompt_mask, self.max_prompt_length, protected ) prompts_text = [re.sub(rf"^({{re.escape(self.pad_token)}})+", "", text) for text in prompts_text] # The chat template inserts a single image token into the prompt text. However, when this text is later # tokenized, the single image token string is expanded into multiple image token IDs, depending on the # image size. Since we're detokenizing here, we may see repeated image tokens in the decoded text. We # collapse them back into a single token string to match the original template. if self.image_token is not None: prompts_text = [ re.sub(rf"({{re.escape(self.image_token)}})+", self.image_token, text) for text in prompts_text ] # Generate completions using either vLLM or regular generation if self.use_vllm:""" function = function.replace(replace_part, new_replacement) # Important note: we disable TRL's importance sampling logic # It is disabled because the LLM path moves left padding to the right. # We must adjust the vLLM sampling_logprob tensor in Unsloth to account for this. string_to_find = "if self.use_vllm and self.vllm_importance_sampling_correction:" replacement_string = "if False and self.use_vllm and self.vllm_importance_sampling_correction:" function = function.replace(string_to_find, replacement_string) string_to_find = """ if "image_sizes" in prompt_inputs: output["image_sizes"] = prompt_inputs["image_sizes"]""" replacement_string = """ if "image_sizes" in prompt_inputs: output["image_sizes"] = prompt_inputs["image_sizes"] if max_left_pad is not None: output["max_left_pad"] = torch.tensor(prompt_ids.shape[0] * [max_left_pad]).unsqueeze(-1) try: if self.use_vllm and getattr(self, "vllm_importance_sampling_correction", False): output["sampling_per_token_logps"] = sampling_per_token_logps except NameError: output["sampling_per_token_logps"] = None""" function = function.replace(string_to_find, replacement_string) # TRL 0.24.0+ extracts prompts = [x["prompt"] for x in inputs], losing metadata # like reasoning_effort. Inject code to store per-sample chat_template_kwargs on self. _metadata_extraction = ( "\n" " # Unsloth: Extract per-sample chat_template_kwargs before metadata is lost\n" " _ct_ = getattr(self.processing_class, 'chat_template', None) or ''\n" " _sk_ = {'prompt', 'chosen', 'rejected', 'completion', 'messages', 'label',\n" " 'images', 'image', 'videos', 'video', 'audios', 'audio'}\n" " self._unsloth_batch_chat_kwargs = []\n" " for _inp_ in inputs:\n" " _kw_ = {}\n" " if isinstance(_inp_, dict):\n" " for _k_ in _inp_.keys() - _sk_:\n" " if _k_ in _ct_ and isinstance(_inp_[_k_], str):\n" " _kw_[_k_] = _inp_[_k_]\n" " self._unsloth_batch_chat_kwargs.append(_kw_)\n" ) # Insert after: prompts = [x["prompt"] for x in inputs] _target_line = 'prompts = [x["prompt"] for x in inputs]' if _target_line in function: function = function.replace( _target_line, _target_line + _metadata_extraction, ) # This path is for TRL 0.24.0 images is a variable exclusive to this version string_to_find = """ if images is not None: output["num_images"] = num_images""" replacement_string = """ if images is not None: output["num_images"] = num_images if max_left_pad is not None: output["max_left_pad"] = torch.tensor(prompt_ids.shape[0] * [max_left_pad]).unsqueeze(-1) try: if self.use_vllm and getattr(self, "vllm_importance_sampling_correction", False): output["sampling_per_token_logps"] = sampling_per_token_logps except NameError: output["sampling_per_token_logps"] = None""" function = function.replace(string_to_find, replacement_string) if trl_version >= Version("0.24.0"): # We replace the call using 'completions' with one using 'completions_text' string_to_find = " rewards_per_func = self._calculate_rewards(inputs, prompts, completions, completion_ids_list)" replacement_string = ( " if images is not None:\n" " rewards_per_func = self._calculate_rewards(inputs, prompts_text, completions_text, completion_ids_list)\n" " else:\n" " rewards_per_func = self._calculate_rewards(inputs, prompts, completions, completion_ids_list)" ) function = function.replace(string_to_find, replacement_string) _generate_return = """ ) = self._generate(prompts)""" if _generate_return in function and "_unsloth_clear_stateful_mrope" not in function: function = function.replace( _generate_return, _generate_return + """ _unsloth_clear_stateful_mrope( self.accelerator.unwrap_model(self.model, keep_fp32_wrapper = False) )""", ) if "wake_up()" not in function: # Sleep functionality has been added to trl in v0.23.0. We do not want to redo this. # https://github.com/huggingface/trl/commit/edbe8234bc7e528f72ac76607de9d3e4753e2709 pattern = re.compile(r".*self\.llm\.generate\(.*\).*", re.MULTILINE) matches = list(pattern.finditer(function)) patched = function # Generally there's only one match. But this is just to make sure we don't miss any. for match in reversed(matches): line = match.group(0) indent_match = re.match(r"(\s*)", line) indent = indent_match.group(1) if indent_match else "" wrapped = ( f"{indent}if hasattr(self, 'llm'):\n" f"{indent} if getattr(self.llm.llm_engine.vllm_config.model_config, 'enable_sleep_mode', False):\n" f"{indent} self.llm.wake_up()\n" f"{line}\n\n" f"{indent}if hasattr(self, 'llm'):\n" f"{indent} if getattr(self.llm.llm_engine.vllm_config.model_config, 'enable_sleep_mode', False):\n" f"{indent} self.llm.sleep(os.environ.get('VLLM_SLEEP_MODE', 1))\n" ) patched = patched[: match.start()] + wrapped + patched[match.end() :] function = patched _mm_alignment = """ if "mm_token_type_ids" in forward_kwargs or "image_grid_thw" in forward_kwargs: _mm_token_type_ids = _unsloth_fix_mm_token_type_ids( self.processing_class, prompt_completion_ids, forward_kwargs.get("mm_token_type_ids", None), completion_ids = completion_ids, ) if _mm_token_type_ids is not None: forward_kwargs["mm_token_type_ids"] = _mm_token_type_ids """ _tool_image_marker = ( " # For VLM tool images: build token type IDs from the full prompt_completion_ids." ) if _tool_image_marker in function: function = function.replace(_tool_image_marker, _mm_alignment + "\n" + _tool_image_marker) else: _tt_search = ( 'if "token_type_ids" in forward_kwargs:\n' ' token_type_ids = forward_kwargs["token_type_ids"]\n' ' forward_kwargs["token_type_ids"] = torch.cat(\n' " [token_type_ids, token_type_ids.new_zeros(completion_ids.shape)], dim=1\n" " )" ) function = function.replace(_tt_search, _tt_search + "\n" + _mm_alignment.rstrip()) _save_search = ( 'if "token_type_ids" in forward_kwargs:\n' ' output["token_type_ids"] = forward_kwargs["token_type_ids"]' ) if 'output["mm_token_type_ids"]' not in function: _save_replace = ( _save_search + "\n" ' if "mm_token_type_ids" in forward_kwargs:\n' ' output["mm_token_type_ids"] = forward_kwargs["mm_token_type_ids"]' ) function = function.replace(_save_search, _save_replace) if re.search(r"\btool_mask\b", function) and 'output["tool_mask"]' not in function: function = function.replace( " return output", " if tool_mask is not None:\n" ' output["tool_mask"] = tool_mask\n' " return output", ) return function RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__generate_and_score_completions) # Fix {"reasoning_effort" : "high"} not applied def grpo_trainer_fix_maybe_apply_chat_template(function_name, function): spaces = function.find("def ") if spaces % 4 != 0: return function spaces += 4 replacement = """ _chat_template_ = getattr(self.processing_class, "chat_template", None) if _chat_template_ is None: _chat_template_ = "" _supported_keys_ = set(("prompt", "chosen", "rejected", "completion", "messages", "label")) _batch_chat_kwargs_ = getattr(self, "_unsloth_batch_chat_kwargs", None) prompts_text = [] for _idx_, _example_ in enumerate(__INPUTS__REPLACEMENT__): _tokenizer_kwargs_ = {} if type(_example_) is not dict: _example_ = {"prompt": _example_} _left_keys_ = _example_.keys() - _supported_keys_ for k in _left_keys_: if k in _chat_template_: v = _example_[k] if type(v) is str: _tokenizer_kwargs_[k] = v if _batch_chat_kwargs_ is not None and _idx_ < len(_batch_chat_kwargs_): for _bk_, _bv_ in _batch_chat_kwargs_[_idx_].items(): if _bk_ not in _tokenizer_kwargs_: _tokenizer_kwargs_[_bk_] = _bv_ _x_ = maybe_apply_chat_template(_example_, self.processing_class, **_tokenizer_kwargs_)["prompt"] prompts_text.append(_x_) """ replacement = textwrap.dedent(replacement).strip() replacement = textwrap.indent(replacement, spaces * " ") replacement = f"\n{replacement}\n" what = 'prompts_text = [maybe_apply_chat_template(example, self.processing_class)["prompt"] for example in inputs]' function = function.replace(what, replacement.replace("__INPUTS__REPLACEMENT__", "inputs")) """prompts_text = [ maybe_apply_chat_template({"prompt": prompt}, self.processing_class)["prompt"] for prompt in prompts ]""" function = re.sub( r"prompts_text = \[" r"[\s]{0,}" r"maybe_apply_chat_template\(\{[\"\']prompt[\"\'][\s]{0,}\:[\s]{0,}prompt[\s]{0,}\}[\s]{0,}\,[\s]{0,}self\.processing_class\)" r"\[[\"\']prompt[\"\']\] for prompt in prompts" r"[\s]{0,}" r"\]", replacement.replace("__INPUTS__REPLACEMENT__", "prompts"), function, ) return function RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer_fix_maybe_apply_chat_template) # Remove _move_model_to_vllm def grpo_trainer__move_model_to_vllm(function_name, function): if function_name != "_move_model_to_vllm": return function def _move_model_to_vllm(self, *args, **kwargs): return None function = inspect.getsource(_move_model_to_vllm) return function RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__move_model_to_vllm) # Edit _get_per_token_logps to handle mixed precision def grpo_trainer__get_per_token_logps(function_name, function): if function_name != "_get_per_token_logps": return function def _get_per_token_logps( self, model, input_ids, attention_mask, logits_to_keep, compute_efficient = False, ): if True: # os.environ.get('UNSLOTH_USE_NEW_MODEL', '0') == '0': return None # Unsloth efficient GRPO # Otherwise, calculate normally: if not hasattr(self, "_autocast_dtype"): self._autocast_dtype = ( torch.float16 if os.environ.get("ACCELERATE_MIXED_PRECISION", "fp16") == "fp16" else torch.bfloat16 ) if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1": self._autocast_dtype = torch.float16 os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" with torch.amp.autocast(device_type = DEVICE_TYPE, dtype = self._autocast_dtype): # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded logits = model( input_ids = input_ids, attention_mask = attention_mask, logits_to_keep = logits_to_keep + 1, ).logits # logits = logits[:, :-1, :] # (B, L-1, V), exclude the last logit: it corresponds to the next token pred return logits # input_ids = input_ids[:, -logits_to_keep:] # For transformers<=4.48, logits_to_keep argument isn't supported, so here we drop logits ourselves. # See https://github.com/huggingface/trl/issues/2770 # logits = logits[:, -logits_to_keep:] # return logits # See https://huggingface.co/blog/the_n_implementation_details_of_rlhf_with_ppo#policy-training-implementation-details # logits = logits / self.temperature # logps = selective_log_softmax(logits, input_ids) # row_indices, col_indices = torch.where(logps < -20) # # Method 1: Check if tensors have elements # if len(row_indices) > 0 and len(col_indices) > 0: # breakpoint() # Breakpoint triggered here # print("Found high values!") # return logps # compute logprobs for the input tokens function = inspect.getsource(_get_per_token_logps) return function RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__get_per_token_logps) def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): if function_name != "_get_per_token_logps_and_entropies": return function # Just copy over from _get_per_token_logps replacement function above. For now this returns None anyway def _get_per_token_logps_and_entropies( self, model, input_ids, attention_mask, logits_to_keep, batch_size = None, compute_entropy = False, compute_efficient = False, *args, **kwargs, ): # All Unsloth code here in this function is licensed under AGPL3 # if True: # os.environ.get('UNSLOTH_USE_NEW_MODEL', '0') == '0': # return None, None # logps, entropies Unsloth efficient GRPO if compute_efficient: return None, None else: if not hasattr(self, "_autocast_dtype"): self._autocast_dtype = ( torch.float16 if os.environ.get("ACCELERATE_MIXED_PRECISION", "fp16") == "fp16" else torch.bfloat16 ) if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1": self._autocast_dtype = torch.float16 compute_aux_loss = kwargs.get("compute_aux_loss", None) pixel_values, image_grid_thw = ( kwargs.get("pixel_values", None), kwargs.get("image_grid_thw", None), ) pixel_attention_mask, image_sizes = ( kwargs.get("pixel_attention_mask", None), kwargs.get("image_sizes", None), ) num_images = kwargs.get("num_images", None) # Transformers 5.x needs token_type_ids/mm_token_type_ids for some vision models token_type_ids = kwargs.get("token_type_ids", None) mm_token_type_ids = kwargs.get("mm_token_type_ids", None) if mm_token_type_ids is not None or image_grid_thw is not None: mm_token_type_ids = _unsloth_fix_mm_token_type_ids( self.processing_class, input_ids, mm_token_type_ids ) unwrapped_model = self.accelerator.unwrap_model(model, keep_fp32_wrapper = False) lm_head = self.model.get_output_embeddings().weight dtype_bytes = 16 if self._autocast_dtype in [torch.float16, torch.bfloat16] else 32 total_rows = input_ids.shape[0] seq_len = input_ids.shape[1] hidden_dim = lm_head.shape[1] vocab_dim = lm_head.shape[0] if self.args.unsloth_grpo_mini_batch is None: B, multiplier = autotune_batch_and_chunks( total_rows, seq_len, hidden_dim, vocab_dim, dtype_bytes, self.args.unsloth_logit_chunk_multiplier, ) B = total_rows // B else: B = self.args.unsloth_grpo_mini_batch if self.args.unsloth_logit_chunk_multiplier is None: multiplier = max(4, seq_len // 4096) else: multiplier = self.args.unsloth_logit_chunk_multiplier all_logprobs_list = [] if pixel_values is None: left_pad_tokens_per_prompt = calculate_pad_tokens_in_prompt( input_ids, logits_to_keep, self.processing_class.pad_token_id ) max_left_pad = torch.max(left_pad_tokens_per_prompt).item() input_ids = left_pack_padding(input_ids, self.processing_class.pad_token_id) attention_mask = input_ids != self.processing_class.pad_token_id attention_mask = attention_mask.to(attention_mask.dtype) else: max_left_pad = 0 def slice_sample_axis(value, start, end): if value is None: return None return value[start:end] import math total_samples = input_ids.shape[0] batch_size = math.ceil(total_samples / B) if isinstance(num_images, torch.Tensor): num_images = num_images.detach().cpu().reshape(-1).tolist() if image_grid_thw is not None and pixel_values is not None and num_images 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]) # why: cum_rows is indexed via .item() inside the per-chunk loop; # keeping it on CPU avoids per-iteration GPU->CPU sync. cum_rows = torch.cat( [ torch.tensor([0], device = rows_per_sample.device), rows_per_sample.cumsum(0), ] ).cpu() cum_imgs = torch.tensor([0] + num_images).cumsum(0) else: cum_rows = None cum_imgs = None def _first_dim_len(value): if value is None: return None if hasattr(value, "shape"): return value.shape[0] try: return len(value) except TypeError: return None total_images = sum(num_images) if num_images is not None else None _image_sizes_n = _first_dim_len(image_sizes) input_ids_chunks = [] attention_mask_chunks = [] pixel_values_chunks = [] image_grid_thw_chunks = [] pixel_attention_mask_chunks = [] image_sizes_chunks = [] token_type_ids_chunks = [] mm_token_type_ids_chunks = [] current_pixel_idx = 0 # TRL 0.23.0 batching logic for start in range(0, total_samples, batch_size): end = min(start + batch_size, total_samples) input_ids_chunks.append(input_ids[start:end]) attention_mask_chunks.append(attention_mask[start:end]) token_type_ids_chunks.append(slice_sample_axis(token_type_ids, start, end)) mm_token_type_ids_chunks.append(slice_sample_axis(mm_token_type_ids, start, end)) if image_grid_thw is not None and pixel_values is not None: if num_images is None: grid_slice = image_grid_thw[start:end] batch_pixel_count = grid_slice.prod(dim = -1).sum().item() start_pixel_idx = current_pixel_idx end_pixel_idx = current_pixel_idx + batch_pixel_count current_pixel_idx = end_pixel_idx img_start = img_end = None else: start_pixel_idx = cum_rows[start].item() end_pixel_idx = cum_rows[end].item() img_start = cum_imgs[start].item() img_end = cum_imgs[end].item() grid_slice = image_grid_thw[img_start:img_end] image_grid_thw_chunks.append(grid_slice) pixel_values_chunks.append(pixel_values[start_pixel_idx:end_pixel_idx]) if image_sizes is None: image_sizes_chunks.append(None) elif ( num_images is not None and _image_sizes_n == total_images and img_start is not None ): image_sizes_chunks.append(image_sizes[img_start:img_end]) else: image_sizes_chunks.append(slice_sample_axis(image_sizes, start, end)) if pixel_attention_mask is None: pixel_attention_mask_chunks.append(None) elif ( num_images is not None and img_start is not None and pixel_attention_mask.shape[0] == image_grid_thw.shape[0] ): pixel_attention_mask_chunks.append(pixel_attention_mask[img_start:img_end]) elif ( pixel_attention_mask.shape[0] == pixel_values.shape[0] and pixel_attention_mask.shape[0] != input_ids.shape[0] ): pixel_attention_mask_chunks.append( pixel_attention_mask[start_pixel_idx:end_pixel_idx] ) else: pixel_attention_mask_chunks.append(pixel_attention_mask[start:end]) else: pixel_values_chunks.append(None) image_grid_thw_chunks.append(None) pixel_attention_mask_chunks.append(None) image_sizes_chunks.append(slice_sample_axis(image_sizes, start, end)) temperature = self.temperature model_config = _unsloth_get_model_config(model) logit_softcapping = _unsloth_get_final_logit_softcapping(model) logit_scale_multiply = getattr(model_config, "logit_scale", 0) if logit_scale_multiply is None: logit_scale_multiply = 0 logit_scale_divide = getattr(model_config, "logits_scaling", 0) if logit_scale_divide is None: logit_scale_divide = 0 zipped_inputs = zip( input_ids_chunks, attention_mask_chunks, pixel_values_chunks, image_grid_thw_chunks, pixel_attention_mask_chunks, image_sizes_chunks, token_type_ids_chunks, mm_token_type_ids_chunks, ) os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" # ---- Sequence packing (default-on; disable with UNSLOTH_GRPO_SEQ_PACKING=0) ---- # One varlen [1, sum L] forward replaces the padded [B, Lmax] loop (also fixes the # left-pad RoPE error). Self-verified against the per-row forward, re-checked as T # grows; falls back if a backend ignores packed_seq_lengths. logprobs = None # ---- PrefixGrouper (GRPO shared-prompt dedup; default ON, exact + self-verified) ---- # G completions per prompt share the prefix; the packed path forwards it G times, # PrefixGrouper stores it once (FlexAttention shared-prefix mask), cutting the trunk # forward from G*(P+R) to P+G*R tokens. Gated by UNSLOTH_GRPO_PREFIX_GROUPER (needs # seq-packing), tok_r auto-gate, and first-use self-verify vs the packed path # (mismatch => fall back + mark unsafe), so a mask/isolation regression cannot ship # silently. When off / ungrouped / unverified, the packed path below runs as before. _pg_result = None _pg_use = False _pg_skip_pk = False # once a shape is PG-verified, skip the full-row forward _pg_forward_fn = None # deferred PG forward (runs at the verify site below) _pg_num_gen = getattr(self, "num_generations", None) # Env gate hoisted to module level (mirrored via RL_PRE_ITEMS). Skip PG under vLLM # (fast_inference=True): the rollout dominates the step, so PG saves little and its # first-use self-verify is net overhead. _pg_engage = ( UNSLOTH_GRPO_PREFIX_GROUPER_ON and not getattr(self, "use_vllm", False) and not getattr(unwrapped_model, "_unsloth_prefix_grouper_nograd_disabled", False) ) if _pg_engage: try: # Skip softcap models (the flex kernel never applies attn_logit_softcapping) # and hybrid SSM / MoE models: only the threaded attention forwards get the # shared-prefix isolation, so a Mamba or MoE decoder that does not forward # prefix_seg_info would leak suffixes across completions. PG also rides on # sequence packing, so it needs the same zoo masked-column guard. _pg_cfg = getattr(unwrapped_model, "config", None) _pg_engage = ( _pg_enabled_fn() and UNSLOTH_ZOO_HAS_MASKED_COL_GUARD and pixel_values is None and token_type_ids is None and mm_token_type_ids is None and _pg_num_gen is not None and _pg_num_gen >= 2 and not getattr(_pg_cfg, "attn_logit_softcapping", None) # normal backends apply config.attention_dropout in training; the flex # path is deterministic, so skip PG when it is set. and not getattr(_pg_cfg, "attention_dropout", 0) and not any( getattr(_pg_cfg, _pg_a, None) is not None for _pg_a in ( "mamba_d_ssm", "mamba_d_state", "mamba_expand", "num_experts", "num_local_experts", "n_routed_experts", "moe_intermediate_size", ) ) ) except Exception: _pg_engage = False if _pg_engage: try: _pg_pad = self.processing_class.pad_token_id # cap the PG span (P+max(R)) at the sliding window, like the packed _pk_sw guard. _pg_sw = getattr( getattr(unwrapped_model, "config", None), "sliding_window", None ) if not (isinstance(_pg_sw, int) and _pg_sw > 0): _pg_sw = None _pg_layout = _pg_build_layout( input_ids, logits_to_keep, _pg_pad, _pg_num_gen, left_pad_tokens_per_prompt, max_segment_cap = _pg_sw, ) _pg_unsafe = getattr( unwrapped_model, "_unsloth_prefix_grouper_nograd_unsafe", None ) if _pg_unsafe is None: _pg_unsafe = set() if _pg_layout is not None and _pg_layout.signature not in _pg_unsafe: _pg_sig = _pg_layout.signature _pg_verified = getattr( unwrapped_model, "_unsloth_prefix_grouper_nograd_verified", None ) if _pg_verified is None: _pg_verified = set() _pg_chunks = max(1, total_rows * multiplier) def _pg_run_forward(_pg_layout = _pg_layout, _pg_chunks = _pg_chunks): with _get_inference_mode_context_manager(model): with torch.amp.autocast( device_type = "cuda", dtype = self._autocast_dtype ): _pg_hidden = unwrapped_model( input_ids = _pg_layout.flat_ids, position_ids = _pg_layout.position_ids, prefix_seg_info = _pg_layout.prefix_seg_info, use_cache = False, ).logits _pg_r = _pg_layout.extract_logps( _pg_hidden, lm_head, chunked_hidden_states_selective_log_softmax, _pg_chunks, logit_scale_multiply, logit_scale_divide, logit_softcapping, temperature, ) _pg_hidden = None # release before any verify forward device_synchronize() # clip to the loss window [B, logits_to_keep+max_left_pad] _pg_w = logits_to_keep + max_left_pad if _pg_r.shape[1] > _pg_w: _pg_r = _pg_r[:, -_pg_w:] return _pg_r # trust only within the verified envelope: re-verify when T or the # longest segment grows, like the packed path _pg_T = int(_pg_layout.flat_ids.shape[1]) _pg_maxseg = int(_pg_layout.position_ids.max()) + 1 _pg_env = ( _pg_verified.get(_pg_sig) if isinstance(_pg_verified, dict) else None ) if (not _pg_verify_on()) or ( _pg_env is not None and _pg_T <= _pg_env[0] and _pg_maxseg <= _pg_env[1] ): # trusted shape: run PG now and skip the full-row forward below _pg_result = _pg_run_forward() _pg_use = True _pg_skip_pk = True else: # unverified shape: defer the forward until the packed reference # exists (verify site below), so a declined packed path never wastes # a whole-batch PG forward _pg_forward_fn = _pg_run_forward except Exception as _pg_err: _pg_result = None _pg_use = False _pg_skip_pk = False _pg_forward_fn = None # A FlexAttention/Triton compile failure or OOM here is GPU-wide, not # layout-specific, so retrying the same PG forward every step just re-pays # the failure. Persistently disable PG (mirrors the seq-packing handler # setting _unsloth_seq_packing_nograd_ok = False); the packed/padded path # below still produces the exact result. unwrapped_model._unsloth_prefix_grouper_nograd_disabled = True if isinstance(_pg_err, torch.cuda.OutOfMemoryError): torch.cuda.empty_cache() os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" if UNSLOTH_ENABLE_LOGGING: print( f"[Unsloth] GRPO PrefixGrouper (no-grad) disabled (fell back to packed): {_pg_err!r}", flush = True, ) # ---- Sequence packing (default-on; disable with UNSLOTH_GRPO_SEQ_PACKING=0) ---- # One varlen [1, sum L] block-diagonal forward replaces the padded [B, Lmax] loop # (exact per-row result; also fixes the padded path's left-pad RoPE error). # Self-verified vs the per-row forward, re-checked as T grows; falls back if a # backend ignores packed_seq_lengths. lm_head runs on completion positions only. _pk_result = None _pk_use = False _pk_enabled = UNSLOTH_GRPO_SEQ_PACKING_ON # Without zoo#840's masked-column guard, zeroed prompt/pad columns turn NaN in exp(). _pk_enabled = _pk_enabled and UNSLOTH_ZOO_HAS_MASKED_COL_GUARD _pk_ok = getattr(unwrapped_model, "_unsloth_seq_packing_nograd_ok", None) if ( _pk_enabled and not _pg_skip_pk and pixel_values is None and token_type_ids is None and mm_token_type_ids is None and _pk_ok is not False ): try: _pk_pad = self.processing_class.pad_token_id _pk_keep = input_ids != _pk_pad _pk_len = _pk_keep.sum(dim = 1) _pk_len_cpu = _pk_len.tolist() # single GPU->CPU sync, reused below _pk_nz_cpu = [_n for _n in _pk_len_cpu if _n > 0] _pk_flat = input_ids[_pk_keep].unsqueeze(0) _pk_T = _pk_flat.shape[1] _pk_L = input_ids.shape[1] _pk_W = logits_to_keep + max_left_pad _pk_maxseg = max(_pk_nz_cpu) if _pk_nz_cpu else 0 # sliding-window models lose the per-sequence local window in a packed stream _pk_sw = getattr( getattr(unwrapped_model, "config", None), "sliding_window", None ) _pk_sw_ok = not (isinstance(_pk_sw, int) and _pk_sw > 0 and _pk_maxseg > _pk_sw) # per-row completion mask (same as the loss); prompt-only rows count as inactive _pk_cmask = create_completion_attention_mask( input_ids[:, -_pk_W:], left_pad_tokens_per_prompt, max_left_pad, _pk_pad ) _pk_active = int(_pk_cmask.any(dim = 1).sum()) # skip the packed forward entirely at known-unsafe lengths (avoids a wasted pass / OOM) _pk_unsafe = getattr( unwrapped_model, "_unsloth_seq_packing_nograd_unsafe_T", None ) # cap the flattened forward at one padded [batch_size, seq_len] mini-batch's # token budget; anything larger uses the chunked padded loop _pk_cap = batch_size * seq_len if ( _pk_T >= 2 and _pk_T <= _pk_cap and len(_pk_nz_cpu) > 0 and _pk_sw_ok and not (_pk_unsafe is not None and _pk_T >= _pk_unsafe) and (_pk_ok is True or _pk_active >= 2) ): # reset 0-based position_ids per segment _pk_pos = (_pk_keep.cumsum(dim = 1) - 1)[_pk_keep].unsqueeze(0) _pk_chunks = max(1, total_rows * multiplier) _pk_nz_idx = _pk_keep.nonzero( as_tuple = False ) # [T, 2] = (row, col), row-major _pk_within = _pk_nz_idx[1:, 0] == _pk_nz_idx[:-1, 0] # [T-1] # per-row completion start after left-packing (matches create_completion_attention_mask) _pk_cstart = (_pk_L - logits_to_keep) - left_pad_tokens_per_prompt # [rows] _pk_ctgt = (_pk_nz_idx[1:, 1] >= _pk_cstart[_pk_nz_idx[1:, 0]]) & _pk_within with _get_inference_mode_context_manager(model): with torch.amp.autocast(device_type = "cuda", dtype = self._autocast_dtype): # use_cache=False: a KV cache silently disables varlen packing _pk_hidden = unwrapped_model( input_ids = _pk_flat, position_ids = _pk_pos, packed_seq_lengths = torch.tensor( _pk_nz_cpu, dtype = torch.int32, device = input_ids.device ), use_cache = False, ).logits _pk_sel = chunked_hidden_states_selective_log_softmax( _pk_hidden[0, :-1, :][_pk_ctgt].unsqueeze(0), lm_head, _pk_flat[0, 1:][_pk_ctgt].unsqueeze(0), _pk_chunks, logit_scale_multiply, logit_scale_divide, logit_softcapping, temperature, )[0] # GPT-OSS offload race guard (matches the padded loop) device_synchronize() # scatter each logprob back to its (row, col) so [:, -_pk_W:] matches padded _pk_tgt = (_pk_nz_idx[1:, 0] * _pk_L + _pk_nz_idx[1:, 1])[_pk_ctgt] _pk_result = ( torch.zeros( total_rows * _pk_L, dtype = torch.float32, device = input_ids.device, ) .index_put((_pk_tgt,), _pk_sel.to(torch.float32)) .view(total_rows, _pk_L)[:, -_pk_W:] ) # re-verify when T or the longest segment grows past what was verified # (a LongRoPE cache switch can change the result) _pk_vT = int( getattr(unwrapped_model, "_unsloth_seq_packing_nograd_verified_T", 0) ) _pk_vS = int( getattr(unwrapped_model, "_unsloth_seq_packing_nograd_verified_seg", 0) ) # debug: hand-edit this condition to force re-verify every step if _pk_ok is True and _pk_T <= _pk_vT and _pk_maxseg <= _pk_vS: _pk_use = True # already verified for this shape else: # verify against the per-row forward (ground truth) _pk_ref = torch.zeros_like(_pk_result) with _get_inference_mode_context_manager(model): with torch.amp.autocast( device_type = "cuda", dtype = self._autocast_dtype ): for _pk_i in range(total_rows): _pk_ni = _pk_len_cpu[_pk_i] if _pk_ni < 2: continue _pk_rmask = _pk_keep[_pk_i] _pk_real = input_ids[_pk_i][_pk_rmask].unsqueeze(0) _pk_rpos = torch.arange( _pk_ni, device = input_ids.device ).unsqueeze(0) _pk_rh = unwrapped_model( input_ids = _pk_real, position_ids = _pk_rpos, use_cache = False, ).logits _pk_rsel = chunked_hidden_states_selective_log_softmax( _pk_rh[:, :-1, :], lm_head, _pk_real[:, 1:], 1, logit_scale_multiply, logit_scale_divide, logit_softcapping, temperature, )[0] _pk_rcols = _pk_rmask.nonzero(as_tuple = False).squeeze(1)[ 1: ] - (_pk_L - _pk_W) _pk_rkeep = _pk_rcols >= 0 _pk_ref[_pk_i, _pk_rcols[_pk_rkeep]] = _pk_rsel[ _pk_rkeep ].to(torch.float32) device_synchronize() # compare over the loss-mask region only _pk_cm = _pk_cmask.float() _pk_diff = float(((_pk_result - _pk_ref).abs() * _pk_cm).max()) if UNSLOTH_ENABLE_LOGGING: print( f"[Unsloth] GRPO seq-packing (no-grad) verify: T={_pk_T} maxseg={_pk_maxseg} packed-vs-perrow max|d|={_pk_diff:.4f}", flush = True, ) # kernel-noise floor ~0.25; cross-sample contamination is >= 2.4 if _pk_diff < 7e-1: unwrapped_model._unsloth_seq_packing_nograd_ok = True # widen the trusted shape only when >= 2 completion rows exercised # cross-sample packing; single-row passes prove nothing if _pk_active >= 2: unwrapped_model._unsloth_seq_packing_nograd_verified_T = max( _pk_vT, _pk_T ) unwrapped_model._unsloth_seq_packing_nograd_verified_seg = max( _pk_vS, _pk_maxseg ) _pk_ok = True _pk_use = True else: _pk_use = False if _pk_diff >= 1.5: # contamination (attention ignores the packed mask): disable packing unwrapped_model._unsloth_seq_packing_nograd_ok = False else: # likely a length boundary (LongRoPE): mark unsafe, keep smaller shapes unwrapped_model._unsloth_seq_packing_nograd_unsafe_T = ( _pk_T if _pk_unsafe is None else min(_pk_unsafe, _pk_T) ) if UNSLOTH_ENABLE_LOGGING: print( f"[Unsloth] GRPO seq-packing (no-grad) fell back at T={_pk_T} (diff={_pk_diff:.3f})", flush = True, ) except Exception as _pk_err: # any failure: drop intermediates, use the padded loop, do not retry _pk_hidden = None _pk_sel = None _pk_result = None _pk_use = False if isinstance(_pk_err, torch.cuda.OutOfMemoryError): torch.cuda.empty_cache() unwrapped_model._unsloth_seq_packing_nograd_ok = False if UNSLOTH_ENABLE_LOGGING: print( f"[Unsloth] GRPO sequence-packing (no-grad) disabled (fell back to padded): {_pk_err!r}", flush = True, ) # ---- PrefixGrouper first-use self-verify (no-grad) ---- # Compare the untrusted PG result to the full-row packed result (itself verified vs # per-row) over the completion mask: < tol_ok -> trust the structure; >= TOL_KILL -> # unsafe forever; borderline -> fall back this shape. if _pg_forward_fn is not None and not _pg_use: if _pk_use and _pk_result is not None: try: # deferred PG forward, run only now that the packed reference exists _pg_result = _pg_forward_fn() _pg_W2 = logits_to_keep + max_left_pad _pg_cm = create_completion_attention_mask( input_ids[:, -_pg_W2:], left_pad_tokens_per_prompt, max_left_pad, self.processing_class.pad_token_id, ).float() _pg_a = _pg_result[:, -_pg_W2:].float() _pg_b = _pk_result[:, -_pg_W2:].float() _pg_diff = float(((_pg_a - _pg_b).abs() * _pg_cm).max()) if UNSLOTH_ENABLE_LOGGING: print( f"[Unsloth] GRPO PrefixGrouper (no-grad) verify: sig={_pg_layout.signature} " f"shared-prefix vs full-row-packed max|d|={_pg_diff:.4f}", flush = True, ) if _pg_diff < _pg_tol_ok(): _pg_v = getattr( unwrapped_model, "_unsloth_prefix_grouper_nograd_verified", None ) if not isinstance(_pg_v, dict): _pg_v = {} _pg_vT = int(_pg_layout.flat_ids.shape[1]) _pg_vS = int(_pg_layout.position_ids.max()) + 1 _pg_old = _pg_v.get(_pg_layout.signature, (0, 0)) _pg_v[_pg_layout.signature] = ( max(_pg_vT, _pg_old[0]), max(_pg_vS, _pg_old[1]), ) unwrapped_model._unsloth_prefix_grouper_nograd_verified = _pg_v _pg_use = True else: _pg_u = getattr( unwrapped_model, "_unsloth_prefix_grouper_nograd_unsafe", None ) if _pg_u is None: _pg_u = set() if _pg_diff >= _PG_TOL_KILL: _pg_u.add(_pg_layout.signature) unwrapped_model._unsloth_prefix_grouper_nograd_unsafe = _pg_u _pg_use = False except Exception as _pg_err3: _pg_result = None _pg_use = False if isinstance(_pg_err3, torch.cuda.OutOfMemoryError): torch.cuda.empty_cache() os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" if UNSLOTH_ENABLE_LOGGING: print( f"[Unsloth] GRPO PrefixGrouper (no-grad) verify failed (fell back to packed): {_pg_err3!r}", flush = True, ) # else: no packed reference (packing off/failed) -> cannot verify; fall back. if _pg_use and _pg_result is not None: logprobs = _pg_result # PrefixGrouper verified/trusted -> skip the loop zipped_inputs = [] elif _pk_use and _pk_result is not None: logprobs = _pk_result # verified -> skip the loop zipped_inputs = [] else: # free packed intermediates before running the padded loop _pk_hidden = _pk_sel = _pk_result = _pk_ref = None with _get_inference_mode_context_manager(model): for ( input_ids_chunk, attention_mask_chunk, pixel_values_chunk, image_grid_thw_chunk, pixel_attention_mask_chunk, image_sizes_chunk, token_type_ids_chunk, mm_token_type_ids_chunk, ) in zipped_inputs: _extra_vision_kwargs = {} if token_type_ids_chunk is not None: _extra_vision_kwargs["token_type_ids"] = token_type_ids_chunk if mm_token_type_ids_chunk is not None: _extra_vision_kwargs["mm_token_type_ids"] = mm_token_type_ids_chunk with torch.amp.autocast(device_type = "cuda", dtype = self._autocast_dtype): if pixel_values is None: outputs = unwrapped_model( input_ids = input_ids_chunk, attention_mask = attention_mask_chunk, pixel_values = pixel_values_chunk, image_grid_thw = image_grid_thw_chunk, pixel_attention_mask = pixel_attention_mask_chunk, image_sizes = image_sizes_chunk, **_extra_vision_kwargs, ) logits_chunk = outputs.logits del outputs # free hidden_states before chunked log-softmax completion_input_ids_chunk = input_ids_chunk[ :, -(logits_to_keep + max_left_pad) : ] logits_chunk = logits_chunk[ :, -(logits_to_keep + max_left_pad + 1) :, : ] logits_chunk = logits_chunk[:, :-1, :] logprobs_chunk = chunked_hidden_states_selective_log_softmax( logits_chunk, lm_head, completion_input_ids_chunk, chunks = input_ids_chunk.shape[0] * multiplier, logit_scale_multiply = logit_scale_multiply, logit_scale_divide = logit_scale_divide, logit_softcapping = logit_softcapping, temperature = temperature, ) else: # Essentially, for VLMs we do not go via the optimized path in models/, # so we don't encounter the Flash Attn left-padding issue. outputs = unwrapped_model( input_ids = input_ids_chunk, attention_mask = attention_mask_chunk, pixel_values = pixel_values_chunk, image_grid_thw = image_grid_thw_chunk, pixel_attention_mask = pixel_attention_mask_chunk, image_sizes = image_sizes_chunk, logits_to_keep = logits_to_keep + 1, **_extra_vision_kwargs, ) logits_chunk = outputs.logits del outputs # free hidden_states before chunked log-softmax logits_chunk = logits_chunk[:, :-1, :] completion_input_ids_chunk = input_ids_chunk[:, -logits_to_keep:] # Guard: check if model returned hidden states or logits if logits_chunk.shape[-1] == lm_head.shape[1]: logprobs_chunk = chunked_hidden_states_selective_log_softmax( logits_chunk, lm_head, completion_input_ids_chunk, chunks = input_ids_chunk.shape[0] * multiplier, logit_scale_multiply = logit_scale_multiply, logit_scale_divide = logit_scale_divide, logit_softcapping = logit_softcapping, temperature = temperature, ) else: # Model returned logits directly - scaling/softcapping already applied by model forward logprobs_chunk = chunked_selective_log_softmax( logits_chunk, completion_input_ids_chunk, temperature, ) # This is needed to avoid race conditions with GPT OSS offload_embbed=True # However, it seems that this line does not slow down or disrupt models. device_synchronize() all_logprobs_list.append(logprobs_chunk) if logprobs is None: # padded fallback when packing was not used logprobs = torch.cat(all_logprobs_list, dim = 0) entropies = None os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "0" # aux loss is unused: it is off by default (router_aux_loss_coef set to 0 in models/rl.py) # and explicit opt-in is rejected at trainer init, so this is always None (kept in the # return for TRL >= 1.7.0's 3-tuple contract). aux_loss = None return logprobs.detach(), entropies, aux_loss # logps, entropies, aux_loss # input_ids = input_ids[:, -logits_to_keep:] # For transformers<=4.48, logits_to_keep argument isn't supported, so here we drop logits ourselves. # See https://github.com/huggingface/trl/issues/2770 # logits = logits[:, -logits_to_keep:] # return logits # See https://huggingface.co/blog/the_n_implementation_details_of_rlhf_with_ppo#policy-training-implementation-details # logits = logits / self.temperature # logps = selective_log_softmax(logits, input_ids) # row_indices, col_indices = torch.where(logps < -20) # # Method 1: Check if tensors have elements # if len(row_indices) > 0 and len(col_indices) > 0: # breakpoint() # Breakpoint triggered here # print("Found high values!") # return logps # compute logprobs for the input tokens function = inspect.getsource(_get_per_token_logps_and_entropies) if trl_version < Version("1.7.0"): # TRL < 1.7.0 unpacks (logps, entropies) at every call site; TRL >= 1.7.0 # always unpacks (logps, entropies, aux_loss). Drop the aux_loss element so # the return arity matches the installed TRL. Regex tolerates comment / # whitespace drift on the return line; fail loud if the anchor ever stops # matching rather than silently shipping a 3-tuple to older TRL. new_function, n = re.subn( r"return (logprobs\.detach\(\), entropies), aux_loss[^\n]*", r"return \1 # logps, entropies", function, ) if n != 1: raise RuntimeError( "Unsloth GRPO: could not downgrade the per-token-logps return to a " f"2-tuple for TRL {trl_version} (matched {n} times, expected 1). The " "return line changed; update the arity gate in rl_replacements.py." ) function = new_function return function RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__get_per_token_logps_and_entropies) def _unsloth_get_model_config(model): """Return HuggingFace model config, unwrapping DDP/Accelerate wrappers.""" config = getattr(model, "config", None) if config is None and hasattr(model, "module"): config = getattr(model.module, "config", None) return config def _unsloth_get_final_logit_softcapping(model): """Return final_logit_softcapping for a model config, falling back to the nested text sub-config for composite models. Handles both: - Gemma-4-style configs where the attribute lives on ``config.text_config`` - T5Gemma-style composite configs where the text sub-config is only reachable via ``config.get_text_config()`` Returns 0 if unset, matching the previous behaviour. """ config = _unsloth_get_model_config(model) if config is None: return 0 softcap = getattr(config, "final_logit_softcapping", None) if softcap is None: text_cfg = getattr(config, "text_config", None) if text_cfg is None: get_text_config = getattr(config, "get_text_config", None) if callable(get_text_config): try: text_cfg = get_text_config() except (TypeError, ValueError): text_cfg = None if text_cfg is not None and text_cfg is not config: softcap = getattr(text_cfg, "final_logit_softcapping", None) return 0 if softcap is None else softcap grpo_compute_loss = RL_REPLACEMENTS["grpo_compute_loss"] grpo_compute_loss_slow = RL_REPLACEMENTS["grpo_compute_loss_slow"] UnslothEfficientGRPO = RL_REPLACEMENTS["UnslothEfficientGRPO"] grpo_accumulated_loss = RL_REPLACEMENTS["grpo_accumulated_loss"] grpo_update_SamplingParams = RL_REPLACEMENTS["grpo_update_SamplingParams"] RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_get_model_config)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_get_final_logit_softcapping)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_get_mm_token_id)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_fix_mm_token_type_ids)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_clear_stateful_mrope)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(grpo_compute_loss)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(UnslothEfficientGRPO)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(grpo_accumulated_loss)) RL_PRE_ITEMS["grpo_trainer"].append(grpo_compute_loss_slow) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(grpo_update_SamplingParams)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_get_inference_mode_context_manager)) # inspect.getsource inlines function bodies but not module imports, so constants the inlined # grpo functions reference (e.g. UNSLOTH_ENABLE_LOGGING) must be redefined in the generated cache. RL_PRE_ITEMS["grpo_trainer"].append( "import os as _unsloth_os\n" "UNSLOTH_ENABLE_LOGGING = _unsloth_os.environ.get('UNSLOTH_ENABLE_LOGGING', '0') in ('1', 'True', 'true')\n" ) # Sequence-packing gates, same values as the module-top constants. RL_PRE_ITEMS["grpo_trainer"].append( "UNSLOTH_GRPO_SEQ_PACKING_ON = _unsloth_os.environ.get('UNSLOTH_GRPO_SEQ_PACKING', '1').lower() not in ('0', 'false', 'no', 'off')\n" ) RL_PRE_ITEMS["grpo_trainer"].append( "try:\n" " import inspect as _unsloth_inspect\n" " from unsloth_zoo.rl_replacements import RL_REPLACEMENTS as _unsloth_zoo_RL\n" " UNSLOTH_ZOO_HAS_MASKED_COL_GUARD = 'torch.where(_keep, new' in _unsloth_inspect.getsource(_unsloth_zoo_RL['grpo_compute_loss'])\n" "except Exception:\n" " UNSLOTH_ZOO_HAS_MASKED_COL_GUARD = False\n" ) # PrefixGrouper gate, same shape as the module-top constants. RL_PRE_ITEMS["grpo_trainer"].append( "_pg_build_layout = _pg_enabled_fn = _pg_verify_on = _pg_tol_ok = _PG_TOL_KILL = None\n" "UNSLOTH_GRPO_PREFIX_GROUPER_ON = _unsloth_os.environ.get('UNSLOTH_GRPO_PREFIX_GROUPER', '1').lower() not in ('0', 'false', 'no', 'off')\n" "if UNSLOTH_GRPO_PREFIX_GROUPER_ON:\n" " try:\n" " from unsloth.utils.prefix_grouper import build_group_layout as _pg_build_layout, prefix_grouper_enabled as _pg_enabled_fn, verify_on as _pg_verify_on, tol_ok as _pg_tol_ok, TOL_KILL as _PG_TOL_KILL\n" " except Exception:\n" " UNSLOTH_GRPO_PREFIX_GROUPER_ON = False\n" ) # Edit _get_per_token_logps to handle mixed precision def grpo_trainer_compute_loss(function_name, function): if function_name != "compute_loss": return function def compute_loss( self, model, inputs, return_outputs = False, num_items_in_batch = None, ): if return_outputs: raise ValueError("The GRPOTrainer does not support returning outputs") # Compute the per-token log probabilities for the model prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"] completion_ids, completion_mask = ( inputs["completion_ids"], inputs["completion_mask"], ) pixel_values, image_grid_thw = ( inputs.get("pixel_values", None), inputs.get("image_grid_thw", None), ) pixel_attention_mask, image_sizes = ( inputs.get("pixel_attention_mask", None), inputs.get("image_sizes", None), ) num_images = inputs.get("num_images", None) # Transformers 5.x needs token_type_ids/mm_token_type_ids for some vision models token_type_ids = inputs.get("token_type_ids", None) mm_token_type_ids = inputs.get("mm_token_type_ids", None) num_items_in_batch = inputs.get("num_items_in_batch", None) sampling_per_token_logps = inputs.get("sampling_per_token_logps", None) tool_mask = inputs.get("tool_mask", None) # Missing when evaluate() runs standalone; eval does not accumulate, so # fall back to 1 to avoid underreporting eval_loss (#2464). current_gradient_accumulation_steps = getattr( self, "current_gradient_accumulation_steps", 1 ) num_processes = self.accelerator.num_processes input_ids = torch.cat([prompt_ids, completion_ids], dim = 1) bsz, qlen = input_ids.shape attention_mask = torch.cat([prompt_mask, completion_mask], dim = 1) if mm_token_type_ids is not None or image_grid_thw is not None: mm_token_type_ids = _unsloth_fix_mm_token_type_ids( self.processing_class, input_ids, mm_token_type_ids, completion_ids = completion_ids, ) # attention_mask = None logits_to_keep = completion_ids.size( 1 ) # we only need to compute the logits for the completion tokens _input_ids = input_ids _logits_to_keep = logits_to_keep get_logps_func = ( lambda model, input_ids, attention_mask, logits_to_keep, batch_size = None, compute_entropy = False, compute_efficient = False: self._get_per_token_logps( model, input_ids, attention_mask, logits_to_keep, compute_efficient ) if hasattr(self, "_get_per_token_logps") else self._get_per_token_logps_and_entropies( model, input_ids, attention_mask, logits_to_keep, batch_size, compute_entropy, compute_efficient, )[0] ) # logps per_token_logps = get_logps_func( model, input_ids, attention_mask, logits_to_keep, compute_efficient = True ) # Compute the KL divergence between the model and the reference model # _prepare_inputs doesn't return reference log probs anymore. We need to calculate it ourselves. # https://github.com/huggingface/trl/blob/05bc43e960396581e458195b8388efe6b82cae1f/trl/trainer/grpo_trainer.py#L1328 # if self.beta != 0.0: # with torch.inference_mode(), model.disable_adapter(): # ref_per_token_logps = per_token_logps = get_logps_func(model, input_ids, attention_mask, logits_to_keep) # else: # ref_per_token_logps = None ref_logps = inputs.get("ref_per_token_logps", None) # per_token_kl = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1 # x - x.detach() allows for preserving gradients from x advantages = inputs["advantages"] # per_token_loss = torch.exp(per_token_logps - per_token_logps.detach()) * advantages.unsqueeze(1) # per_token_loss = -(per_token_loss - self.beta * per_token_kl) # loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean() old_logps = inputs.get("old_per_token_logps", None) input_ids = input_ids[:, -logits_to_keep:] # Get logit softcapping and logit scale model_config = _unsloth_get_model_config(model) logit_softcapping = _unsloth_get_final_logit_softcapping(model) # Gemma logit_scale_multiply = getattr(model_config, "logit_scale", 0) # Cohere if logit_scale_multiply is None: logit_scale_multiply = 0 logit_scale_divide = getattr(model_config, "logits_scaling", 0) # Granite if logit_scale_divide is None: logit_scale_divide = 0 max_left_pad = inputs.get("max_left_pad", 0) if per_token_logps is not None: loss_mask = completion_mask if tool_mask is not None: if tool_mask.shape != completion_mask.shape: raise ValueError( "tool_mask/env_mask must have the same shape as completion_mask" ) loss_mask = completion_mask * tool_mask.to( device = completion_mask.device, dtype = completion_mask.dtype, ) ( loss, completion_length, mean_kl, delta, flat_is_ratio, coef_1, completion_mask, ) = grpo_compute_loss_slow( ref_logps, per_token_logps, old_logps, sampling_per_token_logps, input_ids, loss_mask, self.beta, advantages, pixel_values = pixel_values, image_grid_thw = image_grid_thw, loss_type = self.args.loss_type, importance_sampling_level = self.importance_sampling_level, epsilon_low = self.epsilon_low, epsilon_high = self.epsilon_high, max_completion_length = self.args.max_completion_length, delta = self.args.delta, temperature = self.args.temperature, max_left_pad = max_left_pad, logit_softcapping = logit_softcapping, logit_scale_multiply = logit_scale_multiply, logit_scale_divide = logit_scale_divide, num_items_in_batch = num_items_in_batch, current_gradient_accumulation_steps = current_gradient_accumulation_steps, num_processes = num_processes, ) else: def _unsloth_requires_multi_image_zoo(value): if value is None: return False if isinstance(value, torch.Tensor): counts = value.detach().cpu().reshape(-1).tolist() else: counts = list(value) return any(int(n) != 1 for n in counts) if _unsloth_requires_multi_image_zoo(num_images) and not getattr( self, "_unsloth_grpo_zoo_checked", False ): _supports_num_images = ( "num_images" in inspect.signature(grpo_accumulated_loss).parameters ) if not _supports_num_images: try: _zoo_src = inspect.getsource(grpo_accumulated_loss) except (TypeError, OSError): _zoo_src = "" _supports_num_images = "num_images" in _zoo_src if not _supports_num_images: raise RuntimeError( "Multi-image GRPO requires an unsloth_zoo build whose " "grpo_accumulated_loss handles num_images. Please upgrade " "unsloth_zoo (see https://github.com/unslothai/unsloth-zoo/pull/613)." ) self._unsloth_grpo_zoo_checked = True if tool_mask is not None and not getattr( self, "_unsloth_grpo_tool_mask_zoo_checked", False ): _supports_tool_mask = ( "tool_mask" in inspect.signature(grpo_accumulated_loss).parameters ) if not _supports_tool_mask: try: _zoo_src = inspect.getsource(grpo_accumulated_loss) except (TypeError, OSError): _zoo_src = "" _supports_tool_mask = "tool_mask" in _zoo_src if not _supports_tool_mask: raise RuntimeError( "env_mask/tool_mask GRPO requires an unsloth_zoo build whose " "grpo_accumulated_loss handles tool_mask. Please upgrade " "unsloth_zoo." ) self._unsloth_grpo_tool_mask_zoo_checked = True _grpo_accumulated_loss_kwargs = {} if tool_mask is not None: _grpo_accumulated_loss_kwargs["tool_mask"] = tool_mask if hasattr(self.args, "loss_type"): ( loss, completion_length, mean_kl, delta, flat_is_ratio, coef_1, completion_mask, ) = grpo_accumulated_loss( trainer = self, input_ids = _input_ids, pixel_values = pixel_values, image_grid_thw = image_grid_thw, pixel_attention_mask = pixel_attention_mask, image_sizes = image_sizes, num_images = num_images, logits_to_keep = logits_to_keep, completion_mask = completion_mask, advantages = advantages, old_logps = old_logps, ref_logps = ref_logps, n_chunks = self.args.unsloth_num_chunks, loss_type = self.args.loss_type, importance_sampling_level = self.importance_sampling_level, epsilon_low = self.epsilon_low, epsilon_high = self.epsilon_high, max_completion_length = self.args.max_completion_length, delta = self.args.delta, temperature = self.args.temperature, max_left_pad = max_left_pad, logit_softcapping = logit_softcapping, logit_scale_multiply = logit_scale_multiply, logit_scale_divide = logit_scale_divide, attention_mask = attention_mask, num_items_in_batch = num_items_in_batch, current_gradient_accumulation_steps = current_gradient_accumulation_steps, num_processes = num_processes, sampling_per_token_logps = sampling_per_token_logps, token_type_ids = token_type_ids, mm_token_type_ids = mm_token_type_ids, **_grpo_accumulated_loss_kwargs, ) else: # to ensure backwards compatibility with trl 0.15.2 and maybe even 0.17 loss, completion_length, mean_kl, coef_1, completion_mask = grpo_accumulated_loss( trainer = self, input_ids = _input_ids, pixel_values = pixel_values, image_grid_thw = image_grid_thw, pixel_attention_mask = pixel_attention_mask, image_sizes = image_sizes, num_images = num_images, logits_to_keep = logits_to_keep, completion_mask = completion_mask, advantages = advantages, old_logps = old_logps, ref_logps = ref_logps, n_chunks = self.args.unsloth_num_chunks, temperature = self.args.temperature, logit_softcapping = logit_softcapping, logit_scale_multiply = logit_scale_multiply, logit_scale_divide = logit_scale_divide, attention_mask = attention_mask, token_type_ids = token_type_ids, mm_token_type_ids = mm_token_type_ids, **_grpo_accumulated_loss_kwargs, ) if "train" in self._metrics: mode = "eval" if self.control.should_evaluate else "train" self._metrics[mode]["completion_length"].append(completion_length.item()) self._metrics[mode]["kl"].append(mean_kl.item()) else: self._metrics["completion_length"].append(completion_length.item()) self._metrics["kl"].append(mean_kl.item()) if ( self.use_vllm and delta is not None and getattr(self, "vllm_importance_sampling_correction", False) ): mean_delta = ( torch.mean(delta) if delta.numel() > 0 else torch.tensor(0.0, device = self.model.device) ) max_delta = ( torch.max(delta) if delta.numel() > 0 else torch.tensor(0.0, device = self.model.device) ) self._metrics[mode]["sampling/sampling_logp_difference/mean"].append( self.accelerator.gather(mean_delta).mean().item() ) self._metrics[mode]["sampling/sampling_logp_difference/max"].append( self.accelerator.gather(max_delta).max().item() ) min_importance_sampling_ratio = ( torch.min(flat_is_ratio) if flat_is_ratio.numel() > 0 else torch.tensor(0.0, device = self.model.device) ) mean_importance_sampling_ratio = ( torch.mean(flat_is_ratio) if flat_is_ratio.numel() > 0 else torch.tensor(0.0, device = self.model.device) ) max_importance_sampling_ratio = ( torch.max(flat_is_ratio) if flat_is_ratio.numel() > 0 else torch.tensor(0.0, device = self.model.device) ) self._metrics[mode]["sampling/importance_sampling_ratio/min"].append( self.accelerator.gather(min_importance_sampling_ratio) .nan_to_num(nan = float("inf")) .min() .item() ) self._metrics[mode]["sampling/importance_sampling_ratio/mean"].append( self.accelerator.gather(mean_importance_sampling_ratio).nanmean().item() ) self._metrics[mode]["sampling/importance_sampling_ratio/max"].append( self.accelerator.gather(max_importance_sampling_ratio) .nan_to_num(nan = float("-inf")) .max() .item() ) completion_token_count = completion_mask.sum().clamp(min = 1.0) def masked_batch_mean(x): if x.shape[1] == 1: # when importance_sampling_level == "sequence" return x.mean() else: return (x * completion_mask).sum() / completion_token_count if advantages.dim() == 1: advantages = advantages.unsqueeze(1) if self.loss_type in ["grpo", "bnpo", "dr_grpo", "dapo"]: # Compute the clipped probability ratios 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 low_clip = masked_batch_mean(is_low_clipped.float()) high_clip = masked_batch_mean(is_high_clipped.float()) clip_ratio = masked_batch_mean(is_region_clipped.float()) gathered_low_clip = self.accelerator.gather(low_clip) 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(high_clip) 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(clip_ratio) self._metrics[mode]["clip_ratio/region_mean"].append( gathered_clip_ratio.nanmean().item() ) elif self.loss_type == "cispo": is_cispo_clipped = (coef_1 > self.epsilon_high) & (advantages > 0) cispo_clip_ratio = masked_batch_mean(is_cispo_clipped.float()) gathered_cispo_clip_ratio = self.accelerator.gather(cispo_clip_ratio) self._metrics[mode]["cispo_clip_ratio"].append( gathered_cispo_clip_ratio.nanmean().item() ) return loss function = inspect.getsource(compute_loss) return function RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer_compute_loss) # Fix KTO shape mismatch when Unsloth model forward truncates input_ids # but labels aren't truncated. TRL 0.27.2+ _process_tokens only truncates # completions, not prompts -- so prompts exceeding max_seq_length cause the # model to produce shorter logits than the labels expect. def kto_trainer_get_batch_logps(function_name, function): if function_name != "get_batch_logps": return function # The raise is inside an if block inside the method, so we need # to preserve the exact indentation of the raise statement. old = 'raise ValueError("Logits (batch and sequence length dim) and labels must have the same shape.")' new = ( "# Unsloth: auto-truncate to shorter sequence length (model may have truncated input_ids)\n" " _min_len = min(logits.shape[1], labels.shape[1])\n" " logits = logits[:, :_min_len, :]\n" " labels = labels[:, :_min_len]" ) function = function.replace(old, new) return function RL_FUNCTIONS["kto_trainer"].append(kto_trainer_get_batch_logps) # TRL 1.x dropped KTOTrainer.get_batch_logps and moved the log-prob math into # _compute_logps / compute_ref_log_probs / _compute_kl_logps, which call # selective_log_softmax on completion-only tokens. Same truncation hazard as # above, so clamp logits/ids/mask to the shorter seq length (no-op when equal). _KTO_COMPLETION_RE = re.compile( r"(?P[ \t]*)shift_logits = completion_logits\[:, :-1, :\]\.contiguous\(\)\n" r"(?P=ws)per_token_logps = selective_log_softmax\(\s*shift_logits,\s*" r"(?P\w+)\[[\"']completion_input_ids[\"']\]\[:, 1:\]\.contiguous\(\)\s*\)\n" r"(?P=ws)per_token_logps\[(?P=var)\[[\"']completion_mask[\"']\]\[:, 1:\] == 0\] = 0\.0" ) _KTO_KL_RE = re.compile( r"(?P[ \t]*)shift_KL_logits = KL_logits\[:, :-1, :\]\.contiguous\(\)\n" r"(?P=ws)KL_per_token_logps = selective_log_softmax\(\s*shift_KL_logits,\s*" r"(?P\w+)\[[\"']KL_completion_input_ids[\"']\]\[:, 1:\]\.contiguous\(\)\s*\)\n" r"(?P=ws)KL_per_token_logps\[(?P=var)\[[\"']KL_completion_mask[\"']\]\[:, 1:\] == 0\] = 0\.0" ) def _kto_completion_repl(m): ws, var = m.group("ws"), m.group("var") return ( f"{ws}shift_logits = completion_logits[:, :-1, :].contiguous()\n" f"{ws}# Unsloth: clamp logits/ids/mask to shorter seq len (model may truncate input_ids)\n" f'{ws}_uns_ids = {var}["completion_input_ids"][:, 1:].contiguous()\n' f"{ws}_uns_n = min(shift_logits.shape[1], _uns_ids.shape[1])\n" f"{ws}per_token_logps = selective_log_softmax(shift_logits[:, :_uns_n], _uns_ids[:, :_uns_n])\n" f'{ws}per_token_logps[{var}["completion_mask"][:, 1:][:, :_uns_n] == 0] = 0.0' ) def _kto_kl_repl(m): ws, var = m.group("ws"), m.group("var") return ( f"{ws}shift_KL_logits = KL_logits[:, :-1, :].contiguous()\n" f"{ws}# Unsloth: clamp logits/ids/mask to shorter seq len (model may truncate input_ids)\n" f'{ws}_uns_kl_ids = {var}["KL_completion_input_ids"][:, 1:].contiguous()\n' f"{ws}_uns_kl_n = min(shift_KL_logits.shape[1], _uns_kl_ids.shape[1])\n" f"{ws}KL_per_token_logps = selective_log_softmax(shift_KL_logits[:, :_uns_kl_n], _uns_kl_ids[:, :_uns_kl_n])\n" f'{ws}KL_per_token_logps[{var}["KL_completion_mask"][:, 1:][:, :_uns_kl_n] == 0] = 0.0' ) def kto_trainer_align_completion_logps(function_name, function): if function_name not in ( "_compute_logps", "compute_ref_log_probs", "_compute_kl_logps", ): return function function = _KTO_COMPLETION_RE.sub(_kto_completion_repl, function) function = _KTO_KL_RE.sub(_kto_kl_repl, function) return function RL_FUNCTIONS["kto_trainer"].append(kto_trainer_align_completion_logps) # https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_trainer.py#L356 # TRL warns if batch size is not a multiple of num_generations -> fix this. def grpo_trainer_fix_batch_size(RLTrainer_source, RLConfig_source): if "divisible by the number of generations" not in RLTrainer_source: # in later trl versions this doesn't exist anymore return "" if "num_generations" not in RLConfig_source: return "" check_batch_size = ( "div = per_device_train_batch_size // num_generations\n" "if div * num_generations != per_device_train_batch_size:\n" " print('Unsloth: We now expect `per_device_train_batch_size` to be a multiple of `num_generations`.\\n" "We will change the batch size of ' + str(per_device_train_batch_size) + ' to the `num_generations` of ' + str(num_generations))\n" " per_device_train_batch_size = num_generations\n" ) return check_batch_size RL_CONFIG_CHANGES["grpo_trainer"].append(grpo_trainer_fix_batch_size) # Add other reward function names def grpo_trainer_metrics(RLTrainer_source, RLConfig_source): if "reward_funcs" not in RLTrainer_source: return "" # For new TRL we have /mean and /std use_mean = "rewards/{reward_func_name}/mean" in RLTrainer_source use_std = "rewards/{reward_func_name}/std" in RLTrainer_source if not use_mean: use_normal = "rewards/{reward_func_name}" in RLTrainer_source else: use_normal = False log_metrics = ( "if not isinstance(reward_funcs, list): _reward_funcs = [reward_funcs]\n" "else: _reward_funcs = reward_funcs\n" "for reward_func in _reward_funcs:\n" " try:\n" " reward_func_name = reward_func.__name__\n" f" if {use_mean}:\n" " other_metrics.append(f'rewards/{reward_func_name}/mean')\n" f" if {use_std}:\n" " other_metrics.append(f'rewards/{reward_func_name}/std')\n" f" if {use_normal}:\n" " other_metrics.append(f'rewards/{reward_func_name}')\n" " except: pass\n" ) return log_metrics RL_METRICS_CHANGES["grpo_trainer"].append(grpo_trainer_metrics) def openenv_vllm_reload_weights(): # This function patches the trl openenv generate_rollout_completions function to: # 1. Guard the reload_weights call (skip when sharing weights with vLLM) # 2. Fix wake_up call to be compatible with unsloth (remove tags to wake everything) # # The issue: TRL's wake_up(tags=["kv_cache"]) only wakes kv_cache, leaving is_sleeping=True # at the executor level. This causes unsloth's patched generate to try waking up again, # resulting in double create_and_map on already-mapped handles. # # The fix: Use wake_up() with no tags, which wakes everything. Unsloth's patched # CuMemAllocator.wake_up skips weights anyway, so this is safe. if importlib.util.find_spec("trl") is None: return if Version(importlib_version("trl")) < Version("0.26.0"): return try: import trl.experimental.openenv.utils as openenv_utils import trl.experimental.openenv as openenv except (ImportError, NameError, Exception) as e: logger.info(f"Unsloth: Failed to import trl openenv: {e}") logger.info( "Unsloth: trl.experimental.openenv not available — skipping RL openenv patches." ) return # trl 0.28 changed the function name yet again! Thanks trl :) patch_target_name = "_generate_rollout_completions_colocate" if hasattr(openenv_utils, patch_target_name): patch_target = getattr(openenv_utils, patch_target_name) else: # Older TRL versions may keep sleep/wake logic in the public dispatcher. patch_target_name = "generate_rollout_completions" patch_target = getattr(openenv_utils, patch_target_name) # TRL 0.29.1+ ships some openenv helpers as compiled bytecode without # accessible source on disk; inspect.getsource raises OSError("could # not get source code") in that case. Skip the source-rewrite patch # rather than crash. The unmodified TRL openenv path will run, which # means the duplicate `collective_rpc("reload_weights")` is NOT # stripped (line 1800 below) and `wake_up(tags=["kv_cache"])` is NOT # retagged to `wake_up()` (line 1804). Users who do not use openenv # GRPO are unaffected; openenv GRPO users on this TRL build may see # redundant reload_weights calls or partial wake_up behavior. try: src = inspect.getsource(patch_target) except OSError as e: logger.warning( f"Unsloth: Could not retrieve source for trl openenv " f"{patch_target_name} ({e}); skipping rewrite. The unmodified " f"TRL openenv path will run, so the duplicate reload_weights " f"strip and the wake_up tag rewrite are NOT applied. Open an " f"issue if you see redundant reload_weights or partial wake_up " f"on openenv GRPO with this TRL build." ) return src = textwrap.dedent(src) original_src = src reload_weights_pattern = re.compile( r"^(?P[ \t]*)(?P\S+)\.collective_rpc\(\s*(['\"])reload_weights\3\s*\)\s*$", re.MULTILINE, ) def replace_reload_weights(match): indent = match.group("indent") obj = match.group("obj") return ( f"{indent}if not getattr({obj}, 'shared_weights', False):\n" f'{indent} {obj}.collective_rpc("reload_weights")\n' ) src = reload_weights_pattern.sub(replace_reload_weights, src) # Change wake_up(tags=["kv_cache"]) to wake_up() - wake everything to set is_sleeping=False # This prevents double wake_up issues. Unsloth's allocator skips weights anyway. src = re.sub(r"\.wake_up\(tags=\[.*?\]\)", ".wake_up()", src) if original_src == src: logger.warning("Unsloth: Warning - regex did not match, patch may have failed") return # Execute and explicitly assign to module local_ns = {} exec(compile(src, "", "exec"), openenv_utils.__dict__, local_ns) patched_func = local_ns[patch_target_name] # Patch the target function in utils; if dispatcher was patched also update parent module alias. setattr(openenv_utils, patch_target_name, patched_func) if patch_target_name == "generate_rollout_completions": openenv.generate_rollout_completions = patched_func logger.info(f"Unsloth: Patched trl openenv {patch_target_name}") RL_ADDITIONAL_FUNCTIONS["openenv"].append(openenv_vllm_reload_weights) def vllm_generation_init_patch(): # trl moved vllm stuff to trl/generation/vllm_generation.py # We need to patch it to not instantiate another vLLM instance if we already have one with fast_inference # Edit the TRL source directly and install the patched function in the TRL module. # https://github.com/huggingface/trl/commit/0eb66d8f2fc63b3d00d8dbc18f99c3f48750bd16 # This exists in trl versions 0.28.0 and above if importlib.util.find_spec("trl") is None: return if Version(importlib_version("trl")) < Version("0.28.0"): return try: import trl.generation.vllm_generation as vllm_generation except (ImportError, NameError, Exception) as e: logger.info(f"Unsloth: Failed to import trl.generation.vllm_generation: {e}") return def patch_vllm_generation_method(method_name, transform, marker, filename_suffix): method = getattr(vllm_generation.VLLMGeneration, method_name, None) if method is None: logger.info(f"Unsloth: Could not find VLLMGeneration.{method_name}") return False try: src = inspect.getsource(method) except Exception as e: logger.info(f"Unsloth: Could not get source of VLLMGeneration.{method_name}: {e}") return False src = textwrap.dedent(src) if marker in src: return True src = transform(src) filename = f"" source_lines = [line + "\n" for line in src.splitlines()] linecache.cache[filename] = ( len(src), None, source_lines, filename, ) local_ns = {} exec(compile(src, filename, "exec"), vllm_generation.__dict__, local_ns) setattr(vllm_generation.VLLMGeneration, method_name, local_ns[method_name]) return True # Patch init to remove vLLM.LLM instantiation def patch_init_vllm(src): pattern = re.compile( r"(?P^(?P[ \t]*)self\.llm\s*=\s*LLM\s*\(\n(?:.*\n)*?^(?P=indent)\))", re.MULTILINE, ) def replace_llm_block(match): indent = match.group("indent") llm_block = textwrap.dedent(match.group("llm_block")) return ( f"{indent}if hasattr(model, 'vllm_engine'):\n" f"{indent} # Unsloth already inits vLLM in fast inference mode. Do not redo :)\n" f"{indent} self.llm = model.vllm_engine\n" f"{indent} self.unsloth_fast_inference_lora = getattr(self.llm, 'shared_weights', False)\n" f"{indent} if getattr(self.llm, 'shared_weights', False) and hasattr(model, 'load_lora'):\n" f"{indent} self._unsloth_load_lora = model.load_lora\n" f"{indent}else:\n" + textwrap.indent(llm_block, indent + " ") ) patched_src, num_replacements = pattern.subn(replace_llm_block, src, count = 1) if num_replacements == 0: raise RuntimeError( "Unsloth: Warning - regex did not match, VLLMGeneration._init_vllm patch may have failed" ) return patched_src # has some sync_weights or reload rpc calls. # we patched the grpo_trainer to strip them for prev versions # Ref: grpo_trainer__generate_single_turn above around L270-280 def patch_sync_weights(src): pattern = re.compile( r"^(?Pdef sync_weights\(self\):\n)(?P(?:.*\n)*)", re.MULTILINE, ) def replace_sync_weights(match): body = match.group("body") # Chain getattr so server mode (where self.llm is not set) does # not raise AttributeError before the default kicks in. guard = ( " if getattr(getattr(self, 'llm', None), 'shared_weights', False) or " "getattr(self, 'unsloth_fast_inference_lora', False):\n" " # Unsloth fast inference LoRA shares weights with vLLM already.\n" " return\n\n" ) return match.group("def_line") + guard + body patched_src, num_replacements = pattern.subn(replace_sync_weights, src, count = 1) if num_replacements == 0: raise RuntimeError( "Unsloth: Warning - regex did not match, VLLMGeneration.sync_weights patch may have failed" ) return patched_src def patch_generate(src): pattern = re.compile( r"^(?P[ \t]*)self\.llm\.collective_rpc\(\s*(['\"])reload_weights\2\s*\)\s*$", re.MULTILINE, ) def replace_reload_weights(match): indent = match.group("indent") # Chain getattr so server mode (no self.llm) is safe here too. return ( f"{indent}if not (getattr(getattr(self, 'llm', None), 'shared_weights', False) or " f"getattr(self, 'unsloth_fast_inference_lora', False)):\n" f'{indent} self.llm.collective_rpc("reload_weights")' ) patched_src, num_replacements = pattern.subn(replace_reload_weights, src, count = 1) if num_replacements == 0: raise RuntimeError( "Unsloth: Warning - regex did not match, VLLMGeneration.generate patch may have failed" ) # Inject lora_request when sharing weights (vLLM needs the adapter) lora_generate_pattern = re.compile( r"(self\.llm\.generate\([^\)]+)\)", ) def inject_lora_request(match): return ( f"{match.group(1)}, lora_request=" f"self._unsloth_load_lora('vllm_gen_lora', load_tensors=True) " f"if hasattr(self, '_unsloth_load_lora') else None)" ) patched_src = lora_generate_pattern.sub(inject_lora_request, patched_src) return patched_src try: init_patched = patch_vllm_generation_method( "_init_vllm", patch_init_vllm, "self.unsloth_fast_inference_lora = getattr(self.llm, 'shared_weights', False)", "init_vllm", ) sync_patched = patch_vllm_generation_method( "sync_weights", patch_sync_weights, "if getattr(getattr(self, 'llm', None), 'shared_weights', False) or getattr(self, 'unsloth_fast_inference_lora', False):", "sync_weights", ) generate_patched = patch_vllm_generation_method( "generate", patch_generate, "if not (getattr(getattr(self, 'llm', None), 'shared_weights', False) or getattr(self, 'unsloth_fast_inference_lora', False)):", "generate", ) except RuntimeError as e: logger.warning(str(e)) return if init_patched: logger.info("Unsloth: Patched trl VLLMGeneration._init_vllm") if sync_patched: logger.info("Unsloth: Patched trl VLLMGeneration.sync_weights") if generate_patched: logger.info("Unsloth: Patched trl VLLMGeneration.generate") RL_ADDITIONAL_FUNCTIONS["vllm_generation"].append(vllm_generation_init_patch)