diff --git a/unsloth_compiled_cache/.locks/.lock.AqlmLoraLinear_peft_forward.py b/unsloth_compiled_cache/.locks/.lock.AqlmLoraLinear_peft_forward.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.AwqLoraLinear_peft_forward.py b/unsloth_compiled_cache/.locks/.lock.AwqLoraLinear_peft_forward.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.BatchNorm1d.py b/unsloth_compiled_cache/.locks/.lock.BatchNorm1d.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.BatchNorm2d.py b/unsloth_compiled_cache/.locks/.lock.BatchNorm2d.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.BatchNorm3d.py b/unsloth_compiled_cache/.locks/.lock.BatchNorm3d.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.BlockDiagonalLinear_peft_forward.py b/unsloth_compiled_cache/.locks/.lock.BlockDiagonalLinear_peft_forward.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.Conv1d.py b/unsloth_compiled_cache/.locks/.lock.Conv1d.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.Conv2d.py b/unsloth_compiled_cache/.locks/.lock.Conv2d.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.Conv3d.py b/unsloth_compiled_cache/.locks/.lock.Conv3d.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.ConvTranspose1d.py b/unsloth_compiled_cache/.locks/.lock.ConvTranspose1d.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.ConvTranspose2d.py b/unsloth_compiled_cache/.locks/.lock.ConvTranspose2d.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.ConvTranspose3d.py b/unsloth_compiled_cache/.locks/.lock.ConvTranspose3d.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.GPTQLoraLinear_peft_forward.py b/unsloth_compiled_cache/.locks/.lock.GPTQLoraLinear_peft_forward.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.GroupNorm.py b/unsloth_compiled_cache/.locks/.lock.GroupNorm.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.LayerNorm.py b/unsloth_compiled_cache/.locks/.lock.LayerNorm.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.Linear4bit_peft_forward.py b/unsloth_compiled_cache/.locks/.lock.Linear4bit_peft_forward.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.Linear8bitLt_peft_forward.py b/unsloth_compiled_cache/.locks/.lock.Linear8bitLt_peft_forward.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.Linear_peft_forward.py b/unsloth_compiled_cache/.locks/.lock.Linear_peft_forward.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.LoraParallelLinear_peft_forward.py b/unsloth_compiled_cache/.locks/.lock.LoraParallelLinear_peft_forward.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.RMSNorm.py b/unsloth_compiled_cache/.locks/.lock.RMSNorm.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.TeLinear_peft_forward.py b/unsloth_compiled_cache/.locks/.lock.TeLinear_peft_forward.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.UnslothCPOTrainer.py b/unsloth_compiled_cache/.locks/.lock.UnslothCPOTrainer.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.UnslothDPOTrainer.py b/unsloth_compiled_cache/.locks/.lock.UnslothDPOTrainer.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.UnslothDistillationTrainer.py b/unsloth_compiled_cache/.locks/.lock.UnslothDistillationTrainer.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.UnslothGKDTrainer.py b/unsloth_compiled_cache/.locks/.lock.UnslothGKDTrainer.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.UnslothGOLDTrainer.py b/unsloth_compiled_cache/.locks/.lock.UnslothGOLDTrainer.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.UnslothGRPOTrainer.py b/unsloth_compiled_cache/.locks/.lock.UnslothGRPOTrainer.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.UnslothKTOTrainer.py b/unsloth_compiled_cache/.locks/.lock.UnslothKTOTrainer.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.UnslothORPOTrainer.py b/unsloth_compiled_cache/.locks/.lock.UnslothORPOTrainer.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.UnslothOnlineDPOTrainer.py b/unsloth_compiled_cache/.locks/.lock.UnslothOnlineDPOTrainer.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.UnslothRLOOTrainer.py b/unsloth_compiled_cache/.locks/.lock.UnslothRLOOTrainer.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.UnslothRewardTrainer.py b/unsloth_compiled_cache/.locks/.lock.UnslothRewardTrainer.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.UnslothSDFTTrainer.py b/unsloth_compiled_cache/.locks/.lock.UnslothSDFTTrainer.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.UnslothSFTTrainer.py b/unsloth_compiled_cache/.locks/.lock.UnslothSFTTrainer.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.UnslothServerDistillationTrainer.py b/unsloth_compiled_cache/.locks/.lock.UnslothServerDistillationTrainer.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.unsloth_compiled_module_gemma4.py b/unsloth_compiled_cache/.locks/.lock.unsloth_compiled_module_gemma4.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/.locks/.lock.unsloth_compiled_module_siglip.py b/unsloth_compiled_cache/.locks/.lock.unsloth_compiled_module_siglip.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/unsloth_compiled_cache/AqlmLoraLinear_peft_forward.py b/unsloth_compiled_cache/AqlmLoraLinear_peft_forward.py deleted file mode 100644 index 978edfda31fa83dfe602eb0b76d0263ad0db68c5..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/AqlmLoraLinear_peft_forward.py +++ /dev/null @@ -1,99 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - - -torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 32, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': False, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from unsloth_zoo.temporary_patches.common import torch_compile -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from peft.tuners.lora.aqlm import (torch) - - -torch_addmm = torch.addmm -torch_add = torch.add -torch_is_compiling = torch.compiler.is_compiling -# torch 2.8 / 2.9 inductor miscompiles addmm(alpha = scaling, beta = 1) traced inside a -# compiled block (e.g. the VLM decoder MLP), forward and backward; fixed in 2.10 (unsloth#3271). -torch_lora_no_compiled_addmm = tuple(int(v) for v in torch.__version__.split(".")[:2]) < (2, 10) -# @torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) -def lora_forward(result, lora_A, lora_B, dropout, x, scaling): - # Use result.dtype (bfloat16 from base layer) since x may have been cast to float32 - # by _cast_input_dtype when autocast is disabled - target_dtype = result.dtype - xA = dropout(x).to(target_dtype) @ lora_A.weight.to(target_dtype).t() - if torch_lora_no_compiled_addmm and torch_is_compiling(): - output = result + scaling * (xA @ lora_B.weight.to(target_dtype).t()) - bias = lora_B.bias - if bias is not None: - output = output + scaling * bias.to(target_dtype) - return output - # output = result + scaling * xA @ lora_B.weight.t() - shape = result.shape - output = torch_addmm( - result.view(-1, shape[-1]), - xA.view(-1, xA.shape[-1]), - lora_B.weight.to(target_dtype).t(), - alpha = scaling, - beta = 1, - ).view(shape) - - bias = lora_B.bias - if bias is not None: - output = torch_add( - output, - bias.to(target_dtype), - alpha = scaling, - ) - return output -pass - -def unsloth_forward(self, x: torch.Tensor): - # note: logic differs from default Linear because merging is not supported - result = self.base_layer(x) - - if self.disable_adapters: - return result - - for active_adapter in self.active_adapters: - if active_adapter not in self.lora_A.keys(): - continue - lora_A = self.lora_A[active_adapter] - lora_B = self.lora_B[active_adapter] - dropout = self.lora_dropout[active_adapter] - scaling = self.scaling[active_adapter] - - requires_conversion = not torch.is_autocast_enabled() - if requires_conversion: - expected_dtype = result.dtype - x = self._cast_input_dtype(x, lora_A.weight.dtype) - - output = lora_B(lora_A(dropout(x))) - if requires_conversion: - output = output.to(expected_dtype) - output = output * scaling - result += output - return result diff --git a/unsloth_compiled_cache/AwqLoraLinear_peft_forward.py b/unsloth_compiled_cache/AwqLoraLinear_peft_forward.py deleted file mode 100644 index c0c13a41d26f4cd2ed25d9f80ccee41bd1760792..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/AwqLoraLinear_peft_forward.py +++ /dev/null @@ -1,98 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - - -torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 32, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': False, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from unsloth_zoo.temporary_patches.common import torch_compile -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from peft.tuners.lora.awq import (torch) - - -torch_addmm = torch.addmm -torch_add = torch.add -torch_is_compiling = torch.compiler.is_compiling -# torch 2.8 / 2.9 inductor miscompiles addmm(alpha = scaling, beta = 1) traced inside a -# compiled block (e.g. the VLM decoder MLP), forward and backward; fixed in 2.10 (unsloth#3271). -torch_lora_no_compiled_addmm = tuple(int(v) for v in torch.__version__.split(".")[:2]) < (2, 10) -# @torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) -def lora_forward(result, lora_A, lora_B, dropout, x, scaling): - # Use result.dtype (bfloat16 from base layer) since x may have been cast to float32 - # by _cast_input_dtype when autocast is disabled - target_dtype = result.dtype - xA = dropout(x).to(target_dtype) @ lora_A.weight.to(target_dtype).t() - if torch_lora_no_compiled_addmm and torch_is_compiling(): - output = result + scaling * (xA @ lora_B.weight.to(target_dtype).t()) - bias = lora_B.bias - if bias is not None: - output = output + scaling * bias.to(target_dtype) - return output - # output = result + scaling * xA @ lora_B.weight.t() - shape = result.shape - output = torch_addmm( - result.view(-1, shape[-1]), - xA.view(-1, xA.shape[-1]), - lora_B.weight.to(target_dtype).t(), - alpha = scaling, - beta = 1, - ).view(shape) - - bias = lora_B.bias - if bias is not None: - output = torch_add( - output, - bias.to(target_dtype), - alpha = scaling, - ) - return output -pass - -def unsloth_forward(self, x: torch.Tensor): - result = self.quant_linear_module(x) - - if self.disable_adapters: - return result - - for active_adapter in self.active_adapters: - if active_adapter not in self.lora_A.keys(): - continue - lora_A = self.lora_A[active_adapter] - lora_B = self.lora_B[active_adapter] - dropout = self.lora_dropout[active_adapter] - scaling = self.scaling[active_adapter] - - requires_conversion = not torch.is_autocast_enabled() - if requires_conversion: - expected_dtype = result.dtype - x = self._cast_input_dtype(x, lora_A.weight.dtype) - - output = lora_B(lora_A(dropout(x))) - if requires_conversion: - output = output.to(expected_dtype) - output = output * scaling - result = result + output - return result diff --git a/unsloth_compiled_cache/BatchNorm1d.py b/unsloth_compiled_cache/BatchNorm1d.py deleted file mode 100644 index 44e040f9531607e81dae7bd129d13ac88128a2cd..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/BatchNorm1d.py +++ /dev/null @@ -1,151 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - - -import os -import sys -import torch -import importlib.util -import math -if importlib.util.find_spec("unsloth_studio") is None: - UNSLOTH_STUDIO_ENABLED = False -else: - UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" -pass -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -import math - -UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" -UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" -UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) -UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") -if UNSLOTH_COMPILE_LOCATION not in sys.path: - sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) - -import logging -logger_compiler = logging.getLogger(__name__) -if UNSLOTH_ENABLE_LOGGING: - logger_compiler.setLevel(logging.DEBUG) - -global INFERENCE_RUNS -INFERENCE_RUNS = 0 - -try: - import torch._dynamo.eval_frame as torch_dynamo_eval_frame - torch_dynamo_eval_frame._stance.stance - torch_compiler_set_stance = torch.compiler.set_stance -except: - torch_dynamo_eval_frame = None - torch_compiler_set_stance = None -pass - -# The stance is process wide, so whether Unsloth owns it lives in unsloth_zoo, shared by -# every generated module: inference in one model's module, training in another's. -from unsloth_zoo.temporary_patches.utils import UNSLOTH_EAGER_STANCE_OWNED, unsloth_claim_eager_stance, unsloth_owns_stance, unsloth_release_stance - -def unsloth_training_stance(): - # eager_on_recompile is global: left on by generate() or eval forwards, every recompile - # in a later training step (new sequence length, model.train() without for_training) - # would silently run the eager body. Training forwards restore the default stance, - # unless the user chose eager_on_recompile themselves. Every set_stance call installs a - # fresh stance object, so ours is still active only if the current one `is` it; a user - # set_stance since then (even to eager_on_recompile again) is theirs and stays. - global INFERENCE_RUNS - INFERENCE_RUNS = 0 - if len(UNSLOTH_EAGER_STANCE_OWNED) != 0 and torch_dynamo_eval_frame is not None: - current = torch_dynamo_eval_frame._stance - # Ownership is only dropped once ours is reset: a temporary stance scope's exit restores - # one of our exact objects, which a later training forward must still recognise. - if current.stance == "eager_on_recompile" and unsloth_owns_stance(current): - # Only the reset object: an outer scope may still restore another one of ours. - # Put back what the caller had before Unsloth switched (force_eager, a backend, ...). - previous = unsloth_release_stance(current) - if previous is not None and hasattr(torch_dynamo_eval_frame, "_set_stance"): - torch_dynamo_eval_frame._set_stance(previous) - else: - torch_compiler_set_stance(stance = "default", skip_guard_eval_unsafe = False) -pass - -from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT -from unsloth_zoo.temporary_patches.utils import UNSLOTH_DECODE_COMPILE - -torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 32, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': False, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from transformers.models.siglip.modeling_siglip import (np, nn) - -def forward(self, input: Tensor) -> Tensor: - original_dtype = input.dtype - if self.weight is not None: input = input.to(self.weight.dtype) - self._check_input_dim(input) - - # exponential_average_factor is set to self.momentum - # (when it is available) only so that it gets updated - # in ONNX graph when this node is exported to ONNX. - if self.momentum is None: - exponential_average_factor = 0.0 - else: - exponential_average_factor = self.momentum - - if self.training and self.track_running_stats: - # TODO: if statement only here to tell the jit to skip emitting this when it is None - if self.num_batches_tracked is not None: # type: ignore[has-type] - self.num_batches_tracked.add_(1) # type: ignore[has-type] - if self.momentum is None: # use cumulative moving average - exponential_average_factor = 1.0 / float(self.num_batches_tracked) - else: # use exponential moving average - exponential_average_factor = self.momentum - - r""" - Decide whether the mini-batch stats should be used for normalization rather than the buffers. - Mini-batch stats are used in training mode, and in eval mode when buffers are None. - """ - if self.training: - bn_training = True - else: - bn_training = (self.running_mean is None) and (self.running_var is None) - - r""" - Buffers are only updated if they are to be tracked and we are in training mode. Thus they only need to be - passed when the update should occur (i.e. in training mode when they are tracked), or when buffer stats are - used for normalization (i.e. in eval mode when buffers are not None). - """ - return F.batch_norm( - input, - # If buffers are not to be tracked, ensure that they won't be updated - ( - self.running_mean - if not self.training or self.track_running_stats - else None - ), - self.running_var if not self.training or self.track_running_stats else None, - self.weight, - self.bias, - bn_training, - exponential_average_factor, - self.eps, - ).to(original_dtype) diff --git a/unsloth_compiled_cache/BatchNorm2d.py b/unsloth_compiled_cache/BatchNorm2d.py deleted file mode 100644 index 44e040f9531607e81dae7bd129d13ac88128a2cd..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/BatchNorm2d.py +++ /dev/null @@ -1,151 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - - -import os -import sys -import torch -import importlib.util -import math -if importlib.util.find_spec("unsloth_studio") is None: - UNSLOTH_STUDIO_ENABLED = False -else: - UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" -pass -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -import math - -UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" -UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" -UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) -UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") -if UNSLOTH_COMPILE_LOCATION not in sys.path: - sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) - -import logging -logger_compiler = logging.getLogger(__name__) -if UNSLOTH_ENABLE_LOGGING: - logger_compiler.setLevel(logging.DEBUG) - -global INFERENCE_RUNS -INFERENCE_RUNS = 0 - -try: - import torch._dynamo.eval_frame as torch_dynamo_eval_frame - torch_dynamo_eval_frame._stance.stance - torch_compiler_set_stance = torch.compiler.set_stance -except: - torch_dynamo_eval_frame = None - torch_compiler_set_stance = None -pass - -# The stance is process wide, so whether Unsloth owns it lives in unsloth_zoo, shared by -# every generated module: inference in one model's module, training in another's. -from unsloth_zoo.temporary_patches.utils import UNSLOTH_EAGER_STANCE_OWNED, unsloth_claim_eager_stance, unsloth_owns_stance, unsloth_release_stance - -def unsloth_training_stance(): - # eager_on_recompile is global: left on by generate() or eval forwards, every recompile - # in a later training step (new sequence length, model.train() without for_training) - # would silently run the eager body. Training forwards restore the default stance, - # unless the user chose eager_on_recompile themselves. Every set_stance call installs a - # fresh stance object, so ours is still active only if the current one `is` it; a user - # set_stance since then (even to eager_on_recompile again) is theirs and stays. - global INFERENCE_RUNS - INFERENCE_RUNS = 0 - if len(UNSLOTH_EAGER_STANCE_OWNED) != 0 and torch_dynamo_eval_frame is not None: - current = torch_dynamo_eval_frame._stance - # Ownership is only dropped once ours is reset: a temporary stance scope's exit restores - # one of our exact objects, which a later training forward must still recognise. - if current.stance == "eager_on_recompile" and unsloth_owns_stance(current): - # Only the reset object: an outer scope may still restore another one of ours. - # Put back what the caller had before Unsloth switched (force_eager, a backend, ...). - previous = unsloth_release_stance(current) - if previous is not None and hasattr(torch_dynamo_eval_frame, "_set_stance"): - torch_dynamo_eval_frame._set_stance(previous) - else: - torch_compiler_set_stance(stance = "default", skip_guard_eval_unsafe = False) -pass - -from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT -from unsloth_zoo.temporary_patches.utils import UNSLOTH_DECODE_COMPILE - -torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 32, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': False, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from transformers.models.siglip.modeling_siglip import (np, nn) - -def forward(self, input: Tensor) -> Tensor: - original_dtype = input.dtype - if self.weight is not None: input = input.to(self.weight.dtype) - self._check_input_dim(input) - - # exponential_average_factor is set to self.momentum - # (when it is available) only so that it gets updated - # in ONNX graph when this node is exported to ONNX. - if self.momentum is None: - exponential_average_factor = 0.0 - else: - exponential_average_factor = self.momentum - - if self.training and self.track_running_stats: - # TODO: if statement only here to tell the jit to skip emitting this when it is None - if self.num_batches_tracked is not None: # type: ignore[has-type] - self.num_batches_tracked.add_(1) # type: ignore[has-type] - if self.momentum is None: # use cumulative moving average - exponential_average_factor = 1.0 / float(self.num_batches_tracked) - else: # use exponential moving average - exponential_average_factor = self.momentum - - r""" - Decide whether the mini-batch stats should be used for normalization rather than the buffers. - Mini-batch stats are used in training mode, and in eval mode when buffers are None. - """ - if self.training: - bn_training = True - else: - bn_training = (self.running_mean is None) and (self.running_var is None) - - r""" - Buffers are only updated if they are to be tracked and we are in training mode. Thus they only need to be - passed when the update should occur (i.e. in training mode when they are tracked), or when buffer stats are - used for normalization (i.e. in eval mode when buffers are not None). - """ - return F.batch_norm( - input, - # If buffers are not to be tracked, ensure that they won't be updated - ( - self.running_mean - if not self.training or self.track_running_stats - else None - ), - self.running_var if not self.training or self.track_running_stats else None, - self.weight, - self.bias, - bn_training, - exponential_average_factor, - self.eps, - ).to(original_dtype) diff --git a/unsloth_compiled_cache/BatchNorm3d.py b/unsloth_compiled_cache/BatchNorm3d.py deleted file mode 100644 index 44e040f9531607e81dae7bd129d13ac88128a2cd..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/BatchNorm3d.py +++ /dev/null @@ -1,151 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - - -import os -import sys -import torch -import importlib.util -import math -if importlib.util.find_spec("unsloth_studio") is None: - UNSLOTH_STUDIO_ENABLED = False -else: - UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" -pass -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -import math - -UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" -UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" -UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) -UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") -if UNSLOTH_COMPILE_LOCATION not in sys.path: - sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) - -import logging -logger_compiler = logging.getLogger(__name__) -if UNSLOTH_ENABLE_LOGGING: - logger_compiler.setLevel(logging.DEBUG) - -global INFERENCE_RUNS -INFERENCE_RUNS = 0 - -try: - import torch._dynamo.eval_frame as torch_dynamo_eval_frame - torch_dynamo_eval_frame._stance.stance - torch_compiler_set_stance = torch.compiler.set_stance -except: - torch_dynamo_eval_frame = None - torch_compiler_set_stance = None -pass - -# The stance is process wide, so whether Unsloth owns it lives in unsloth_zoo, shared by -# every generated module: inference in one model's module, training in another's. -from unsloth_zoo.temporary_patches.utils import UNSLOTH_EAGER_STANCE_OWNED, unsloth_claim_eager_stance, unsloth_owns_stance, unsloth_release_stance - -def unsloth_training_stance(): - # eager_on_recompile is global: left on by generate() or eval forwards, every recompile - # in a later training step (new sequence length, model.train() without for_training) - # would silently run the eager body. Training forwards restore the default stance, - # unless the user chose eager_on_recompile themselves. Every set_stance call installs a - # fresh stance object, so ours is still active only if the current one `is` it; a user - # set_stance since then (even to eager_on_recompile again) is theirs and stays. - global INFERENCE_RUNS - INFERENCE_RUNS = 0 - if len(UNSLOTH_EAGER_STANCE_OWNED) != 0 and torch_dynamo_eval_frame is not None: - current = torch_dynamo_eval_frame._stance - # Ownership is only dropped once ours is reset: a temporary stance scope's exit restores - # one of our exact objects, which a later training forward must still recognise. - if current.stance == "eager_on_recompile" and unsloth_owns_stance(current): - # Only the reset object: an outer scope may still restore another one of ours. - # Put back what the caller had before Unsloth switched (force_eager, a backend, ...). - previous = unsloth_release_stance(current) - if previous is not None and hasattr(torch_dynamo_eval_frame, "_set_stance"): - torch_dynamo_eval_frame._set_stance(previous) - else: - torch_compiler_set_stance(stance = "default", skip_guard_eval_unsafe = False) -pass - -from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT -from unsloth_zoo.temporary_patches.utils import UNSLOTH_DECODE_COMPILE - -torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 32, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': False, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from transformers.models.siglip.modeling_siglip import (np, nn) - -def forward(self, input: Tensor) -> Tensor: - original_dtype = input.dtype - if self.weight is not None: input = input.to(self.weight.dtype) - self._check_input_dim(input) - - # exponential_average_factor is set to self.momentum - # (when it is available) only so that it gets updated - # in ONNX graph when this node is exported to ONNX. - if self.momentum is None: - exponential_average_factor = 0.0 - else: - exponential_average_factor = self.momentum - - if self.training and self.track_running_stats: - # TODO: if statement only here to tell the jit to skip emitting this when it is None - if self.num_batches_tracked is not None: # type: ignore[has-type] - self.num_batches_tracked.add_(1) # type: ignore[has-type] - if self.momentum is None: # use cumulative moving average - exponential_average_factor = 1.0 / float(self.num_batches_tracked) - else: # use exponential moving average - exponential_average_factor = self.momentum - - r""" - Decide whether the mini-batch stats should be used for normalization rather than the buffers. - Mini-batch stats are used in training mode, and in eval mode when buffers are None. - """ - if self.training: - bn_training = True - else: - bn_training = (self.running_mean is None) and (self.running_var is None) - - r""" - Buffers are only updated if they are to be tracked and we are in training mode. Thus they only need to be - passed when the update should occur (i.e. in training mode when they are tracked), or when buffer stats are - used for normalization (i.e. in eval mode when buffers are not None). - """ - return F.batch_norm( - input, - # If buffers are not to be tracked, ensure that they won't be updated - ( - self.running_mean - if not self.training or self.track_running_stats - else None - ), - self.running_var if not self.training or self.track_running_stats else None, - self.weight, - self.bias, - bn_training, - exponential_average_factor, - self.eps, - ).to(original_dtype) diff --git a/unsloth_compiled_cache/BlockDiagonalLinear_peft_forward.py b/unsloth_compiled_cache/BlockDiagonalLinear_peft_forward.py deleted file mode 100644 index ad7982613d10d8ae2cf594c7f6f292643233a784..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/BlockDiagonalLinear_peft_forward.py +++ /dev/null @@ -1,85 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - - -torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 32, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': False, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from unsloth_zoo.temporary_patches.common import torch_compile -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from peft.tuners.lora.variants import (F, torch) - - -torch_addmm = torch.addmm -torch_add = torch.add -torch_is_compiling = torch.compiler.is_compiling -# torch 2.8 / 2.9 inductor miscompiles addmm(alpha = scaling, beta = 1) traced inside a -# compiled block (e.g. the VLM decoder MLP), forward and backward; fixed in 2.10 (unsloth#3271). -torch_lora_no_compiled_addmm = tuple(int(v) for v in torch.__version__.split(".")[:2]) < (2, 10) -# @torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) -def lora_forward(result, lora_A, lora_B, dropout, x, scaling): - # Use result.dtype (bfloat16 from base layer) since x may have been cast to float32 - # by _cast_input_dtype when autocast is disabled - target_dtype = result.dtype - xA = dropout(x).to(target_dtype) @ lora_A.weight.to(target_dtype).t() - if torch_lora_no_compiled_addmm and torch_is_compiling(): - output = result + scaling * (xA @ lora_B.weight.to(target_dtype).t()) - bias = lora_B.bias - if bias is not None: - output = output + scaling * bias.to(target_dtype) - return output - # output = result + scaling * xA @ lora_B.weight.t() - shape = result.shape - output = torch_addmm( - result.view(-1, shape[-1]), - xA.view(-1, xA.shape[-1]), - lora_B.weight.to(target_dtype).t(), - alpha = scaling, - beta = 1, - ).view(shape) - - bias = lora_B.bias - if bias is not None: - output = torch_add( - output, - bias.to(target_dtype), - alpha = scaling, - ) - return output -pass - -def unsloth_forward(self, x: torch.Tensor) -> torch.Tensor: - first_dims = x.shape[:-1] - if x.dim() != 2: - x = x.reshape(-1, x.shape[-1]) - B = x.shape[0] - nb = self.nblocks - m = x.shape[-1] // nb - n = self.out_features // nb - x = x.reshape(B, nb, m) - w = self.weight.view(nb, n, m) - out = torch.einsum("bim,inm->bin", x, w) - return out.reshape(*first_dims, -1) diff --git a/unsloth_compiled_cache/Conv1d.py b/unsloth_compiled_cache/Conv1d.py deleted file mode 100644 index 0662e91f7ef326d3f6a111b64c426b4c3c7e6591..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/Conv1d.py +++ /dev/null @@ -1,104 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - - -import os -import sys -import torch -import importlib.util -import math -if importlib.util.find_spec("unsloth_studio") is None: - UNSLOTH_STUDIO_ENABLED = False -else: - UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" -pass -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -import math - -UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" -UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" -UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) -UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") -if UNSLOTH_COMPILE_LOCATION not in sys.path: - sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) - -import logging -logger_compiler = logging.getLogger(__name__) -if UNSLOTH_ENABLE_LOGGING: - logger_compiler.setLevel(logging.DEBUG) - -global INFERENCE_RUNS -INFERENCE_RUNS = 0 - -try: - import torch._dynamo.eval_frame as torch_dynamo_eval_frame - torch_dynamo_eval_frame._stance.stance - torch_compiler_set_stance = torch.compiler.set_stance -except: - torch_dynamo_eval_frame = None - torch_compiler_set_stance = None -pass - -# The stance is process wide, so whether Unsloth owns it lives in unsloth_zoo, shared by -# every generated module: inference in one model's module, training in another's. -from unsloth_zoo.temporary_patches.utils import UNSLOTH_EAGER_STANCE_OWNED, unsloth_claim_eager_stance, unsloth_owns_stance, unsloth_release_stance - -def unsloth_training_stance(): - # eager_on_recompile is global: left on by generate() or eval forwards, every recompile - # in a later training step (new sequence length, model.train() without for_training) - # would silently run the eager body. Training forwards restore the default stance, - # unless the user chose eager_on_recompile themselves. Every set_stance call installs a - # fresh stance object, so ours is still active only if the current one `is` it; a user - # set_stance since then (even to eager_on_recompile again) is theirs and stays. - global INFERENCE_RUNS - INFERENCE_RUNS = 0 - if len(UNSLOTH_EAGER_STANCE_OWNED) != 0 and torch_dynamo_eval_frame is not None: - current = torch_dynamo_eval_frame._stance - # Ownership is only dropped once ours is reset: a temporary stance scope's exit restores - # one of our exact objects, which a later training forward must still recognise. - if current.stance == "eager_on_recompile" and unsloth_owns_stance(current): - # Only the reset object: an outer scope may still restore another one of ours. - # Put back what the caller had before Unsloth switched (force_eager, a backend, ...). - previous = unsloth_release_stance(current) - if previous is not None and hasattr(torch_dynamo_eval_frame, "_set_stance"): - torch_dynamo_eval_frame._set_stance(previous) - else: - torch_compiler_set_stance(stance = "default", skip_guard_eval_unsafe = False) -pass - -from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT -from unsloth_zoo.temporary_patches.utils import UNSLOTH_DECODE_COMPILE - -torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 32, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': False, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from transformers.models.siglip.modeling_siglip import (np) - -def forward(self, input: Tensor) -> Tensor: - original_dtype = input.dtype - input = input.to(self.weight.dtype) - return self._conv_forward(input, self.weight, self.bias).to(original_dtype) diff --git a/unsloth_compiled_cache/Conv2d.py b/unsloth_compiled_cache/Conv2d.py deleted file mode 100644 index 0662e91f7ef326d3f6a111b64c426b4c3c7e6591..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/Conv2d.py +++ /dev/null @@ -1,104 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - - -import os -import sys -import torch -import importlib.util -import math -if importlib.util.find_spec("unsloth_studio") is None: - UNSLOTH_STUDIO_ENABLED = False -else: - UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" -pass -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -import math - -UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" -UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" -UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) -UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") -if UNSLOTH_COMPILE_LOCATION not in sys.path: - sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) - -import logging -logger_compiler = logging.getLogger(__name__) -if UNSLOTH_ENABLE_LOGGING: - logger_compiler.setLevel(logging.DEBUG) - -global INFERENCE_RUNS -INFERENCE_RUNS = 0 - -try: - import torch._dynamo.eval_frame as torch_dynamo_eval_frame - torch_dynamo_eval_frame._stance.stance - torch_compiler_set_stance = torch.compiler.set_stance -except: - torch_dynamo_eval_frame = None - torch_compiler_set_stance = None -pass - -# The stance is process wide, so whether Unsloth owns it lives in unsloth_zoo, shared by -# every generated module: inference in one model's module, training in another's. -from unsloth_zoo.temporary_patches.utils import UNSLOTH_EAGER_STANCE_OWNED, unsloth_claim_eager_stance, unsloth_owns_stance, unsloth_release_stance - -def unsloth_training_stance(): - # eager_on_recompile is global: left on by generate() or eval forwards, every recompile - # in a later training step (new sequence length, model.train() without for_training) - # would silently run the eager body. Training forwards restore the default stance, - # unless the user chose eager_on_recompile themselves. Every set_stance call installs a - # fresh stance object, so ours is still active only if the current one `is` it; a user - # set_stance since then (even to eager_on_recompile again) is theirs and stays. - global INFERENCE_RUNS - INFERENCE_RUNS = 0 - if len(UNSLOTH_EAGER_STANCE_OWNED) != 0 and torch_dynamo_eval_frame is not None: - current = torch_dynamo_eval_frame._stance - # Ownership is only dropped once ours is reset: a temporary stance scope's exit restores - # one of our exact objects, which a later training forward must still recognise. - if current.stance == "eager_on_recompile" and unsloth_owns_stance(current): - # Only the reset object: an outer scope may still restore another one of ours. - # Put back what the caller had before Unsloth switched (force_eager, a backend, ...). - previous = unsloth_release_stance(current) - if previous is not None and hasattr(torch_dynamo_eval_frame, "_set_stance"): - torch_dynamo_eval_frame._set_stance(previous) - else: - torch_compiler_set_stance(stance = "default", skip_guard_eval_unsafe = False) -pass - -from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT -from unsloth_zoo.temporary_patches.utils import UNSLOTH_DECODE_COMPILE - -torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 32, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': False, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from transformers.models.siglip.modeling_siglip import (np) - -def forward(self, input: Tensor) -> Tensor: - original_dtype = input.dtype - input = input.to(self.weight.dtype) - return self._conv_forward(input, self.weight, self.bias).to(original_dtype) diff --git a/unsloth_compiled_cache/Conv3d.py b/unsloth_compiled_cache/Conv3d.py deleted file mode 100644 index 0662e91f7ef326d3f6a111b64c426b4c3c7e6591..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/Conv3d.py +++ /dev/null @@ -1,104 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - - -import os -import sys -import torch -import importlib.util -import math -if importlib.util.find_spec("unsloth_studio") is None: - UNSLOTH_STUDIO_ENABLED = False -else: - UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" -pass -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -import math - -UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" -UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" -UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) -UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") -if UNSLOTH_COMPILE_LOCATION not in sys.path: - sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) - -import logging -logger_compiler = logging.getLogger(__name__) -if UNSLOTH_ENABLE_LOGGING: - logger_compiler.setLevel(logging.DEBUG) - -global INFERENCE_RUNS -INFERENCE_RUNS = 0 - -try: - import torch._dynamo.eval_frame as torch_dynamo_eval_frame - torch_dynamo_eval_frame._stance.stance - torch_compiler_set_stance = torch.compiler.set_stance -except: - torch_dynamo_eval_frame = None - torch_compiler_set_stance = None -pass - -# The stance is process wide, so whether Unsloth owns it lives in unsloth_zoo, shared by -# every generated module: inference in one model's module, training in another's. -from unsloth_zoo.temporary_patches.utils import UNSLOTH_EAGER_STANCE_OWNED, unsloth_claim_eager_stance, unsloth_owns_stance, unsloth_release_stance - -def unsloth_training_stance(): - # eager_on_recompile is global: left on by generate() or eval forwards, every recompile - # in a later training step (new sequence length, model.train() without for_training) - # would silently run the eager body. Training forwards restore the default stance, - # unless the user chose eager_on_recompile themselves. Every set_stance call installs a - # fresh stance object, so ours is still active only if the current one `is` it; a user - # set_stance since then (even to eager_on_recompile again) is theirs and stays. - global INFERENCE_RUNS - INFERENCE_RUNS = 0 - if len(UNSLOTH_EAGER_STANCE_OWNED) != 0 and torch_dynamo_eval_frame is not None: - current = torch_dynamo_eval_frame._stance - # Ownership is only dropped once ours is reset: a temporary stance scope's exit restores - # one of our exact objects, which a later training forward must still recognise. - if current.stance == "eager_on_recompile" and unsloth_owns_stance(current): - # Only the reset object: an outer scope may still restore another one of ours. - # Put back what the caller had before Unsloth switched (force_eager, a backend, ...). - previous = unsloth_release_stance(current) - if previous is not None and hasattr(torch_dynamo_eval_frame, "_set_stance"): - torch_dynamo_eval_frame._set_stance(previous) - else: - torch_compiler_set_stance(stance = "default", skip_guard_eval_unsafe = False) -pass - -from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT -from unsloth_zoo.temporary_patches.utils import UNSLOTH_DECODE_COMPILE - -torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 32, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': False, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from transformers.models.siglip.modeling_siglip import (np) - -def forward(self, input: Tensor) -> Tensor: - original_dtype = input.dtype - input = input.to(self.weight.dtype) - return self._conv_forward(input, self.weight, self.bias).to(original_dtype) diff --git a/unsloth_compiled_cache/ConvTranspose1d.py b/unsloth_compiled_cache/ConvTranspose1d.py deleted file mode 100644 index 6c4af18fe9408b0af8a8d788a4a89977e56db6c9..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/ConvTranspose1d.py +++ /dev/null @@ -1,132 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - - -import os -import sys -import torch -import importlib.util -import math -if importlib.util.find_spec("unsloth_studio") is None: - UNSLOTH_STUDIO_ENABLED = False -else: - UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" -pass -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -import math - -UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" -UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" -UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) -UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") -if UNSLOTH_COMPILE_LOCATION not in sys.path: - sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) - -import logging -logger_compiler = logging.getLogger(__name__) -if UNSLOTH_ENABLE_LOGGING: - logger_compiler.setLevel(logging.DEBUG) - -global INFERENCE_RUNS -INFERENCE_RUNS = 0 - -try: - import torch._dynamo.eval_frame as torch_dynamo_eval_frame - torch_dynamo_eval_frame._stance.stance - torch_compiler_set_stance = torch.compiler.set_stance -except: - torch_dynamo_eval_frame = None - torch_compiler_set_stance = None -pass - -# The stance is process wide, so whether Unsloth owns it lives in unsloth_zoo, shared by -# every generated module: inference in one model's module, training in another's. -from unsloth_zoo.temporary_patches.utils import UNSLOTH_EAGER_STANCE_OWNED, unsloth_claim_eager_stance, unsloth_owns_stance, unsloth_release_stance - -def unsloth_training_stance(): - # eager_on_recompile is global: left on by generate() or eval forwards, every recompile - # in a later training step (new sequence length, model.train() without for_training) - # would silently run the eager body. Training forwards restore the default stance, - # unless the user chose eager_on_recompile themselves. Every set_stance call installs a - # fresh stance object, so ours is still active only if the current one `is` it; a user - # set_stance since then (even to eager_on_recompile again) is theirs and stays. - global INFERENCE_RUNS - INFERENCE_RUNS = 0 - if len(UNSLOTH_EAGER_STANCE_OWNED) != 0 and torch_dynamo_eval_frame is not None: - current = torch_dynamo_eval_frame._stance - # Ownership is only dropped once ours is reset: a temporary stance scope's exit restores - # one of our exact objects, which a later training forward must still recognise. - if current.stance == "eager_on_recompile" and unsloth_owns_stance(current): - # Only the reset object: an outer scope may still restore another one of ours. - # Put back what the caller had before Unsloth switched (force_eager, a backend, ...). - previous = unsloth_release_stance(current) - if previous is not None and hasattr(torch_dynamo_eval_frame, "_set_stance"): - torch_dynamo_eval_frame._set_stance(previous) - else: - torch_compiler_set_stance(stance = "default", skip_guard_eval_unsafe = False) -pass - -from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT -from unsloth_zoo.temporary_patches.utils import UNSLOTH_DECODE_COMPILE - -torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 32, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': False, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from transformers.models.siglip.modeling_siglip import (np, nn) - -def forward(self, input: Tensor, output_size: list[int] | None = None) -> Tensor: - original_dtype = input.dtype - input = input.to(self.weight.dtype) - if self.padding_mode != "zeros": - raise ValueError( - "Only `zeros` padding mode is supported for ConvTranspose1d" - ) - - if not isinstance(self.padding, tuple): - raise AssertionError("self.padding must be a tuple") - # One cannot replace List by Tuple or Sequence in "_output_padding" because - # TorchScript does not support `Sequence[T]` or `Tuple[T, ...]`. - num_spatial_dims = 1 - output_padding = self._output_padding( - input, - output_size, - self.stride, # type: ignore[arg-type] - self.padding, # type: ignore[arg-type] - self.kernel_size, # type: ignore[arg-type] - num_spatial_dims, - self.dilation, # type: ignore[arg-type] - ) - return F.conv_transpose1d( - input, - self.weight, - self.bias, - self.stride, - self.padding, - output_padding, - self.groups, - self.dilation, - ).to(original_dtype) diff --git a/unsloth_compiled_cache/ConvTranspose2d.py b/unsloth_compiled_cache/ConvTranspose2d.py deleted file mode 100644 index 3730b0def5d4b18eb581e7c39cabfa6a016bf541..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/ConvTranspose2d.py +++ /dev/null @@ -1,141 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - - -import os -import sys -import torch -import importlib.util -import math -if importlib.util.find_spec("unsloth_studio") is None: - UNSLOTH_STUDIO_ENABLED = False -else: - UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" -pass -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -import math - -UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" -UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" -UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) -UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") -if UNSLOTH_COMPILE_LOCATION not in sys.path: - sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) - -import logging -logger_compiler = logging.getLogger(__name__) -if UNSLOTH_ENABLE_LOGGING: - logger_compiler.setLevel(logging.DEBUG) - -global INFERENCE_RUNS -INFERENCE_RUNS = 0 - -try: - import torch._dynamo.eval_frame as torch_dynamo_eval_frame - torch_dynamo_eval_frame._stance.stance - torch_compiler_set_stance = torch.compiler.set_stance -except: - torch_dynamo_eval_frame = None - torch_compiler_set_stance = None -pass - -# The stance is process wide, so whether Unsloth owns it lives in unsloth_zoo, shared by -# every generated module: inference in one model's module, training in another's. -from unsloth_zoo.temporary_patches.utils import UNSLOTH_EAGER_STANCE_OWNED, unsloth_claim_eager_stance, unsloth_owns_stance, unsloth_release_stance - -def unsloth_training_stance(): - # eager_on_recompile is global: left on by generate() or eval forwards, every recompile - # in a later training step (new sequence length, model.train() without for_training) - # would silently run the eager body. Training forwards restore the default stance, - # unless the user chose eager_on_recompile themselves. Every set_stance call installs a - # fresh stance object, so ours is still active only if the current one `is` it; a user - # set_stance since then (even to eager_on_recompile again) is theirs and stays. - global INFERENCE_RUNS - INFERENCE_RUNS = 0 - if len(UNSLOTH_EAGER_STANCE_OWNED) != 0 and torch_dynamo_eval_frame is not None: - current = torch_dynamo_eval_frame._stance - # Ownership is only dropped once ours is reset: a temporary stance scope's exit restores - # one of our exact objects, which a later training forward must still recognise. - if current.stance == "eager_on_recompile" and unsloth_owns_stance(current): - # Only the reset object: an outer scope may still restore another one of ours. - # Put back what the caller had before Unsloth switched (force_eager, a backend, ...). - previous = unsloth_release_stance(current) - if previous is not None and hasattr(torch_dynamo_eval_frame, "_set_stance"): - torch_dynamo_eval_frame._set_stance(previous) - else: - torch_compiler_set_stance(stance = "default", skip_guard_eval_unsafe = False) -pass - -from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT -from unsloth_zoo.temporary_patches.utils import UNSLOTH_DECODE_COMPILE - -torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 32, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': False, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from transformers.models.siglip.modeling_siglip import (np, nn) - -def forward(self, input: Tensor, output_size: list[int] | None = None) -> Tensor: - original_dtype = input.dtype - input = input.to(self.weight.dtype) - """ - Performs the forward pass. - - Attributes: - input (Tensor): The input tensor. - output_size (list[int], optional): A list of integers representing - the size of the output tensor. Default is None. - """ - if self.padding_mode != "zeros": - raise ValueError( - "Only `zeros` padding mode is supported for ConvTranspose2d" - ) - - if not isinstance(self.padding, tuple): - raise AssertionError("self.padding must be a tuple") - # One cannot replace List by Tuple or Sequence in "_output_padding" because - # TorchScript does not support `Sequence[T]` or `Tuple[T, ...]`. - num_spatial_dims = 2 - output_padding = self._output_padding( - input, - output_size, - self.stride, # type: ignore[arg-type] - self.padding, # type: ignore[arg-type] - self.kernel_size, # type: ignore[arg-type] - num_spatial_dims, - self.dilation, # type: ignore[arg-type] - ) - - return F.conv_transpose2d( - input, - self.weight, - self.bias, - self.stride, - self.padding, - output_padding, - self.groups, - self.dilation, - ).to(original_dtype) diff --git a/unsloth_compiled_cache/ConvTranspose3d.py b/unsloth_compiled_cache/ConvTranspose3d.py deleted file mode 100644 index 589f3869d332c44c0704e4d78041f6f094abfb78..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/ConvTranspose3d.py +++ /dev/null @@ -1,133 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - - -import os -import sys -import torch -import importlib.util -import math -if importlib.util.find_spec("unsloth_studio") is None: - UNSLOTH_STUDIO_ENABLED = False -else: - UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" -pass -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -import math - -UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" -UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" -UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) -UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") -if UNSLOTH_COMPILE_LOCATION not in sys.path: - sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) - -import logging -logger_compiler = logging.getLogger(__name__) -if UNSLOTH_ENABLE_LOGGING: - logger_compiler.setLevel(logging.DEBUG) - -global INFERENCE_RUNS -INFERENCE_RUNS = 0 - -try: - import torch._dynamo.eval_frame as torch_dynamo_eval_frame - torch_dynamo_eval_frame._stance.stance - torch_compiler_set_stance = torch.compiler.set_stance -except: - torch_dynamo_eval_frame = None - torch_compiler_set_stance = None -pass - -# The stance is process wide, so whether Unsloth owns it lives in unsloth_zoo, shared by -# every generated module: inference in one model's module, training in another's. -from unsloth_zoo.temporary_patches.utils import UNSLOTH_EAGER_STANCE_OWNED, unsloth_claim_eager_stance, unsloth_owns_stance, unsloth_release_stance - -def unsloth_training_stance(): - # eager_on_recompile is global: left on by generate() or eval forwards, every recompile - # in a later training step (new sequence length, model.train() without for_training) - # would silently run the eager body. Training forwards restore the default stance, - # unless the user chose eager_on_recompile themselves. Every set_stance call installs a - # fresh stance object, so ours is still active only if the current one `is` it; a user - # set_stance since then (even to eager_on_recompile again) is theirs and stays. - global INFERENCE_RUNS - INFERENCE_RUNS = 0 - if len(UNSLOTH_EAGER_STANCE_OWNED) != 0 and torch_dynamo_eval_frame is not None: - current = torch_dynamo_eval_frame._stance - # Ownership is only dropped once ours is reset: a temporary stance scope's exit restores - # one of our exact objects, which a later training forward must still recognise. - if current.stance == "eager_on_recompile" and unsloth_owns_stance(current): - # Only the reset object: an outer scope may still restore another one of ours. - # Put back what the caller had before Unsloth switched (force_eager, a backend, ...). - previous = unsloth_release_stance(current) - if previous is not None and hasattr(torch_dynamo_eval_frame, "_set_stance"): - torch_dynamo_eval_frame._set_stance(previous) - else: - torch_compiler_set_stance(stance = "default", skip_guard_eval_unsafe = False) -pass - -from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT -from unsloth_zoo.temporary_patches.utils import UNSLOTH_DECODE_COMPILE - -torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 32, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': False, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from transformers.models.siglip.modeling_siglip import (np, nn) - -def forward(self, input: Tensor, output_size: list[int] | None = None) -> Tensor: - original_dtype = input.dtype - input = input.to(self.weight.dtype) - if self.padding_mode != "zeros": - raise ValueError( - "Only `zeros` padding mode is supported for ConvTranspose3d" - ) - - if not isinstance(self.padding, tuple): - raise AssertionError("self.padding must be a tuple") - # One cannot replace List by Tuple or Sequence in "_output_padding" because - # TorchScript does not support `Sequence[T]` or `Tuple[T, ...]`. - num_spatial_dims = 3 - output_padding = self._output_padding( - input, - output_size, - self.stride, # type: ignore[arg-type] - self.padding, # type: ignore[arg-type] - self.kernel_size, # type: ignore[arg-type] - num_spatial_dims, - self.dilation, # type: ignore[arg-type] - ) - - return F.conv_transpose3d( - input, - self.weight, - self.bias, - self.stride, - self.padding, - output_padding, - self.groups, - self.dilation, - ).to(original_dtype) diff --git a/unsloth_compiled_cache/GPTQLoraLinear_peft_forward.py b/unsloth_compiled_cache/GPTQLoraLinear_peft_forward.py deleted file mode 100644 index 9cb2ef7f61040db84856dd17e3826a13371ba699..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/GPTQLoraLinear_peft_forward.py +++ /dev/null @@ -1,106 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - - -torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 32, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': False, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from unsloth_zoo.temporary_patches.common import torch_compile -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from peft.tuners.lora.gptq import (torch) - - -torch_addmm = torch.addmm -torch_add = torch.add -torch_is_compiling = torch.compiler.is_compiling -# torch 2.8 / 2.9 inductor miscompiles addmm(alpha = scaling, beta = 1) traced inside a -# compiled block (e.g. the VLM decoder MLP), forward and backward; fixed in 2.10 (unsloth#3271). -torch_lora_no_compiled_addmm = tuple(int(v) for v in torch.__version__.split(".")[:2]) < (2, 10) -# @torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) -def lora_forward(result, lora_A, lora_B, dropout, x, scaling): - # Use result.dtype (bfloat16 from base layer) since x may have been cast to float32 - # by _cast_input_dtype when autocast is disabled - target_dtype = result.dtype - xA = dropout(x).to(target_dtype) @ lora_A.weight.to(target_dtype).t() - if torch_lora_no_compiled_addmm and torch_is_compiling(): - output = result + scaling * (xA @ lora_B.weight.to(target_dtype).t()) - bias = lora_B.bias - if bias is not None: - output = output + scaling * bias.to(target_dtype) - return output - # output = result + scaling * xA @ lora_B.weight.t() - shape = result.shape - output = torch_addmm( - result.view(-1, shape[-1]), - xA.view(-1, xA.shape[-1]), - lora_B.weight.to(target_dtype).t(), - alpha = scaling, - beta = 1, - ).view(shape) - - bias = lora_B.bias - if bias is not None: - output = torch_add( - output, - bias.to(target_dtype), - alpha = scaling, - ) - return output -pass - -def unsloth_forward(self, x: torch.Tensor): - # note: logic differs from default Linear because merging is not supported - result = self.quant_linear_module(x) - - if self.disable_adapters: - return result - - lora_A_keys = self.lora_A.keys() - - for active_adapter in self.active_adapters: - if active_adapter not in lora_A_keys: - continue - torch_result_dtype = result.dtype - - lora_A = self.lora_A[active_adapter] - lora_B = self.lora_B[active_adapter] - dropout = self.lora_dropout[active_adapter] - scaling = self.scaling[active_adapter] - - if not torch.is_autocast_enabled(): result, x = result.to(lora_A.weight.dtype), x.to(lora_A.weight.dtype) - - if active_adapter not in self.lora_variant: # vanilla LoRA - return lora_forward(result, lora_A, lora_B, dropout, x, scaling).to(torch_result_dtype) - else: - result = self.lora_variant[active_adapter].forward( - self, - active_adapter=active_adapter, - x=x, - result=result, - ) - - result = result.to(torch_result_dtype) - return result diff --git a/unsloth_compiled_cache/GroupNorm.py b/unsloth_compiled_cache/GroupNorm.py deleted file mode 100644 index 6f3919c480d92173a17914ac0cfe3852c5c644db..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/GroupNorm.py +++ /dev/null @@ -1,102 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - - -import os -import sys -import torch -import importlib.util -import math -if importlib.util.find_spec("unsloth_studio") is None: - UNSLOTH_STUDIO_ENABLED = False -else: - UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" -pass -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -import math - -UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" -UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" -UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) -UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") -if UNSLOTH_COMPILE_LOCATION not in sys.path: - sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) - -import logging -logger_compiler = logging.getLogger(__name__) -if UNSLOTH_ENABLE_LOGGING: - logger_compiler.setLevel(logging.DEBUG) - -global INFERENCE_RUNS -INFERENCE_RUNS = 0 - -try: - import torch._dynamo.eval_frame as torch_dynamo_eval_frame - torch_dynamo_eval_frame._stance.stance - torch_compiler_set_stance = torch.compiler.set_stance -except: - torch_dynamo_eval_frame = None - torch_compiler_set_stance = None -pass - -# The stance is process wide, so whether Unsloth owns it lives in unsloth_zoo, shared by -# every generated module: inference in one model's module, training in another's. -from unsloth_zoo.temporary_patches.utils import UNSLOTH_EAGER_STANCE_OWNED, unsloth_claim_eager_stance, unsloth_owns_stance, unsloth_release_stance - -def unsloth_training_stance(): - # eager_on_recompile is global: left on by generate() or eval forwards, every recompile - # in a later training step (new sequence length, model.train() without for_training) - # would silently run the eager body. Training forwards restore the default stance, - # unless the user chose eager_on_recompile themselves. Every set_stance call installs a - # fresh stance object, so ours is still active only if the current one `is` it; a user - # set_stance since then (even to eager_on_recompile again) is theirs and stays. - global INFERENCE_RUNS - INFERENCE_RUNS = 0 - if len(UNSLOTH_EAGER_STANCE_OWNED) != 0 and torch_dynamo_eval_frame is not None: - current = torch_dynamo_eval_frame._stance - # Ownership is only dropped once ours is reset: a temporary stance scope's exit restores - # one of our exact objects, which a later training forward must still recognise. - if current.stance == "eager_on_recompile" and unsloth_owns_stance(current): - # Only the reset object: an outer scope may still restore another one of ours. - # Put back what the caller had before Unsloth switched (force_eager, a backend, ...). - previous = unsloth_release_stance(current) - if previous is not None and hasattr(torch_dynamo_eval_frame, "_set_stance"): - torch_dynamo_eval_frame._set_stance(previous) - else: - torch_compiler_set_stance(stance = "default", skip_guard_eval_unsafe = False) -pass - -from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT -from unsloth_zoo.temporary_patches.utils import UNSLOTH_DECODE_COMPILE - -torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 32, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': False, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from transformers.models.siglip.modeling_siglip import (np) - -def forward(self, input: Tensor) -> Tensor: - return F.group_norm(input, self.num_groups, self.weight, self.bias, self.eps).to(input.dtype) diff --git a/unsloth_compiled_cache/LayerNorm.py b/unsloth_compiled_cache/LayerNorm.py deleted file mode 100644 index 774449c097e578e067f7bd0535c6a4d5795206e0..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/LayerNorm.py +++ /dev/null @@ -1,104 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - - -import os -import sys -import torch -import importlib.util -import math -if importlib.util.find_spec("unsloth_studio") is None: - UNSLOTH_STUDIO_ENABLED = False -else: - UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" -pass -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -import math - -UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" -UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" -UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) -UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") -if UNSLOTH_COMPILE_LOCATION not in sys.path: - sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) - -import logging -logger_compiler = logging.getLogger(__name__) -if UNSLOTH_ENABLE_LOGGING: - logger_compiler.setLevel(logging.DEBUG) - -global INFERENCE_RUNS -INFERENCE_RUNS = 0 - -try: - import torch._dynamo.eval_frame as torch_dynamo_eval_frame - torch_dynamo_eval_frame._stance.stance - torch_compiler_set_stance = torch.compiler.set_stance -except: - torch_dynamo_eval_frame = None - torch_compiler_set_stance = None -pass - -# The stance is process wide, so whether Unsloth owns it lives in unsloth_zoo, shared by -# every generated module: inference in one model's module, training in another's. -from unsloth_zoo.temporary_patches.utils import UNSLOTH_EAGER_STANCE_OWNED, unsloth_claim_eager_stance, unsloth_owns_stance, unsloth_release_stance - -def unsloth_training_stance(): - # eager_on_recompile is global: left on by generate() or eval forwards, every recompile - # in a later training step (new sequence length, model.train() without for_training) - # would silently run the eager body. Training forwards restore the default stance, - # unless the user chose eager_on_recompile themselves. Every set_stance call installs a - # fresh stance object, so ours is still active only if the current one `is` it; a user - # set_stance since then (even to eager_on_recompile again) is theirs and stays. - global INFERENCE_RUNS - INFERENCE_RUNS = 0 - if len(UNSLOTH_EAGER_STANCE_OWNED) != 0 and torch_dynamo_eval_frame is not None: - current = torch_dynamo_eval_frame._stance - # Ownership is only dropped once ours is reset: a temporary stance scope's exit restores - # one of our exact objects, which a later training forward must still recognise. - if current.stance == "eager_on_recompile" and unsloth_owns_stance(current): - # Only the reset object: an outer scope may still restore another one of ours. - # Put back what the caller had before Unsloth switched (force_eager, a backend, ...). - previous = unsloth_release_stance(current) - if previous is not None and hasattr(torch_dynamo_eval_frame, "_set_stance"): - torch_dynamo_eval_frame._set_stance(previous) - else: - torch_compiler_set_stance(stance = "default", skip_guard_eval_unsafe = False) -pass - -from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT -from unsloth_zoo.temporary_patches.utils import UNSLOTH_DECODE_COMPILE - -torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 32, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': False, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from transformers.models.siglip.modeling_siglip import (np) - -def forward(self, input: Tensor) -> Tensor: - return F.layer_norm( - input, self.normalized_shape, self.weight, self.bias, self.eps - ).to(input.dtype) diff --git a/unsloth_compiled_cache/Linear4bit_peft_forward.py b/unsloth_compiled_cache/Linear4bit_peft_forward.py deleted file mode 100644 index 3b962abdc39085b34be7a421a6118d690871392c..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/Linear4bit_peft_forward.py +++ /dev/null @@ -1,138 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - - -try: - from peft.tuners.lora.layer import VARIANT_KWARG_KEYS -except ImportError: - VARIANT_KWARG_KEYS = ['alora_offsets'] -torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 32, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': False, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} - -from unsloth_zoo.temporary_patches.misc import _lora_integer_input as _unsloth_lora_integer_input -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from unsloth_zoo.temporary_patches.common import torch_compile -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from peft.tuners.lora.bnb import (VARIANT_KWARG_KEYS, torch) - - -torch_addmm = torch.addmm -torch_add = torch.add -torch_is_compiling = torch.compiler.is_compiling -# torch 2.8 / 2.9 inductor miscompiles addmm(alpha = scaling, beta = 1) traced inside a -# compiled block (e.g. the VLM decoder MLP), forward and backward; fixed in 2.10 (unsloth#3271). -torch_lora_no_compiled_addmm = tuple(int(v) for v in torch.__version__.split(".")[:2]) < (2, 10) -# @torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) -def lora_forward(result, lora_A, lora_B, dropout, x, scaling): - # Use result.dtype (bfloat16 from base layer) since x may have been cast to float32 - # by _cast_input_dtype when autocast is disabled - target_dtype = result.dtype - xA = dropout(x).to(target_dtype) @ lora_A.weight.to(target_dtype).t() - if torch_lora_no_compiled_addmm and torch_is_compiling(): - output = result + scaling * (xA @ lora_B.weight.to(target_dtype).t()) - bias = lora_B.bias - if bias is not None: - output = output + scaling * bias.to(target_dtype) - return output - # output = result + scaling * xA @ lora_B.weight.t() - shape = result.shape - output = torch_addmm( - result.view(-1, shape[-1]), - xA.view(-1, xA.shape[-1]), - lora_B.weight.to(target_dtype).t(), - alpha = scaling, - beta = 1, - ).view(shape) - - bias = lora_B.bias - if bias is not None: - output = torch_add( - output, - bias.to(target_dtype), - alpha = scaling, - ) - return output -pass - -def unsloth_forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor: - if not x.is_floating_point(): x = _unsloth_lora_integer_input(self, x) - adapter_names = kwargs.pop("adapter_names", None) - variant_kwargs = {k: kwargs.pop(k, None) for k in VARIANT_KWARG_KEYS} # don't pass these to base_layer - - if self.disable_adapters: - if self.merged: - self.unmerge() - if not torch.is_autocast_enabled() and hasattr(self.base_layer, 'weight') and self.base_layer.weight is not None and not hasattr(self.base_layer.weight, 'quant_state') and self.base_layer.weight.is_floating_point() and self.base_layer.weight.element_size() > 1 and x.dtype != self.base_layer.weight.dtype: - x = x.to(self.base_layer.weight.dtype) - result = self.base_layer(x, *args, **kwargs) - elif adapter_names is not None: - result = self._mixed_batch_forward(x, *args, adapter_names=adapter_names, **variant_kwargs, **kwargs) - elif self.merged: - if not torch.is_autocast_enabled() and hasattr(self.base_layer, 'weight') and self.base_layer.weight is not None and not hasattr(self.base_layer.weight, 'quant_state') and self.base_layer.weight.is_floating_point() and self.base_layer.weight.element_size() > 1 and x.dtype != self.base_layer.weight.dtype: - x = x.to(self.base_layer.weight.dtype) - result = self.base_layer(x, *args, **kwargs) - else: - if not torch.is_autocast_enabled() and hasattr(self.base_layer, 'weight') and self.base_layer.weight is not None and not hasattr(self.base_layer.weight, 'quant_state') and self.base_layer.weight.is_floating_point() and self.base_layer.weight.element_size() > 1 and x.dtype != self.base_layer.weight.dtype: - x = x.to(self.base_layer.weight.dtype) - result = self.base_layer(x, *args, **kwargs) - # As per Tim Dettmers, for 4bit, we need to defensively clone here. - # The reason is that in some cases, an error can occur that backprop - # does not work on a manipulated view. This issue may be solved with - # newer PyTorch versions but this would need extensive testing to be - # sure. - - - for active_adapter in self.active_adapters: - if active_adapter not in self.lora_A.keys(): - continue - lora_A = self.lora_A[active_adapter] - lora_B = self.lora_B[active_adapter] - dropout = self.lora_dropout[active_adapter] - scaling = self.scaling[active_adapter] - - requires_conversion = not torch.is_autocast_enabled() - if requires_conversion: - expected_dtype = result.dtype - x = self._cast_input_dtype(x, lora_A.weight.dtype) - - if active_adapter not in self.lora_variant: # vanilla LoRA - return lora_forward(result, lora_A, lora_B, dropout, x, scaling).to(result.dtype) - if requires_conversion: - output = output.to(expected_dtype) - result = result + output - else: - result = self.lora_variant[active_adapter].forward( - self, - active_adapter=active_adapter, - x=x, - result=result, - **variant_kwargs, - **kwargs, - ) - if requires_conversion: - result = result.to(expected_dtype) - - return result diff --git a/unsloth_compiled_cache/Linear8bitLt_peft_forward.py b/unsloth_compiled_cache/Linear8bitLt_peft_forward.py deleted file mode 100644 index 5925d3a7ad733103d790fc273afb11c1da92c423..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/Linear8bitLt_peft_forward.py +++ /dev/null @@ -1,128 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - - -try: - from peft.tuners.lora.layer import VARIANT_KWARG_KEYS -except ImportError: - VARIANT_KWARG_KEYS = ['alora_offsets'] -torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 32, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': False, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} - -import torch._dynamo -@torch._dynamo.disable -def _call_8bit_base_layer(base_layer, x, *args, **kwargs): - return base_layer(x, *args, **kwargs) -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from unsloth_zoo.temporary_patches.common import torch_compile -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from peft.tuners.lora.bnb import (VARIANT_KWARG_KEYS, torch) - - -torch_addmm = torch.addmm -torch_add = torch.add -torch_is_compiling = torch.compiler.is_compiling -# torch 2.8 / 2.9 inductor miscompiles addmm(alpha = scaling, beta = 1) traced inside a -# compiled block (e.g. the VLM decoder MLP), forward and backward; fixed in 2.10 (unsloth#3271). -torch_lora_no_compiled_addmm = tuple(int(v) for v in torch.__version__.split(".")[:2]) < (2, 10) -# @torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) -def lora_forward(result, lora_A, lora_B, dropout, x, scaling): - # Use result.dtype (bfloat16 from base layer) since x may have been cast to float32 - # by _cast_input_dtype when autocast is disabled - target_dtype = result.dtype - xA = dropout(x).to(target_dtype) @ lora_A.weight.to(target_dtype).t() - if torch_lora_no_compiled_addmm and torch_is_compiling(): - output = result + scaling * (xA @ lora_B.weight.to(target_dtype).t()) - bias = lora_B.bias - if bias is not None: - output = output + scaling * bias.to(target_dtype) - return output - # output = result + scaling * xA @ lora_B.weight.t() - shape = result.shape - output = torch_addmm( - result.view(-1, shape[-1]), - xA.view(-1, xA.shape[-1]), - lora_B.weight.to(target_dtype).t(), - alpha = scaling, - beta = 1, - ).view(shape) - - bias = lora_B.bias - if bias is not None: - output = torch_add( - output, - bias.to(target_dtype), - alpha = scaling, - ) - return output -pass - -def unsloth_forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor: - - adapter_names = kwargs.pop("adapter_names", None) - variant_kwargs = {k: kwargs.pop(k, None) for k in VARIANT_KWARG_KEYS} # don't pass these to base_layer - - if self.disable_adapters: - if self.merged: - self.unmerge() - result = _call_8bit_base_layer(self.base_layer, x, *args, **kwargs) - elif adapter_names is not None: - result = self._mixed_batch_forward(x, *args, adapter_names=adapter_names, **variant_kwargs, **kwargs) - elif self.merged: - result = _call_8bit_base_layer(self.base_layer, x, *args, **kwargs) - else: - result = _call_8bit_base_layer(self.base_layer, x, *args, **kwargs) - for active_adapter in self.active_adapters: - if active_adapter not in self.lora_A.keys(): - continue - lora_A = self.lora_A[active_adapter] - lora_B = self.lora_B[active_adapter] - dropout = self.lora_dropout[active_adapter] - scaling = self.scaling[active_adapter] - - requires_conversion = not torch.is_autocast_enabled() - if requires_conversion: - expected_dtype = result.dtype - x = self._cast_input_dtype(x, lora_A.weight.dtype) - - if active_adapter not in self.lora_variant: # vanilla LoRA - return lora_forward(result, lora_A, lora_B, dropout, x, scaling).to(result.dtype) - if requires_conversion: - output = output.to(expected_dtype) - result = result + output - else: - result = self.lora_variant[active_adapter].forward( - self, - active_adapter=active_adapter, - x=x, - result=result, - **variant_kwargs, - **kwargs, - ) - if requires_conversion: - result = result.to(expected_dtype) - - return result diff --git a/unsloth_compiled_cache/Linear_peft_forward.py b/unsloth_compiled_cache/Linear_peft_forward.py deleted file mode 100644 index f07657bb461fb4e1a31ae39928b28076cf1f1992..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/Linear_peft_forward.py +++ /dev/null @@ -1,125 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - - -try: - from peft.tuners.lora.layer import VARIANT_KWARG_KEYS -except ImportError: - VARIANT_KWARG_KEYS = ['alora_offsets'] -torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 32, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': False, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from unsloth_zoo.temporary_patches.common import torch_compile -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from peft.tuners.lora.inc import (torch) - - -torch_addmm = torch.addmm -torch_add = torch.add -torch_is_compiling = torch.compiler.is_compiling -# torch 2.8 / 2.9 inductor miscompiles addmm(alpha = scaling, beta = 1) traced inside a -# compiled block (e.g. the VLM decoder MLP), forward and backward; fixed in 2.10 (unsloth#3271). -torch_lora_no_compiled_addmm = tuple(int(v) for v in torch.__version__.split(".")[:2]) < (2, 10) -# @torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) -def lora_forward(result, lora_A, lora_B, dropout, x, scaling): - # Use result.dtype (bfloat16 from base layer) since x may have been cast to float32 - # by _cast_input_dtype when autocast is disabled - target_dtype = result.dtype - xA = dropout(x).to(target_dtype) @ lora_A.weight.to(target_dtype).t() - if torch_lora_no_compiled_addmm and torch_is_compiling(): - output = result + scaling * (xA @ lora_B.weight.to(target_dtype).t()) - bias = lora_B.bias - if bias is not None: - output = output + scaling * bias.to(target_dtype) - return output - # output = result + scaling * xA @ lora_B.weight.t() - shape = result.shape - output = torch_addmm( - result.view(-1, shape[-1]), - xA.view(-1, xA.shape[-1]), - lora_B.weight.to(target_dtype).t(), - alpha = scaling, - beta = 1, - ).view(shape) - - bias = lora_B.bias - if bias is not None: - output = torch_add( - output, - bias.to(target_dtype), - alpha = scaling, - ) - return output -pass - -def unsloth_forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: - - adapter_names = kwargs.pop("adapter_names", None) - variant_kwargs = {k: kwargs.pop(k, None) for k in VARIANT_KWARG_KEYS} # don't pass these to base_layer - - if self.disable_adapters: - if self.merged: - self.unmerge() - if not torch.is_autocast_enabled() and hasattr(self.base_layer, 'weight') and self.base_layer.weight is not None and not hasattr(self.base_layer.weight, 'quant_state') and self.base_layer.weight.is_floating_point() and self.base_layer.weight.element_size() > 1 and x.dtype != self.base_layer.weight.dtype: - x = x.to(self.base_layer.weight.dtype) - result = self.base_layer(x, *args, **kwargs) - elif adapter_names is not None: - result = self._mixed_batch_forward(x, *args, adapter_names=adapter_names, **variant_kwargs, **kwargs) - elif self.merged: - if not torch.is_autocast_enabled() and hasattr(self.base_layer, 'weight') and self.base_layer.weight is not None and not hasattr(self.base_layer.weight, 'quant_state') and self.base_layer.weight.is_floating_point() and self.base_layer.weight.element_size() > 1 and x.dtype != self.base_layer.weight.dtype: - x = x.to(self.base_layer.weight.dtype) - result = self.base_layer(x, *args, **kwargs) - else: - if not torch.is_autocast_enabled() and hasattr(self.base_layer, 'weight') and self.base_layer.weight is not None and not hasattr(self.base_layer.weight, 'quant_state') and self.base_layer.weight.is_floating_point() and self.base_layer.weight.element_size() > 1 and x.dtype != self.base_layer.weight.dtype: - x = x.to(self.base_layer.weight.dtype) - result = self.base_layer(x, *args, **kwargs) - torch_result_dtype = result.dtype - - lora_A_keys = self.lora_A.keys() - for active_adapter in self.active_adapters: - if active_adapter not in lora_A_keys: - continue - - lora_A = self.lora_A[active_adapter] - lora_B = self.lora_B[active_adapter] - dropout = self.lora_dropout[active_adapter] - scaling = self.scaling[active_adapter] - if not torch.is_autocast_enabled(): result, x = result.to(lora_A.weight.dtype), x.to(lora_A.weight.dtype) - if active_adapter not in self.lora_variant: # vanilla LoRA - return lora_forward(result, lora_A, lora_B, dropout, x, scaling).to(torch_result_dtype) - else: - result = self.lora_variant[active_adapter].forward( - self, - active_adapter=active_adapter, - x=x, - result=result, - **variant_kwargs, - **kwargs, - ) - - result = result.to(torch_result_dtype) - - return result diff --git a/unsloth_compiled_cache/LoraParallelLinear_peft_forward.py b/unsloth_compiled_cache/LoraParallelLinear_peft_forward.py deleted file mode 100644 index 1b639f512c14b4594c02bc77f2e02d545fc152bc..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/LoraParallelLinear_peft_forward.py +++ /dev/null @@ -1,102 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - - -torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 32, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': False, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from unsloth_zoo.temporary_patches.common import torch_compile -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from peft.tuners.lora.tp_layer import (Any, __name__, torch) - - -torch_addmm = torch.addmm -torch_add = torch.add -torch_is_compiling = torch.compiler.is_compiling -# torch 2.8 / 2.9 inductor miscompiles addmm(alpha = scaling, beta = 1) traced inside a -# compiled block (e.g. the VLM decoder MLP), forward and backward; fixed in 2.10 (unsloth#3271). -torch_lora_no_compiled_addmm = tuple(int(v) for v in torch.__version__.split(".")[:2]) < (2, 10) -# @torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) -def lora_forward(result, lora_A, lora_B, dropout, x, scaling): - # Use result.dtype (bfloat16 from base layer) since x may have been cast to float32 - # by _cast_input_dtype when autocast is disabled - target_dtype = result.dtype - xA = dropout(x).to(target_dtype) @ lora_A.weight.to(target_dtype).t() - if torch_lora_no_compiled_addmm and torch_is_compiling(): - output = result + scaling * (xA @ lora_B.weight.to(target_dtype).t()) - bias = lora_B.bias - if bias is not None: - output = output + scaling * bias.to(target_dtype) - return output - # output = result + scaling * xA @ lora_B.weight.t() - shape = result.shape - output = torch_addmm( - result.view(-1, shape[-1]), - xA.view(-1, xA.shape[-1]), - lora_B.weight.to(target_dtype).t(), - alpha = scaling, - beta = 1, - ).view(shape) - - bias = lora_B.bias - if bias is not None: - output = torch_add( - output, - bias.to(target_dtype), - alpha = scaling, - ) - return output -pass - -def unsloth_forward(self, x: torch.Tensor, *args: Any, **kwargs: Any): - - adapter_names = kwargs.pop("adapter_names", None) - # If weight is used for matrix multiplication here, the final aggregation operation of the original - # parallel_linear layer will be missing, so we need to directly call its forward function to obtain the - # output of the original parallel_linear layer. - if self.disable_adapters: - if self.merged: - self.unmerge() - result, bias = self.base_layer(x, *args, **kwargs) - elif adapter_names is not None: - raise ValueError(f"{self.__class__.__name__} does not support mixed_batch_forward yet.") - elif self.merged: - result, bias = self.base_layer(x, *args, **kwargs) - else: - result, bias = self.base_layer(x, *args, **kwargs) - torch_result_dtype = result.dtype - for active_adapter in self.active_adapters: - if active_adapter not in self.lora_A.keys(): - continue - lora_A = self.lora_A[active_adapter] - lora_B = self.lora_B[active_adapter] - dropout = self.lora_dropout[active_adapter] - scaling = self.scaling[active_adapter] - if not torch.is_autocast_enabled(): result, x = result.to(lora_A.weight.dtype), x.to(lora_A.weight.dtype) - return lora_forward(result, lora_A, lora_B, dropout, x, scaling).to(torch_result_dtype) - - result = result.to(torch_result_dtype) - return result, bias diff --git a/unsloth_compiled_cache/RMSNorm.py b/unsloth_compiled_cache/RMSNorm.py deleted file mode 100644 index 821cd668480e828a0994d7dff01c51533e2b3bfb..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/RMSNorm.py +++ /dev/null @@ -1,105 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - - -import os -import sys -import torch -import importlib.util -import math -if importlib.util.find_spec("unsloth_studio") is None: - UNSLOTH_STUDIO_ENABLED = False -else: - UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" -pass -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -import math - -UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" -UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" -UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) -UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") -if UNSLOTH_COMPILE_LOCATION not in sys.path: - sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) - -import logging -logger_compiler = logging.getLogger(__name__) -if UNSLOTH_ENABLE_LOGGING: - logger_compiler.setLevel(logging.DEBUG) - -global INFERENCE_RUNS -INFERENCE_RUNS = 0 - -try: - import torch._dynamo.eval_frame as torch_dynamo_eval_frame - torch_dynamo_eval_frame._stance.stance - torch_compiler_set_stance = torch.compiler.set_stance -except: - torch_dynamo_eval_frame = None - torch_compiler_set_stance = None -pass - -# The stance is process wide, so whether Unsloth owns it lives in unsloth_zoo, shared by -# every generated module: inference in one model's module, training in another's. -from unsloth_zoo.temporary_patches.utils import UNSLOTH_EAGER_STANCE_OWNED, unsloth_claim_eager_stance, unsloth_owns_stance, unsloth_release_stance - -def unsloth_training_stance(): - # eager_on_recompile is global: left on by generate() or eval forwards, every recompile - # in a later training step (new sequence length, model.train() without for_training) - # would silently run the eager body. Training forwards restore the default stance, - # unless the user chose eager_on_recompile themselves. Every set_stance call installs a - # fresh stance object, so ours is still active only if the current one `is` it; a user - # set_stance since then (even to eager_on_recompile again) is theirs and stays. - global INFERENCE_RUNS - INFERENCE_RUNS = 0 - if len(UNSLOTH_EAGER_STANCE_OWNED) != 0 and torch_dynamo_eval_frame is not None: - current = torch_dynamo_eval_frame._stance - # Ownership is only dropped once ours is reset: a temporary stance scope's exit restores - # one of our exact objects, which a later training forward must still recognise. - if current.stance == "eager_on_recompile" and unsloth_owns_stance(current): - # Only the reset object: an outer scope may still restore another one of ours. - # Put back what the caller had before Unsloth switched (force_eager, a backend, ...). - previous = unsloth_release_stance(current) - if previous is not None and hasattr(torch_dynamo_eval_frame, "_set_stance"): - torch_dynamo_eval_frame._set_stance(previous) - else: - torch_compiler_set_stance(stance = "default", skip_guard_eval_unsafe = False) -pass - -from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT -from unsloth_zoo.temporary_patches.utils import UNSLOTH_DECODE_COMPILE - -torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 32, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': False, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from transformers.models.siglip.modeling_siglip import (torch) - -def forward(self, x: torch.Tensor) -> torch.Tensor: - """ - Runs the forward pass. - """ - return F.rms_norm(x, self.normalized_shape, self.weight, self.eps).to(x.dtype) diff --git a/unsloth_compiled_cache/TeLinear_peft_forward.py b/unsloth_compiled_cache/TeLinear_peft_forward.py deleted file mode 100644 index 52f5229c9322778211b2c0c9f4931af2d5373d8c..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/TeLinear_peft_forward.py +++ /dev/null @@ -1,104 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - - -torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 32, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': False, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from unsloth_zoo.temporary_patches.common import torch_compile -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from peft.tuners.lora.te import (Any, __name__, torch) - - -torch_addmm = torch.addmm -torch_add = torch.add -torch_is_compiling = torch.compiler.is_compiling -# torch 2.8 / 2.9 inductor miscompiles addmm(alpha = scaling, beta = 1) traced inside a -# compiled block (e.g. the VLM decoder MLP), forward and backward; fixed in 2.10 (unsloth#3271). -torch_lora_no_compiled_addmm = tuple(int(v) for v in torch.__version__.split(".")[:2]) < (2, 10) -# @torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options) -def lora_forward(result, lora_A, lora_B, dropout, x, scaling): - # Use result.dtype (bfloat16 from base layer) since x may have been cast to float32 - # by _cast_input_dtype when autocast is disabled - target_dtype = result.dtype - xA = dropout(x).to(target_dtype) @ lora_A.weight.to(target_dtype).t() - if torch_lora_no_compiled_addmm and torch_is_compiling(): - output = result + scaling * (xA @ lora_B.weight.to(target_dtype).t()) - bias = lora_B.bias - if bias is not None: - output = output + scaling * bias.to(target_dtype) - return output - # output = result + scaling * xA @ lora_B.weight.t() - shape = result.shape - output = torch_addmm( - result.view(-1, shape[-1]), - xA.view(-1, xA.shape[-1]), - lora_B.weight.to(target_dtype).t(), - alpha = scaling, - beta = 1, - ).view(shape) - - bias = lora_B.bias - if bias is not None: - output = torch_add( - output, - bias.to(target_dtype), - alpha = scaling, - ) - return output -pass - -def unsloth_forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: - - adapter_names = kwargs.pop("adapter_names", None) - - if self.disable_adapters: - if not torch.is_autocast_enabled() and hasattr(self.base_layer, 'weight') and self.base_layer.weight is not None and not hasattr(self.base_layer.weight, 'quant_state') and self.base_layer.weight.is_floating_point() and self.base_layer.weight.element_size() > 1 and x.dtype != self.base_layer.weight.dtype: - x = x.to(self.base_layer.weight.dtype) - result = self.base_layer(x, *args, **kwargs) - elif adapter_names is not None: - raise ValueError(f"{self.__class__.__name__} does not support mixed_batch_forward yet.") - else: - if not torch.is_autocast_enabled() and hasattr(self.base_layer, 'weight') and self.base_layer.weight is not None and not hasattr(self.base_layer.weight, 'quant_state') and self.base_layer.weight.is_floating_point() and self.base_layer.weight.element_size() > 1 and x.dtype != self.base_layer.weight.dtype: - x = x.to(self.base_layer.weight.dtype) - result = self.base_layer(x, *args, **kwargs) - torch_result_dtype = result.dtype - - lora_A_keys = self.lora_A.keys() - for active_adapter in self.active_adapters: - if active_adapter not in lora_A_keys: - continue - - lora_A = self.lora_A[active_adapter] - lora_B = self.lora_B[active_adapter] - dropout = self.lora_dropout[active_adapter] - scaling = self.scaling[active_adapter] - if not torch.is_autocast_enabled(): result, x = result.to(lora_A.weight.dtype), x.to(lora_A.weight.dtype) - return lora_forward(result, lora_A, lora_B, dropout, x, scaling).to(torch_result_dtype) - - result = result.to(torch_result_dtype) - - return result diff --git a/unsloth_compiled_cache/UnslothCPOTrainer.py b/unsloth_compiled_cache/UnslothCPOTrainer.py deleted file mode 100644 index a916172506b1797462dd7b944742a59295ef24ff..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/UnslothCPOTrainer.py +++ /dev/null @@ -1,2655 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from unsloth_zoo.temporary_patches.common import torch_compile -from unsloth_zoo.temporary_patches.common import _maybe_compile -import functools -import transformers -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from trl.experimental.cpo.cpo_trainer import (Any, AutoModelForCausalLM, AutoTokenizer, BaseImageProcessor, CPOConfig, CPOTrainer, Callable, DPODataCollatorWithPadding, DataCollator, DataLoader, Dataset, EvalLoopOutput, F, FeatureExtractionMixin, Literal, Path, PeftConfig, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, TrainerCallback, Version, _BaseTrainer, add_bos_token_if_needed, add_eos_token_if_needed, autocast, defaultdict, disable_dropout_in_model, get_config_model_id, get_peft_model, global_then_local_main_first, inspect, is_comet_available, is_peft_available, is_peft_model, is_torch_fx_proxy, is_wandb_available, log_table_to_comet_experiment, logger, maybe_apply_chat_template, maybe_extract_prompt, nn, np, nullcontext, pad_to_length, pd, peft_module_casting_to_bf16, prepare_model_for_kbit_training, random, selective_log_softmax, set_seed, textwrap, torch, transformers, AutoModelForCausalLM, AutoTokenizer, BaseImageProcessor, CPOConfig, CPOTrainer, Callable, DPODataCollatorWithPadding, DataCollator, Dataset, EvalLoopOutput, F, FeatureExtractionMixin, PeftConfig, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, TrainerCallback, Version, autocast, defaultdict, disable_dropout_in_model, get_config_model_id, get_peft_model, global_then_local_main_first, inspect, is_comet_available, is_peft_available, is_peft_model, is_wandb_available, logger, maybe_apply_chat_template, maybe_extract_prompt, nn, np, pd, peft_module_casting_to_bf16, prepare_model_for_kbit_training, random, set_seed, torch, transformers, F, PreTrainedModel, is_peft_available, logger, torch, F, np, F, nn, np, torch, F, PreTrainedModel, add_bos_token_if_needed, add_eos_token_if_needed, logger, nn, np, torch) - - -import os -import math -import logging -from typing import * -from dataclasses import dataclass, field -from packaging.version import Version -import torch -import numpy as np -from contextlib import nullcontext -from torch.nn import functional as F -import inspect -from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling -from transformers.training_args import ParallelMode -from unsloth_zoo.device_type import DEVICE_TYPE, DEVICE_TYPE_TORCH, device_synchronize - -# Wrap trainer with padding to right and enable training mode -import functools -from types import MethodType -try: - from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers -except: - def reset_unsloth_gradient_checkpointing_buffers(): pass -# Canonical reset lives in unsloth.models._utils so the SFT auto-packing wrapper and the plain -# Trainer loop can import the same helper; fall back to a no-op only if it can't be imported. -try: - from unsloth.models._utils import _unsloth_reset_stray_compile_cache -except Exception: - def _unsloth_reset_stray_compile_cache(self): pass -try: - from unsloth.models._utils import _unsloth_dataset_column_names -except Exception: - def _unsloth_dataset_column_names(dataset): return dataset.column_names -# Drops/renames config arguments the installed TRL no longer accepts, so a -# script pinned to an older TRL keeps working after an upgrade. Falls back to -# the historical raw passthrough so this can never break trainer construction. -try: - from unsloth.models.rl_config_compat import filter_config_init_kwargs as _unsloth_filter_config_init_kwargs - # A cache file generated here can be imported by an older Unsloth whose filter - # predates `mirrored_from`, so drop the argument rather than raise TypeError. - if "mirrored_from" not in inspect.signature(_unsloth_filter_config_init_kwargs).parameters: - _unsloth_filter_config_init_kwargs_old = _unsloth_filter_config_init_kwargs - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): - return _unsloth_filter_config_init_kwargs_old(config_class, kwargs) -except Exception: - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): return kwargs -def prepare_for_training_mode(f): - @functools.wraps(f) - def wrapper(self, *args, **kwargs): - # Drop any torch.compile graph cache poisoned by a stray pre-train forward. - try: - _unsloth_reset_stray_compile_cache(self) - except Exception: - pass - # Finish the previous W&B run if this is a subsequent train() call. - # We do this at the START of train() (not the end) so that - # evaluate() / log() still work after train() completes. - # HF's WandbCallback.setup() will call wandb.init() for the new run. - # See: https://github.com/unslothai/unsloth/issues/3954 - if getattr(self, '_unsloth_training_completed', False): - try: - import wandb - if wandb.run is not None: - wandb.finish() - # Reset HF's WandbCallback so it calls wandb.init() for the new run - for cb in self.callback_handler.callbacks: - if type(cb).__name__ == 'WandbCallback': - cb._initialized = False - break - except: - pass - # Enable training mode - _was_training = None - # Restore the GC mode the model was configured with at setup; fall back to - # the training args only when it wasn't recorded (issue #4735). Use hasattr, - # not a None sentinel, so a deliberately-recorded None is restored verbatim. - _model = getattr(self, 'model', None) - if hasattr(_model, '_unsloth_gradient_checkpointing'): - use_gc = _model._unsloth_gradient_checkpointing - else: - use_gc = getattr(self.args, 'gradient_checkpointing', True) - if hasattr(self, 'model') and hasattr(self.model, "training"): - _was_training = self.model.training - if hasattr(self, 'model') and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - output = f(self, *args, **kwargs) - # Restore previous mode when possible - if hasattr(self, 'model') and hasattr(self.model, "for_inference"): - if _was_training is False: - self.model.for_inference() - elif _was_training is True and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - # Reset gradient checkpointing buffers to free memory while staying ready for next run - try: - reset_unsloth_gradient_checkpointing_buffers() - except: - pass - # Mark that training completed so the next train() call can - # finish this W&B run before starting a new one - self._unsloth_training_completed = True - return output - return wrapper -pass - -torch_compile_options = { - "epilogue_fusion" : True, - "max_autotune" : False, - "shape_padding" : True, - "trace.enabled" : False, - "triton.cudagraphs" : False, -} - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_hidden_states_selective_log_softmax( - hidden_states: torch.Tensor, - lm_head: torch.Tensor, - index: torch.Tensor, - chunks: int = 4, - logit_scale_multiply: float = 0.0, - logit_scale_divide: float = 0.0, - logit_softcapping: float = 0.0, - temperature: float = 1.0, - # Rows per chunk cap. Read HERE, in the default, not in the body: the body - # is traced with fullgraph = True, and `os.environ` is an unsupported op - # there on torch 2.4 -- Dynamo raises - # torch._dynamo.exc.Unsupported: const method call bytes.decode - # from os._Environ.__getitem__, which the eager fallback does not catch - # (it only catches recompile-limit and disabled-hook breaks), so the very - # first call would die even with the variable unset. A default is evaluated - # once when the def runs, which is import time, outside any traced region. - # It is also a plain int argument, so Dynamo guards on it instead of - # constant-folding an unguarded read (2.7+ never notice a later change). - # A non-numeric value is ignored rather than raised on, so a typo cannot - # break the import. 0 keeps the previous chunk boundaries exactly. - max_rows_per_chunk: int = ( - int(os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip()) - if os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip().isdigit() - else 0 - ), - lm_head_bias: Optional[torch.Tensor] = None, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, - # Read in the default for the same Dynamo reason as max_rows_per_chunk. Pairs with vLLM's - # head_dtype=float32 that load_vllm requests under the same flag. - fp32_lm_head: bool = ( - os.environ.get("UNSLOTH_VLLM_FP32_LM_HEAD", "0").strip().lower() in ("1", "true", "on", "yes") - ), - # Older Unsloth scores old / ref log-probs here and the policy with the fused kernel: one scorer keeps ratio at 1. - fused_linear: bool = ( - os.environ.get("UNSLOTH_FUSED_LINEAR", "1") != "0" - and os.environ.get("UNSLOTH_GRPO_CHUNKED_LOGPROBS", "1") != "0" - ), -) -> torch.Tensor: - # All Unsloth Zoo code licensed under AGPL3 - from unsloth_zoo.fused_losses import chunked_logprob as _chunked_logprob - if (fused_linear and _chunked_logprob._FUSED_LINEAR and not fp32_lm_head - and support_top_k == 0 and support_top_p == 1.0 and support_min_p == 0.0 - and support_excluded_ids is None): - return _chunked_logprob.unsloth_chunked_logprobs( - hidden_states, lm_head, index, - temperature = temperature, - logit_scale_multiply = logit_scale_multiply, - logit_scale_divide = logit_scale_divide, - logit_softcapping = logit_softcapping, - lm_head_bias = lm_head_bias, - max_rows_per_chunk = max_rows_per_chunk if max_rows_per_chunk > 0 else None, - ).reshape((hidden_states.shape[0], hidden_states.shape[1])) - # Reshape on this tensor's own last dim: a no-op, so a wrong-width caller - # cannot have its row count silently rewritten and instead fails at the - # matmul below, which prints both operands. Do not swap in a bare - # torch._check: it reports only "Expected cond to be True", naming neither - # operand, and Dynamo rejects a message-carrying one. Callers dispatch on - # the width first -- see `compute_logprobs_chunk`, the packed path and - # `_pg_grad_forward`. - flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) - flat_index = index.reshape(-1) - - # Each chunk materialises rows x vocab logits and then a float32 copy of - # them, all on the device holding the output head. With a large vocabulary - # and a fixed chunk count that grows with the batch, so the peak scales with - # the batch rather than staying bounded. max_rows_per_chunk caps the rows - # per chunk instead, which is pure loop splitting: more, smaller chunks, - # same concatenated result. 0 (the default) keeps the previous chunk - # boundaries exactly. - if max_rows_per_chunk > 0: - n_rows = flat_hidden_states.shape[0] - chunks = max(chunks, -(-n_rows // max_rows_per_chunk)) - chunks = min(chunks, max(n_rows, 1)) - - # A one-row chunk's head gradient is a K=1 matmul Inductor mis-lowers under dynamic shapes (illegal - # memory access or NaN head gradients, torch 2.13), so keep chunks at 2+ rows when the head trains. - # Exactly `chunks` row ranges, so the unrolled loop length never depends on T: torch.chunk - # returns ceil(T / ceil(T / chunks)) pieces and recompiled per new count; tensor_split guards - # on T % chunks. At most rows // 2 chunks when the head trains keeps every piece at 2+ rows. - lm_head_grad = torch.is_grad_enabled() and lm_head.requires_grad - if lm_head_grad and flat_hidden_states.shape[0] < 2 * chunks: - chunks = max(flat_hidden_states.shape[0] // 2, 1) - - n_rows = flat_hidden_states.shape[0] - chunks = max(1, min(chunks, n_rows)) - bounds = [(i * n_rows) // chunks for i in range(chunks + 1)] - chunked_hidden_states = [flat_hidden_states[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - - # Cast once, outside the loop, so every chunk shares one fp32 head copy (also for autograd). - if fp32_lm_head: - lm_head = lm_head.to(torch.float32) - if lm_head_bias is not None: - lm_head_bias = lm_head_bias.to(torch.float32) - - all_per_token_logps = [] - - for chunk_hidden_states, chunk_index in zip(chunked_hidden_states, chunked_index): - # When the model is dispatched over several devices, the output head can - # sit on a different one from the hidden states, because accelerate - # places the tail of the model on the last device it fills. Co-locate on - # the head's device before the matmul, otherwise this raises - # Unhandled FakeTensor Device Propagation for aten.mm.default, - # found two different devices cuda:0, cuda:1 - # On a single device every .to() here is a no-op and the result is - # bit-identical to before. - chunk_hidden_states = chunk_hidden_states.to(device = lm_head.device, dtype = lm_head.dtype) - chunk_index = chunk_index.to(lm_head.device) - if fp32_lm_head: - # The GRPO loss runs this under bf16 / fp16 autocast, which would downcast the matmul. - with torch.autocast(device_type = chunk_hidden_states.device.type, enabled = False): - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - elif lm_head_bias is None: - chunk_logits = chunk_hidden_states @ lm_head.t() - else: - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - - if logit_scale_multiply != 0.0: - chunk_logits = chunk_logits * logit_scale_multiply - if logit_scale_divide != 0.0: - chunk_logits = chunk_logits / logit_scale_divide - if logit_softcapping != 0.0: - chunk_logits = logit_softcapping * torch.tanh(chunk_logits / logit_softcapping) - - chunk_logits = chunk_logits.to(torch.float32) - - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - per_token_logps = _truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - ) - else: - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - # Return to the caller's device so the concatenation below and every - # downstream consumer see the device they started on. - all_per_token_logps.append(per_token_logps.to(hidden_states.device)) - - all_per_token_logps = torch.concat(all_per_token_logps) - - all_per_token_logps = all_per_token_logps.reshape((hidden_states.shape[0], hidden_states.shape[1])) - return all_per_token_logps - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_selective_log_softmax( - logits, - index, - temperature: float = 1.0, - chunks: int = 4, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, -): - # Fixed `chunks` row ranges; see chunked_hidden_states_selective_log_softmax. - flat_logits = logits.reshape(-1, logits.shape[-1]) - flat_index = index.reshape(-1) - # Fewer rows than chunks: one row per range, as torch.chunk did, never empty ranges. - chunks = max(1, min(chunks, flat_logits.shape[0])) - bounds = [(i * flat_logits.shape[0]) // chunks for i in range(chunks + 1)] - chunked_logits = [flat_logits[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - all_per_token_logps = [] - # Per-chunk selective_log_softmax. - for chunk_logits, chunk_index in zip(chunked_logits, chunked_index): - chunk_logits = chunk_logits.to(torch.float32) - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - all_per_token_logps.append(_truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - )) - continue - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - all_per_token_logps.append(per_token_logps) - pass - all_per_token_logps = torch.concat(all_per_token_logps) - all_per_token_logps = all_per_token_logps.reshape((logits.shape[0], logits.shape[1])) - return all_per_token_logps - -def calculate_pad_tokens_in_prompt( - input_ids: torch.Tensor, - logits_to_keep: int, - pad_token_id: int -) -> torch.Tensor: - """Count left-padded tokens per sequence, e.g. [pad, pad, pad, cat] -> 3.""" - if logits_to_keep >= input_ids.shape[1]: - raise ValueError("logits_to_keep must be smaller than the sequence length.") - - prompt_section = input_ids[:, :-logits_to_keep] - - padding_mask = (prompt_section == pad_token_id) - - pad_token_counts = padding_mask.sum(dim=1) - - return pad_token_counts - -def create_completion_attention_mask( - completion_input_ids: torch.Tensor, - left_pad_tokens_per_prompt: torch.Tensor, - max_left_pad: int, - pad_token_id: int -) -> torch.Tensor: - """Build a completion mask that zeros leading prompt and trailing pad tokens. - - For [p,p,p,c,c,c,pad,pad,pad] (p=sliced prompt, c=completion, pad=padding) - this returns [0,0,0,1,1,1,0,0,0]. - """ - batch_size, completion_len = completion_input_ids.shape - device = completion_input_ids.device - - num_tokens_to_mask = max_left_pad - left_pad_tokens_per_prompt - - indices = torch.arange(completion_len, device=device).unsqueeze(0) - shift_mask = indices >= num_tokens_to_mask.unsqueeze(1) - - non_padding_mask = (completion_input_ids != pad_token_id) - - final_mask = shift_mask & non_padding_mask - - return final_mask - -def left_pack_padding(tensor: torch.Tensor, pad_id: int) -> torch.Tensor: - """Move all padding tokens in each sequence to the right.""" - mask = (tensor != pad_id) - # stable=True since the binary mask is unordered. - sorted_indices = torch.argsort(mask, dim=1, descending=True, stable=True) - packed_tensor = torch.gather(tensor, 1, sorted_indices) - return packed_tensor - -def align_logprobs_with_mask( - logprob_tensor: torch.Tensor, - attention_mask: torch.Tensor, - pad_value: float = 0.0 -) -> torch.Tensor: - """Align a log probability tensor with a given attention mask.""" - - device = logprob_tensor.device - batch_size, logprob_seq_len = logprob_tensor.shape - mask_seq_len = attention_mask.shape[1] - - padded_logprobs = torch.full( - attention_mask.shape, - fill_value=pad_value, - dtype=logprob_tensor.dtype, - device=device - ) - - left_pad_counts = torch.argmax(attention_mask, dim=1) - - cols = torch.arange(logprob_seq_len, device=device) - dest_indices = left_pad_counts.unsqueeze(1) + cols - - # Destination row indices, shape [batch_size, logprob_seq_len]. - row_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand_as(dest_indices) - - # Keep only in-bounds destinations, then scatter via advanced indexing. - valid_mask = dest_indices < mask_seq_len - valid_rows = row_indices[valid_mask] - valid_cols = dest_indices[valid_mask] - valid_vals = logprob_tensor[valid_mask] - padded_logprobs[valid_rows, valid_cols] = valid_vals - - return padded_logprobs - -def align_completion_tool_mask( - tool_mask: torch.Tensor, - completion_mask: torch.Tensor, -) -> torch.Tensor: - """Align a raw completion-length tool/env mask with Unsloth's repacked loss mask.""" - if tool_mask is None: - return completion_mask - if tool_mask.shape[0] != completion_mask.shape[0]: - raise ValueError("tool_mask batch size must match completion_mask batch size.") - - tool_mask = tool_mask.to(device=completion_mask.device) - if tool_mask.shape == completion_mask.shape: - aligned_tool_mask = tool_mask - else: - aligned_tool_mask = align_logprobs_with_mask( - tool_mask, - completion_mask, - pad_value=0, - ) - return completion_mask * aligned_tool_mask.to(dtype=completion_mask.dtype) - -def autotune_batch_and_chunks( - total_input_rows, - seq_len, - hidden_size, - vocab_size, - dtype_bytes=16, - multiplier=None -): - if multiplier is None: - final_m = max(4, seq_len // 4096) - else: - final_m = multiplier - - if torch.cuda.is_available(): - free_bytes, _ = torch.cuda.mem_get_info() - limit_gb = (free_bytes / (1024**3))*.80 - elif hasattr(torch, "xpu") and torch.xpu.is_available(): - # XPU: estimate free memory as total - reserved. - total_mem = torch.xpu.get_device_properties(0).total_memory - reserved_mem = torch.xpu.memory_reserved() - free_bytes = total_mem - reserved_mem - limit_gb = (free_bytes / (1024**3)) * 0.80 - else: - # Fallback: assume 8GB available. - limit_gb = 8.0 - - bytes_to_gb = 1024**3 - - b_vals = torch.arange(total_input_rows, 0, -1, device='cpu', dtype=torch.float32) - - hidden_gb = (b_vals * seq_len * hidden_size * dtype_bytes) / bytes_to_gb - - base_logits = ((b_vals/total_input_rows) * b_vals * seq_len * vocab_size * dtype_bytes) / bytes_to_gb - logits_gb = base_logits / final_m - - total_mem_gb = hidden_gb + logits_gb - - valid_mask = total_mem_gb <= limit_gb - valid_indices = torch.nonzero(valid_mask, as_tuple=False) - - if valid_indices.shape[0] == 0: - #This means your GPU will OOM - # Capped at the row count: unsloth's no-grad pass divides rows by this without max(1, ...). - return max(1, min(4, total_input_rows)), final_m - - best_idx = valid_indices[0].item() - final_b = int(b_vals[best_idx].item()) - - return final_b, final_m - -def sanitize_logprob(logprob): - """Local port of trl.scripts.vllm_serve.sanitize_logprob. - Filters NaN logprobs from vLLM outputs.""" - value = logprob.logprob - if math.isnan(value): - logging.getLogger(__name__).warning( - f"Generated NaN logprob, token logprob '{logprob}' will be ignored" - ) - return None - return value -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_text_configs(config): - """``config`` and its text sub-config, the two places a transform can live: Gemma-4-style configs keep it on ``config.text_config``, T5Gemma-style ones only reach it via ``config.get_text_config()``.""" - if config is None: - return [] - holders = [config] - 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: - holders.append(text_cfg) - return holders - -def _unsloth_get_final_logit_softcapping(model): - """The soft cap the loss applies, under any of its three spellings: ``final_logit_softcapping`` (Gemma), ``logits_soft_cap`` (RecurrentGemma), ``output_logit_soft_cap`` (xLSTM). Returns 0 if unset.""" - config = _unsloth_get_model_config(model) - if config is None: - return 0 - for holder in _unsloth_text_configs(config): - for name in ("final_logit_softcapping", "logits_soft_cap", "output_logit_soft_cap"): - softcap = getattr(holder, name, None) - if softcap: - return softcap - return 0 - -def _unsloth_resolve_logit_scales(model_config): - """``(multiply, divide)`` for the logits, read with the same field table and per-family overrides ``detect_logit_transforms`` uses. - - Only reached when the installed unsloth_zoo predates that helper. Must stay in step with ``resolve_logit_transforms`` in unsloth/models/llama.py: if the forward applies a transform GRPO does not, the policy log-probabilities come from different logits than the ones generated, which silently shifts every importance ratio. - """ - # Granite divides, HyperCLOVA X multiplies, MiniCPM3 scales hidden states (not a logit transform). - overrides = { - ("logits_scaling", "hyperclovax"): "multiply", - ("logits_scaling", "minicpm3"): None, - } - found = {"multiply": 0, "divide": 0} - for holder in _unsloth_text_configs(model_config): - model_type = getattr(holder, "model_type", "") or "" - for bucket, names in ( - ("multiply", ("logit_scale", "lm_head_multiplier", "output_multiplier")), - ("divide", ("logits_scaling",)), - ): - for name in names: - target = overrides.get((name, model_type), bucket) - if target is None or found[target]: - continue - value = getattr(holder, name, 0) or 0 - if value: - found[target] = value - break - return found["multiply"], found["divide"] - -def _unsloth_grpo_returns_hidden_states(model, tensor, lm_head): - """Does ``tensor`` (a forward's ``.logits``) carry hidden states or real logits? - - ``_get_per_token_logps_and_entropies`` sets ``UNSLOTH_RETURN_HIDDEN_STATES=1``, but only a forward that honours the name hands hidden states back as ``.logits``; any other forward returns a real ``[.., vocab]`` tensor that must not reach the ``lm_head`` matmul. - - The primary test is an explicit signal that the forward honours the flag, both set outside this file: ``__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__`` on the generated class, written by ``unsloth_zoo.compiler.create_standalone_class`` exactly when ``apply_fused_lm_head`` gave that forward its own branch; and ``_unsloth_grpo_hidden_states_forward_wrapped``, set by ``_install_grpo_hidden_states_forward_wrapper`` in ``unsloth/models/rl.py`` for models the compiler did not rewrite. That wrapper degrades to real logits when the model cannot produce hidden states and records whether it did so in ``_unsloth_grpo_hidden_states_degraded`` before returning, so reading the pair after a forward describes the call that just finished; degradation is per call, not per model. - - The width comparison stays as the fallback, for an ``unsloth_zoo`` old enough that it never writes the marker. It is decisive whenever ``vocab_size != hidden_size``, and the signal may only overrule it when it is not, which is the one case the shape cannot answer. - """ - if tensor.shape[-1] != lm_head.shape[1]: - return False # vocab-wide: real logits, whatever any signal claims - if lm_head.shape[0] != lm_head.shape[1]: - return True # hidden-wide and vocab_size != hidden_size: hidden states - return _unsloth_grpo_hidden_states_signal(model) is not False - -def _unsloth_grpo_hidden_states_signal(model): - """``True``/``False`` if the forward honours ``UNSLOTH_RETURN_HIDDEN_STATES``, ``None`` when neither marker is present. See ``_unsloth_grpo_returns_hidden_states`` for where each marker is set. Walks the wrapper chain because the markers are set on whichever object the trainer saw, which may be the DDP module or the PEFT base model rather than the object handed to the logprob loop.""" - candidates = [] - pending = [model] - while pending and len(candidates) < 8: - candidate = pending.pop(0) - if candidate is None or any(candidate is seen for seen in candidates): - continue - candidates.append(candidate) - get_base_model = getattr(candidate, "get_base_model", None) - if callable(get_base_model): - try: - pending.append(get_base_model()) - except Exception: - pass - for _attr in ("module", "base_model", "model"): - child = getattr(candidate, _attr, None) - if child is not None and hasattr(child, "forward"): - pending.append(child) - for candidate in candidates: - if getattr(candidate, "__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__", False): - return True - if getattr(type(candidate), "__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__", False): - return True - if any( - getattr(candidate, "_unsloth_grpo_hidden_states_forward_wrapped", False) - for candidate in candidates - ): - # The wrapper honours the flag unless it recorded that this call could not. - if any( - hasattr(candidate, "_unsloth_grpo_hidden_states_degraded") for candidate in candidates - ): - return not any( - getattr(candidate, "_unsloth_grpo_hidden_states_degraded", False) - for candidate in candidates - ) - # An unsloth/models/rl.py predating the per-call attribute set only the warn-once flag; it is the best signal such a wrapper offers. - return not any( - getattr(candidate, "_unsloth_grpo_hidden_states_warning_issued", False) - for candidate in candidates - ) - return None - -def _unsloth_pref_returns_hidden_states(model, tensor, lm_head): - """``_unsloth_grpo_returns_hidden_states`` for the preference routes, which need a zoo that writes the hidden-state - markers: a square head with no signal is a forward that ignored the flag, so its output is real logits.""" - if tensor.shape[-1] != lm_head.shape[1]: - return False - if lm_head.shape[0] != lm_head.shape[1]: - return True - return _unsloth_grpo_hidden_states_signal(model) is True - -def _unsloth_gkd_dense_head(model): - """The output head when it is a plain dense ``[vocab, hidden]`` projection the chunked loss can read directly. - Rejects bnb heads (``nn.Linear`` subclasses with packed weights), PEFT heads and ZeRO-3 partitioned weights. - """ - get_output_embeddings = getattr(model, "get_output_embeddings", None) - if not callable(get_output_embeddings): - return None - try: - head = get_output_embeddings() - except Exception: - return None - if type(head) is not torch.nn.Linear: - return None - # An lm_head that is not the output embeddings runs before them (ModernBERT decoder, RoBERTa-style heads). - lm_head = getattr(model, "lm_head", None) - if isinstance(lm_head, torch.nn.Module) and lm_head is not head: - return None - weight = getattr(head, "weight", None) - if not isinstance(weight, torch.Tensor) or weight.dim() != 2 or weight.numel() == 0: - return None - if weight.dtype not in (torch.float16, torch.bfloat16, torch.float32, torch.float64): - return None - if type(weight) not in (torch.Tensor, torch.nn.Parameter): - return None - bias = getattr(head, "bias", None) - if bias is not None and (bias.dim() != 1 or bias.shape[0] != weight.shape[0]): - return None - if weight.is_meta or (bias is not None and bias.is_meta): - return None - return head - -def _unsloth_stock_logps_hook( - self, - method, - hooks = ("get_batch_logps",), -): - """False when a subclass overrides one of ``hooks`` below the class whose ``method`` is the rewritten one: TRL's - body would call the override (or it reads the full logits the chunked path replaces with stand-ins).""" - owner = None - for klass in type(self).__mro__: - # The rewritten body calls a chunked or fused-pass helper; a subclass method wrapping it via super() does not. - names = getattr(getattr(vars(klass).get(method), "__code__", None), "co_names", ()) - if any( - name.startswith("_unsloth_") and ("chunked" in name or "fused_pass" in name) - for name in names - ): - owner = klass - break - if owner is None: - return True - for name in hooks: - hook = next((klass for klass in type(self).__mro__ if name in vars(klass)), None) - if hook is not None and hook not in owner.__mro__: - return False - return True - -def _unsloth_pref_chunked_head(self, model): - """The dense output head when this concatenated_forward can score from hidden states, else ``None``.""" - if unsloth_chunked_logprobs is None or os.environ.get("UNSLOTH_FUSED_LINEAR", "1") == "0": - return None - if os.environ.get("UNSLOTH_PREFERENCE_CHUNKED", "1").lower() in ("0", "false", "no", "off"): - return None - if getattr(self, "is_encoder_decoder", False) or getattr(self, "aux_loss_enabled", False): - return None - # Listing the method itself rejects a subclass wrapping it: it would see stand-in logits. - if not _unsloth_stock_logps_hook( - self, - "concatenated_forward", - ("get_batch_logps", "get_batch_loss_metrics", "concatenated_forward"), - ): - return None - # TRL <= 0.25 scores labels against label_pad_token_id but nn.CrossEntropyLoss always ignores -100. - if getattr(self, "label_pad_token_id", -100) != -100: - return None - if getattr(self, "is_fsdp_enabled", False) or getattr(self, "is_deepspeed_enabled", False): - return None - try: - unwrapped = self.accelerator.unwrap_model(model) - except Exception: - return None - head = _unsloth_gkd_dense_head(unwrapped) - if head is None: - return None - # The kernel reads weight / bias directly: a hooked head (accelerate, forward / backward hooks) is not - # F.linear(h, W, b), and its backward hooks would never see the gradient. - if ( - hasattr(head, "_hf_hook") - or "forward" in vars(head) - or head._forward_hooks - or head._forward_pre_hooks - or head._backward_hooks - or getattr(head, "_backward_pre_hooks", None) - ): - return None - # MiniCPM3 divides hidden states by logits_scaling before the head. - for config in _unsloth_text_configs(_unsloth_get_model_config(unwrapped)): - if getattr(config, "model_type", None) == "minicpm3": - return None - # DDP(find_unused_parameters=True) marks a head skipped in forward as unused, then its grad hook fires twice. - if getattr(model, "find_unused_parameters", False) and any( - p is not None and p.requires_grad for p in (head.weight, head.bias) - ): - return None - return head - -def _unsloth_gkd_logit_transforms(model): - """``(scale, softcap)`` exactly as the model's own forward applies them to its logits.""" - model_config = _unsloth_get_model_config(model) - if detect_logit_transforms is not None: - _transforms = detect_logit_transforms(model_config) - logit_softcapping = _transforms["logit_softcapping"] - logit_scale_multiply = _transforms["logit_scale_multiply"] - logit_scale_divide = _transforms["logit_scale_divide"] - else: - logit_softcapping = _unsloth_get_final_logit_softcapping(model) - logit_scale_multiply, logit_scale_divide = _unsloth_resolve_logit_scales(model_config) - scale = 1.0 - if logit_scale_multiply: - scale = scale * float(logit_scale_multiply) - if logit_scale_divide: - scale = scale / float(logit_scale_divide) - return scale, float(logit_softcapping or 0.0) - -def _unsloth_pref_logit_transforms(model): - """``(multiply, divide, softcap)`` for the chunked kernel, from the GKD resolver (one scale, 0.0 = off).""" - scale, softcap = _unsloth_gkd_logit_transforms(model) - return (0.0 if scale == 1.0 else scale), 0.0, softcap - -def _unsloth_pref_chunked_forward(self, model, head, concatenated_batch, len_chosen, states, kind): - """``(outputs, all_logits)``: ``outputs`` is TRL's return tuple computed from hidden states, or ``None`` with - ``all_logits`` holding real logits for TRL's own body (the forward ignored the flag, or a shape TRL rejects).""" - unwrapped = self.accelerator.unwrap_model(model) - multiply, divide, softcap = _unsloth_pref_logit_transforms(unwrapped) - if not _unsloth_pref_returns_hidden_states(unwrapped, states, head.weight): - return None, states - input_ids = concatenated_batch["concatenated_input_ids"] - labels = concatenated_batch["concatenated_labels"] - if states.dim() != 3 or states.shape[:2] != labels.shape or labels.shape != input_ids.shape: - # TRL raises or reads these itself: hand it the logits it expects. - logits = torch.nn.functional.linear( - states.to(device = head.weight.device, dtype = head.weight.dtype), head.weight, head.bias - ) - if multiply: - logits = logits * multiply - if divide: - logits = logits / divide - if softcap: - logits = softcap * torch.tanh(logits / softcap) - return None, logits - - device = states.device - labels = labels.to(device) - input_ids = input_ids.to(device) - shifted = labels[:, 1:] - label_mask = shifted != -100 - # cat, not boolean slice assignment: CPU Inductor fails to vectorise that (VecMask vs VecMask). - rejected_off = torch.zeros_like(label_mask[len_chosen:]) - if kind == "orpo": - # ORPO's NLL scores every attended token of prompt + chosen, from input_ids. - nll_targets = input_ids[:, 1:] - attention_mask = concatenated_batch["concatenated_attention_mask"].to(device) - nll_mask = torch.cat((attention_mask[:len_chosen, 1:] == 1, rejected_off), dim = 0) - average_log_prob = True - # Checking for a second NLL pass needs a host sync, so compiled graphs always take it (correct either way). - is_compiling = getattr(getattr(torch, "compiler", None), "is_compiling", None) - if is_compiling is not None and is_compiling(): - split_nll = True - else: - split_nll = bool((label_mask & nll_mask & (shifted != nll_targets)).any()) - else: - nll_targets = shifted - if self.cpo_alpha != 0: - nll_mask = torch.cat((label_mask[:len_chosen], rejected_off), dim = 0) - else: - nll_mask = torch.zeros_like(label_mask) - average_log_prob = self.loss_type in ["ipo", "simpo"] - split_nll = False - targets = torch.where(label_mask, shifted, torch.where(nll_mask, nll_targets, 0)) - - # CPO's logits metric averages every position, ORPO's all but the last. - rows = states if kind == "cpo" else states[:, :-1] - kwargs = dict( - logit_scale_multiply = multiply, - logit_scale_divide = divide, - logit_softcapping = softcap, - lm_head_bias = head.bias, - ) - if _unsloth_pref_masked_logprobs is not None and _unsloth_linear_mean_logits is not None: - # Prompt and padding rows (often most of each row) never reach the head; they only feed the mean metric. - per_token = _unsloth_pref_masked_logprobs( - states[:, :-1], - head.weight, - targets, - mask = label_mask if split_nll else label_mask | nll_mask, - **kwargs, - ) - mean_logits = _unsloth_linear_mean_logits( - rows, head.weight, head.bias, multiply, divide, softcap - ) - if split_nll: - nll_per_token = _unsloth_pref_masked_logprobs( - states[:len_chosen, :-1], - head.weight, - nll_targets[:len_chosen], - mask = nll_mask[:len_chosen], - **kwargs, - ) - else: - nll_per_token = per_token[:len_chosen] - else: - if kind == "cpo": - targets = torch.nn.functional.pad(targets, (0, 1)) - per_token, mean_logits = unsloth_chunked_logprobs( - rows, head.weight, targets, return_mean_logits = True, **kwargs - ) - if kind == "cpo": - per_token = per_token[:, :-1] - if split_nll: - nll_per_token = unsloth_chunked_logprobs( - states[:len_chosen, :-1], - head.weight, - torch.where(nll_mask, nll_targets, 0)[:len_chosen], - **kwargs, - ) - else: - nll_per_token = per_token[:len_chosen] - - loss_mask = label_mask.to(per_token.dtype) - all_logps = (per_token * loss_mask).sum(-1) - if average_log_prob: - all_logps = all_logps / loss_mask.sum(-1) - chosen_nll_mask = nll_mask[:len_chosen].to(per_token.dtype) - if kind == "cpo" and self.cpo_alpha == 0: - nll_loss = torch.tensor(0.0).to(self.accelerator.device) - else: - nll_loss = -(nll_per_token * chosen_nll_mask).sum() / chosen_nll_mask.sum() - try: - self._unsloth_pref_chunked_calls = getattr(self, "_unsloth_pref_chunked_calls", 0) + 1 - except Exception: - pass - # Stand-ins: TRL only reads chosen/rejected logits as .detach().mean(). - return ( - all_logps[:len_chosen], - all_logps[len_chosen:], - mean_logits[:len_chosen], - mean_logits[len_chosen:], - nll_loss, - ), states - -try: - from unsloth_zoo.device_map_planner import detect_logit_transforms -except Exception: - detect_logit_transforms = None -try: - from unsloth_zoo.fused_losses.chunked_logprob import unsloth_chunked_logprobs -except Exception: - unsloth_chunked_logprobs = None -try: - from unsloth_zoo.fused_losses.chunked_logprob import unsloth_masked_logprobs as _unsloth_pref_masked_logprobs - from unsloth_zoo.fused_losses.preference_loss import linear_mean_logits as _unsloth_linear_mean_logits -except Exception: - _unsloth_pref_masked_logprobs = _unsloth_linear_mean_logits = None -@dataclass -class UnslothCPOConfig(CPOConfig): - """ - - Configuration class for the [`experimental.cpo.CPOTrainer`]. - - This class includes only the parameters that are specific to CPO training. For a full list of training arguments, - please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may - differ from those in [`~transformers.TrainingArguments`]. - - Using [`~transformers.HfArgumentParser`] we can turn this class into - [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the - command line. - - Parameters: - max_length (`int` or `None`, *optional*, defaults to `1024`): - Maximum length of the sequences (prompt + completion) in the batch. This argument is required if you want - to use the default data collator. - max_completion_length (`int`, *optional*): - Maximum length of the completion. This argument is required if you want to use the default data collator - and your model is an encoder-decoder. - beta (`float`, *optional*, defaults to `0.1`): - Parameter controlling the deviation from the reference model. Higher β means less deviation from the - reference model. For the IPO loss (`loss_type="ipo"`), β is the regularization parameter denoted by τ in - the [paper](https://huggingface.co/papers/2310.12036). - label_smoothing (`float`, *optional*, defaults to `0.0`): - Label smoothing factor. This argument is required if you want to use the default data collator. - loss_type (`str`, *optional*, defaults to `"sigmoid"`): - Type of loss to use. Possible values are: - - - `"sigmoid"`: sigmoid loss from the original [DPO](https://huggingface.co/papers/2305.18290) paper. - - `"hinge"`: hinge loss on the normalized likelihood from the - [SLiC](https://huggingface.co/papers/2305.10425) paper. - - `"ipo"`: IPO loss from the [IPO](https://huggingface.co/papers/2310.12036) paper. - - `"simpo"`: SimPO loss from the [SimPO](https://huggingface.co/papers/2405.14734) paper. - - `"alphapo"`: AlphaPO loss from the [AlphaPO](https://huggingface.co/papers/2501.03884) paper. This - automatically sets `loss_type="simpo"` and `cpo_alpha=0.0`. - - disable_dropout (`bool`, *optional*, defaults to `True`): - Whether to disable dropout in the model. - cpo_alpha (`float`, *optional*, defaults to `1.0`): - Weight of the BC regularizer in CPO training. - simpo_gamma (`float`, *optional*, defaults to `0.5`): - Target reward margin for the SimPO loss, used only when the `loss_type="simpo"`. - alpha (`float`, *optional*, defaults to `0.0`): - Alpha parameter that controls reward function shape across all loss types. When alpha=0 (default), uses - standard log probability rewards. When `alpha != 0`, applies AlphaPO transformation: `r = (1 - p^(-alpha)) - / alpha` from the [AlphaPO paper](https://huggingface.co/papers/2501.03884). This parameter works with all - loss types. - generate_during_eval (`bool`, *optional*, defaults to `False`): - If `True`, generates and logs completions from the model to W&B or Comet during evaluation. - is_encoder_decoder (`bool`, *optional*): - When using the `model_init` argument (callable) to instantiate the model instead of the `model` argument, - you need to specify if the model returned by the callable is an encoder-decoder model. - model_init_kwargs (`dict[str, Any]`, *optional*): - Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model from a - string. The `revision` value is also used when loading the processing class. - trust_remote_code (`bool`, *optional*, defaults to `False`): - Whether to allow loading models that ship custom Python code from the Hub. Forwarded to - [`~transformers.AutoModelForCausalLM.from_pretrained`]. - dataset_num_proc (`int`, *optional*): - Number of processes to use for processing the dataset. - - > [!NOTE] - > These parameters have default values different from [`~transformers.TrainingArguments`]: - > - `logging_steps`: Defaults to `10` instead of `500`. - > - `gradient_checkpointing`: Defaults to `True` instead of `False`. - > - `bf16`: Defaults to `True` if `fp16` is not set, instead of `False`. - > - `learning_rate`: Defaults to `1e-6` instead of `5e-5`. - - """ - _unsloth_accepts_adamw_flash = True - vllm_sampling_params: Optional[Any] = field( - default = None, - metadata = {'help': 'vLLM SamplingParams'}, - ) - unsloth_num_chunks : Optional[int] = field( - default = -1, - metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'}, - ) - unsloth_logit_chunk_multiplier : Optional[int] = field( - default = None, - metadata = {'help': 'Multiplier for chunked logit computations.'}, - ) - unsloth_grpo_mini_batch : Optional[int] = field( - default = None, - metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'}, - ) - max_seq_length : Optional[int] = field( - default = None, - metadata = {'help': 'Maximum sequence length to truncate to.'}, - ) - def __init__( - self, - output_dir = None, - per_device_train_batch_size = 4, - num_train_epochs = 3.0, - max_steps = -1, - learning_rate = 5e-05, - lr_scheduler_type = 'linear', - lr_scheduler_kwargs = None, - warmup_steps = 0.1, - optim = 'adamw_8bit', - optim_args = None, - weight_decay = 0.001, - adam_beta1 = 0.9, - adam_beta2 = 0.999, - adam_epsilon = 1e-08, - optim_target_modules = None, - gradient_accumulation_steps = 2, - average_tokens_across_devices = True, - max_grad_norm = 1.0, - label_smoothing_factor = 0.0, - bf16 = False, - fp16 = False, - bf16_full_eval = False, - fp16_full_eval = False, - tf32 = None, - gradient_checkpointing = True, - gradient_checkpointing_kwargs = None, - torch_compile = False, - torch_compile_backend = None, - torch_compile_mode = None, - use_liger_kernel = False, - liger_kernel_config = None, - use_cache = False, - neftune_noise_alpha = None, - torch_empty_cache_steps = 250, - auto_find_batch_size = False, - logging_strategy = 'steps', - logging_steps = 1, - logging_first_step = False, - log_on_each_node = True, - logging_nan_inf_filter = False, - include_num_input_tokens_seen = False, - log_level = 'passive', - log_level_replica = 'warning', - disable_tqdm = None, - report_to = 'none', - run_name = None, - project = 'huggingface', - trackio_space_id = 'trackio', - eval_strategy = 'no', - eval_steps = None, - eval_delay = 0, - per_device_eval_batch_size = 4, - prediction_loss_only = False, - eval_on_start = False, - eval_do_concat_batches = True, - eval_use_gather_object = False, - eval_accumulation_steps = 2, - batch_eval_metrics = False, - save_only_model = False, - save_strategy = 'steps', - save_steps = 500, - save_on_each_node = False, - save_total_limit = None, - enable_jit_checkpoint = False, - push_to_hub = False, - hub_token = None, - hub_private_repo = None, - hub_model_id = None, - hub_strategy = 'every_save', - hub_always_push = False, - hub_revision = None, - load_best_model_at_end = False, - metric_for_best_model = None, - greater_is_better = None, - ignore_data_skip = False, - restore_callback_states_from_checkpoint = False, - full_determinism = False, - seed = 3407, - data_seed = None, - use_cpu = False, - accelerator_config = None, - parallelism_config = None, - dataloader_drop_last = False, - dataloader_num_workers = 0, - dataloader_pin_memory = True, - dataloader_persistent_workers = False, - dataloader_prefetch_factor = None, - remove_unused_columns = True, - label_names = None, - train_sampling_strategy = 'random', - length_column_name = 'length', - ddp_find_unused_parameters = None, - ddp_bucket_cap_mb = None, - ddp_broadcast_buffers = None, - ddp_backend = None, - ddp_timeout = 1800, - fsdp = None, - fsdp_config = None, - deepspeed = None, - debug = '', - skip_memory_metrics = True, - do_train = False, - do_eval = False, - do_predict = False, - resume_from_checkpoint = None, - warmup_ratio = None, - logging_dir = None, - local_rank = -1, - max_length = 1024, - max_completion_length = None, - beta = 0.1, - label_smoothing = 0.0, - loss_type = 'sigmoid', - disable_dropout = True, - cpo_alpha = 1.0, - simpo_gamma = 0.5, - alpha = 0.0, - generate_during_eval = False, - is_encoder_decoder = None, - model_init_kwargs = None, - trust_remote_code = False, - dataset_num_proc = None, - vllm_sampling_params = None, - unsloth_num_chunks = -1, - unsloth_logit_chunk_multiplier = None, - unsloth_grpo_mini_batch = None, - max_seq_length = None, - **kwargs, - ): - if learning_rate < 1e-7: print(f'Unsloth: Your learning rate of `{learning_rate}` is too small and less than 1e-7! Consider increasing it, otherwise gradient updates will be close to 0!') - if learning_rate > 1: print(f'Unsloth: Your learning rate of `{learning_rate}` is way too larger > 1! Consider decreasing it to 1e-1, otherwise gradient updates will explode!') - _unsloth_fsdp2 = os.environ.get('FSDP_VERSION', '').strip() == '2' - _unsloth_fsdp_cfg = fsdp_config - if fsdp and isinstance(_unsloth_fsdp_cfg, str) and os.path.isfile(_unsloth_fsdp_cfg): - try: - import json as _unsloth_json - with open(_unsloth_fsdp_cfg, encoding = 'utf-8') as _unsloth_f: - _unsloth_fsdp_cfg = _unsloth_json.load(_unsloth_f) - except Exception: - _unsloth_fsdp_cfg = None - if fsdp and _unsloth_fsdp_cfg is None: - _unsloth_fsdp_cfg = {} - if fsdp and isinstance(_unsloth_fsdp_cfg, dict): - _unsloth_fsdp2 = _unsloth_fsdp2 or str(_unsloth_fsdp_cfg.get('fsdp_version', _unsloth_fsdp_cfg.get('version', 1))) == '2' - _unsloth_flash_adamw = str(getattr(optim, 'value', optim)).strip().lower() == 'adamw_flash' - if _unsloth_flash_adamw: optim = 'adamw_torch' - if _unsloth_fsdp2 and str(getattr(optim, 'value', optim)) == 'adamw_8bit': - print("Unsloth: FSDP2 shards parameters as DTensors, which bitsandbytes' 8-bit optimizers cannot step. Switching optim = 'adamw_8bit' to 'adamw_torch_fused'.") - optim = 'adamw_torch_fused' - if num_train_epochs is None: - num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override - if output_dir is None and save_strategy == 'steps' and save_steps == 500: - output_dir = 'unsloth_training_checkpoints' - save_strategy = 'no' - try: - from unsloth_zoo.dataset_num_proc import get_dataset_num_proc as _unsloth_get_dataset_num_proc - except Exception: - try: - from unsloth.dataset_num_proc import get_dataset_num_proc as _unsloth_get_dataset_num_proc - except Exception: - _unsloth_get_dataset_num_proc = None - if _unsloth_get_dataset_num_proc is not None: - dataset_num_proc = _unsloth_get_dataset_num_proc(dataset_num_proc, serial_as_none = True) - if eval_steps is not None and eval_strategy != 'steps': - print(f'Unsloth: `eval_steps = {eval_steps}` is ignored because `eval_strategy` is {getattr(eval_strategy, "value", eval_strategy)!r}. Set `eval_strategy = "steps"` to evaluate every `eval_steps` steps.') - - - # One dict so the filter sees the mirrored parameters AND `**kwargs`: - # filtering kwargs alone would double-bind any argument TRL renamed, - # since the new name is itself a mirrored parameter. - _unsloth_config_arguments = dict( - output_dir = output_dir, - per_device_train_batch_size = per_device_train_batch_size, - num_train_epochs = num_train_epochs, - max_steps = max_steps, - learning_rate = learning_rate, - lr_scheduler_type = lr_scheduler_type, - lr_scheduler_kwargs = lr_scheduler_kwargs, - warmup_steps = warmup_steps, - optim = optim, - optim_args = optim_args, - weight_decay = weight_decay, - adam_beta1 = adam_beta1, - adam_beta2 = adam_beta2, - adam_epsilon = adam_epsilon, - optim_target_modules = optim_target_modules, - gradient_accumulation_steps = gradient_accumulation_steps, - average_tokens_across_devices = average_tokens_across_devices, - max_grad_norm = max_grad_norm, - label_smoothing_factor = label_smoothing_factor, - bf16 = bf16, - fp16 = fp16, - bf16_full_eval = bf16_full_eval, - fp16_full_eval = fp16_full_eval, - tf32 = tf32, - gradient_checkpointing = gradient_checkpointing, - gradient_checkpointing_kwargs = gradient_checkpointing_kwargs, - torch_compile = torch_compile, - torch_compile_backend = torch_compile_backend, - torch_compile_mode = torch_compile_mode, - use_liger_kernel = use_liger_kernel, - liger_kernel_config = liger_kernel_config, - use_cache = use_cache, - neftune_noise_alpha = neftune_noise_alpha, - torch_empty_cache_steps = torch_empty_cache_steps, - auto_find_batch_size = auto_find_batch_size, - logging_strategy = logging_strategy, - logging_steps = logging_steps, - logging_first_step = logging_first_step, - log_on_each_node = log_on_each_node, - logging_nan_inf_filter = logging_nan_inf_filter, - include_num_input_tokens_seen = include_num_input_tokens_seen, - log_level = log_level, - log_level_replica = log_level_replica, - disable_tqdm = disable_tqdm, - report_to = report_to, - run_name = run_name, - project = project, - trackio_space_id = trackio_space_id, - eval_strategy = eval_strategy, - eval_steps = eval_steps, - eval_delay = eval_delay, - per_device_eval_batch_size = per_device_eval_batch_size, - prediction_loss_only = prediction_loss_only, - eval_on_start = eval_on_start, - eval_do_concat_batches = eval_do_concat_batches, - eval_use_gather_object = eval_use_gather_object, - eval_accumulation_steps = eval_accumulation_steps, - batch_eval_metrics = batch_eval_metrics, - save_only_model = save_only_model, - save_strategy = save_strategy, - save_steps = save_steps, - save_on_each_node = save_on_each_node, - save_total_limit = save_total_limit, - enable_jit_checkpoint = enable_jit_checkpoint, - push_to_hub = push_to_hub, - hub_token = hub_token, - hub_private_repo = hub_private_repo, - hub_model_id = hub_model_id, - hub_strategy = hub_strategy, - hub_always_push = hub_always_push, - hub_revision = hub_revision, - load_best_model_at_end = load_best_model_at_end, - metric_for_best_model = metric_for_best_model, - greater_is_better = greater_is_better, - ignore_data_skip = ignore_data_skip, - restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint, - full_determinism = full_determinism, - seed = seed, - data_seed = data_seed, - use_cpu = use_cpu, - accelerator_config = accelerator_config, - parallelism_config = parallelism_config, - dataloader_drop_last = dataloader_drop_last, - dataloader_num_workers = dataloader_num_workers, - dataloader_pin_memory = dataloader_pin_memory, - dataloader_persistent_workers = dataloader_persistent_workers, - dataloader_prefetch_factor = dataloader_prefetch_factor, - remove_unused_columns = remove_unused_columns, - label_names = label_names, - train_sampling_strategy = train_sampling_strategy, - length_column_name = length_column_name, - ddp_find_unused_parameters = ddp_find_unused_parameters, - ddp_bucket_cap_mb = ddp_bucket_cap_mb, - ddp_broadcast_buffers = ddp_broadcast_buffers, - ddp_backend = ddp_backend, - ddp_timeout = ddp_timeout, - fsdp = fsdp, - fsdp_config = fsdp_config, - deepspeed = deepspeed, - debug = debug, - skip_memory_metrics = skip_memory_metrics, - do_train = do_train, - do_eval = do_eval, - do_predict = do_predict, - resume_from_checkpoint = resume_from_checkpoint, - warmup_ratio = warmup_ratio, - logging_dir = logging_dir, - local_rank = local_rank, - max_length = max_length, - max_completion_length = max_completion_length, - beta = beta, - label_smoothing = label_smoothing, - loss_type = loss_type, - disable_dropout = disable_dropout, - cpo_alpha = cpo_alpha, - simpo_gamma = simpo_gamma, - alpha = alpha, - generate_during_eval = generate_during_eval, - is_encoder_decoder = is_encoder_decoder, - model_init_kwargs = model_init_kwargs, - trust_remote_code = trust_remote_code, - dataset_num_proc = dataset_num_proc,**kwargs) - super().__init__(**_unsloth_filter_config_init_kwargs(CPOConfig, _unsloth_config_arguments, mirrored_from = __class__)) - self.vllm_sampling_params = vllm_sampling_params - self.unsloth_num_chunks = unsloth_num_chunks - if unsloth_grpo_mini_batch is not None: - if self.generation_batch_size >= unsloth_grpo_mini_batch: - self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch - else: - raise ValueError( - f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, " - f"which is self.per_device_train_batch_size * gradient_accumulation_steps." - ) - self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier - self.max_seq_length = max_seq_length - # Unsloth: keep the reentrant checkpoint path - if getattr(self, 'gradient_checkpointing', False): - _gc_kwargs = getattr(self, 'gradient_checkpointing_kwargs', None) or {} - if _gc_kwargs.get('context_fn') is None and not _gc_kwargs.get('debug', False): - _gc_kwargs['use_reentrant'] = True - self.gradient_checkpointing_kwargs = _gc_kwargs - from unsloth.optimizers.flash_adamw import is_flash_adamw_marked, mark_flash_adamw, unmark_flash_adamw - if _unsloth_flash_adamw or is_flash_adamw_marked(self): - mark_flash_adamw(self) - else: - unmark_flash_adamw(self) - -pass - -class _UnslothCPOTrainer(_BaseTrainer): - r"""""" - - loss_is_scaled_for_ga = False - - _tag_names = ["trl", "cpo"] - _name = "CPO" - _paper = { - "title": "Contrastive Preference Optimization: Pushing the Boundaries of LLM Performance in Machine Translation", - "id": "2401.08417", - # docstyle-ignore - "citation": textwrap.dedent("""\ - @inproceedings{xu2024contrastive, - title = {{Contrastive Preference Optimization: Pushing the Boundaries of LLM Performance in Machine Translation}}, - author = {Haoran Xu and Amr Sharaf and Yunmo Chen and Weiting Tan and Lingfeng Shen and Benjamin Van Durme and Kenton Murray and Young Jin Kim}, - year = 2024, - booktitle = {Forty-first International Conference on Machine Learning, {ICML} 2024, Vienna, Austria, July 21-27, 2024}, - publisher = {OpenReview.net}, - url = {https://openreview.net/forum?id=51iwkioZpn} - }"""), - } - - def __init__( - self, - model: PreTrainedModel | nn.Module | str | None = None, - args: CPOConfig | None = None, - data_collator: DataCollator | None = None, - train_dataset: Dataset | None = None, - eval_dataset: Dataset | dict[str, Dataset] | None = None, - processing_class: PreTrainedTokenizerBase - | BaseImageProcessor - | FeatureExtractionMixin - | ProcessorMixin - | None = None, - model_init: Callable[[], PreTrainedModel] | None = None, - callbacks: list[TrainerCallback] | None = None, - optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), - preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, - peft_config: "PeftConfig | None" = None, - compute_metrics: Callable[[EvalLoopOutput], dict] | None = None, - ): - if train_dataset is None: - raise ValueError("`train_dataset` is required") - - if args.model_init_kwargs is None: - model_init_kwargs = {} - elif not isinstance(model, str): - raise ValueError("You passed model_kwargs to the CPOTrainer. But your model is already instantiated.") - else: - model_init_kwargs = args.model_init_kwargs - dtype = model_init_kwargs.get("dtype", "auto") - if dtype is not None: - # Convert to `torch.dtype` if an str is passed - if isinstance(dtype, str) and dtype != "auto": - dtype = getattr(torch, dtype) - if dtype != "auto" and not isinstance(dtype, torch.dtype): - raise ValueError( - f"Invalid `dtype` passed to the CPOConfig. Expected a string with either `torch.dtype` or 'auto', but got {dtype}." - ) - model_init_kwargs["dtype"] = dtype - model_init_kwargs["device_map"] = model_init_kwargs.get("device_map", "auto") - - model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) - model_revision = model_init_kwargs.get("revision") if isinstance(model, str) else None - - # PEFT initializes the adapter weights randomly, so set_seed must be done before creating the model to ensure - # reproducibility. - set_seed(args.seed) - if isinstance(model, str): - model = AutoModelForCausalLM.from_pretrained(model, **model_init_kwargs) - - # PEFT - # Initialize this variable to False. This helps tracking the case when `peft_module_casting_to_bf16` - # has been called in order to properly call autocast if needed. - self._peft_has_been_casted_to_bf16 = False - if False: - if not is_peft_available(): - raise ImportError( - "You passed `peft_config` but the `peft` library is not installed. " - "Install it with `pip install trl[peft]`." - ) - if not isinstance(peft_config, PeftConfig): - raise TypeError( - f"`peft_config` must be a `peft.PeftConfig` instance (e.g. `peft.LoraConfig`), " - f"got {type(peft_config).__name__}." - ) - if is_peft_model(model): - raise ValueError( - "You passed a `PeftModel` instance together with a `peft_config` to the trainer. Please first " - "merge and unload the existing adapter, save the resulting base model, and then pass that base " - "model along with the new `peft_config` to the trainer." - ) - - if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False): - _support_gc_kwargs = hasattr( - args, "gradient_checkpointing_kwargs" - ) and "gradient_checkpointing_kwargs" in list( - inspect.signature(prepare_model_for_kbit_training).parameters - ) - - prepare_model_kwargs = {"use_gradient_checkpointing": args.gradient_checkpointing} - - if _support_gc_kwargs: - prepare_model_kwargs["gradient_checkpointing_kwargs"] = args.gradient_checkpointing_kwargs - - model = prepare_model_for_kbit_training(model, **prepare_model_kwargs) - elif args.gradient_checkpointing: - # For backward compatibility with older versions of transformers - if hasattr(model, "enable_input_require_grads"): - model.enable_input_require_grads() - else: - - def make_inputs_require_grad(module, input, output): - output.requires_grad_(True) - - model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) - - # get peft model with the given config - # ZeRO-3 + PEFT for non-quantized models: - # - PEFT's default autocast_adapter_dtype=True upcasts LoRA adapter params to fp32 even when the base model is bf16. - # - ZeRO-3's _allgather_params_coalesced allocates output buffers using the dtype of the first persistent parameter, - # so mixed-dtype persistent_parameters [bf16 base + fp32 LoRA] cause a TypeError on the first optimizer step. - # - Passing autocast_adapter_dtype=False keeps adapter params in the base model dtype [bf16], fixing the mismatch. - # - This is safe: the fp32 upcast is a QLoRA-specific concern [low-bit quantized base models], not needed for - # non-quantized bf16 training. - # - See: - # - TRL issue: https://github.com/huggingface/trl/issues/6089 - # - Upstream issue: https://github.com/deepspeedai/DeepSpeed/issues/8072 - _is_quantized_model = getattr(model, "is_loaded_in_4bit", False) or getattr( - model, "is_loaded_in_8bit", False - ) - get_peft_model_kwargs = {} - if args.deepspeed_plugin is not None and args.deepspeed_plugin.zero_stage == 3 and not _is_quantized_model: - get_peft_model_kwargs["autocast_adapter_dtype"] = False - model = get_peft_model(model, peft_config, **get_peft_model_kwargs) - if args.bf16 and getattr(model, "is_loaded_in_4bit", False): - peft_module_casting_to_bf16(model) - # If args.bf16 we need to explicitly call `generate` with torch amp autocast context manager - self._peft_has_been_casted_to_bf16 = True - - # For models that use gradient_checkpointing, we need to attach a hook that enables input - # to explicitly have `requires_grad=True`, otherwise training will either silently - # fail or completely fail. - elif args.gradient_checkpointing: - # For backward compatibility with older versions of transformers - if hasattr(model, "enable_input_require_grads"): - model.enable_input_require_grads() - else: - - def make_inputs_require_grad(module, input, output): - output.requires_grad_(True) - - model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) - - if args.generate_during_eval and not (is_wandb_available() or is_comet_available()): - raise ValueError( - "`generate_during_eval=True` requires Weights and Biases or Comet to be installed." - " Please install `wandb` or `comet-ml` to resolve." - ) - - if model is not None: - self.is_encoder_decoder = model.config.is_encoder_decoder - elif args.is_encoder_decoder is None: - raise ValueError("When no model is provided, you need to pass the parameter is_encoder_decoder.") - else: - self.is_encoder_decoder = args.is_encoder_decoder - - if self.is_encoder_decoder: - self.decoder_start_token_id = model.config.decoder_start_token_id - self.pad_token_id = model.config.pad_token_id - - if processing_class is None: - processing_class = AutoTokenizer.from_pretrained( - get_config_model_id(model.config), revision=model_revision, trust_remote_code=args.trust_remote_code - ) - if args.max_length is None: - logger.warning( - "`max_length` is not set in the CPOConfig's init" - " it will default to `512` by default, but you should do it yourself in the future.", - ) - max_length = 512 - else: - max_length = args.max_length - - if args.max_completion_length is None and self.is_encoder_decoder: - logger.warning( - "When using an encoder decoder architecture, you should set `max_completion_length` in the CPOConfig's init" - " it will default to `128` by default, but you should do it yourself in the future.", - ) - max_completion_length = 128 - else: - max_completion_length = args.max_completion_length - - if data_collator is None: - data_collator = DPODataCollatorWithPadding( - pad_token_id=(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)), - is_encoder_decoder=self.is_encoder_decoder, - ) - - if args.remove_unused_columns: - args.remove_unused_columns = False - # warn users - logger.warning( - "When using DPODataCollatorWithPadding, you should set `remove_unused_columns=False` in your TrainingArguments" - " we have set it for you, but you should do it yourself in the future.", - ) - - self.use_dpo_data_collator = True - else: - self.use_dpo_data_collator = False - - # Disable dropout in the model - if args.disable_dropout: - disable_dropout_in_model(model) - - self.max_length = max_length - self.generate_during_eval = args.generate_during_eval - self.max_completion_length = max_completion_length - self.processing_class = processing_class - - _unsloth_proc_tok = getattr(processing_class, 'tokenizer', processing_class) - if getattr(_unsloth_proc_tok, 'pad_token', None) is None: - _unsloth_proc_tok.pad_token = getattr(_unsloth_proc_tok, 'eos_token', None) - # The model must agree with the tokenizer on the pad token from construction, so mirror it onto the model - # configs. - model.config.get_text_config().pad_token_id = (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)) - model.generation_config.pad_token_id = (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)) - self.pad_token_id = (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)) - - if args.loss_type in ["hinge", "ipo"] and args.label_smoothing > 0: - logger.warning( - f"You are using the {args.loss_type} loss type that does not support label smoothing. The " - "`label_smoothing` parameter will be ignored. Set `label_smoothing` to `0.0` to remove this warning.", - ) - if args.loss_type == "kto_pair": - raise ValueError("Support for kto_pair has been removed in CPOTrainer. Please use KTOTrainer.") - - self.beta = args.beta - self.label_smoothing = args.label_smoothing - self.loss_type = args.loss_type - self.cpo_alpha = args.cpo_alpha - self.aux_loss_enabled = getattr(model.config, "output_router_logits", False) - self.aux_loss_coef = getattr(model.config, "router_aux_loss_coef", 0.0) - if self.aux_loss_enabled and self.aux_loss_coef == 0.0: - logger.warning( - "You set `output_router_logits` to `True` in the model config, but `router_aux_loss_coef` is set to " - "`0.0`, meaning the auxiliary loss will not be used. Either set `router_aux_loss_coef` to a value " - "greater than `0.0`, or set `output_router_logits` to `False` if you don't want to use the auxiliary " - "loss.", - ) - - if args.loss_type == "simpo": - self.simpo_gamma = args.simpo_gamma - - # AlphaPO parameter for reward shaping - self.alpha = args.alpha - - self._stored_metrics = defaultdict(lambda: defaultdict(list)) - - with global_then_local_main_first(): - # Extract the prompt if needed, and apply the chat template if needed - train_dataset = train_dataset.map(maybe_extract_prompt, num_proc=args.dataset_num_proc) - train_dataset = train_dataset.map( - maybe_apply_chat_template, - fn_kwargs={"processing_class": processing_class}, - num_proc=args.dataset_num_proc, - ) - if eval_dataset is not None: - eval_dataset = eval_dataset.map(maybe_extract_prompt, num_proc=args.dataset_num_proc) - eval_dataset = eval_dataset.map( - maybe_apply_chat_template, - fn_kwargs={"processing_class": processing_class}, - num_proc=args.dataset_num_proc, - ) - - # tokenize the dataset - train_dataset = train_dataset.map(self.tokenize_row, num_proc=args.dataset_num_proc) - if eval_dataset is not None: - eval_dataset = eval_dataset.map(self.tokenize_row, num_proc=args.dataset_num_proc) - - # Transformers explicitly set use_reentrant=True in the past to silence a PyTorch warning, but the default was - # never updated once PyTorch switched to recommending use_reentrant=False. Until that change lands upstream - # [see https://github.com/huggingface/transformers/pull/43203] and is released [most likely in 5.0.0], we - # default to the recommended non-reentrant behavior here, while preserving any user-provided value. - if args.gradient_checkpointing and Version(transformers.__version__) < Version("5.0.0"): - args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} - args.gradient_checkpointing_kwargs.setdefault("use_reentrant", False) - - super().__init__( - model=model, - args=args, - data_collator=data_collator, - train_dataset=train_dataset, - eval_dataset=eval_dataset, - processing_class=processing_class, - model_init=model_init, - compute_metrics=compute_metrics, - callbacks=callbacks, - optimizers=optimizers, - preprocess_logits_for_metrics=preprocess_logits_for_metrics, - ) - - # Add tags for models that have been loaded with the correct transformers version - if hasattr(self.model, "add_model_tags"): - self.model.add_model_tags(self._tag_names) - - if not hasattr(self, "accelerator"): - raise AttributeError( - "Your `Trainer` does not have an `accelerator` object. Consider upgrading `transformers`." - ) - - def build_tokenized_answer(self, prompt, answer): - """ - Llama tokenizer does satisfy `enc(a + b) = enc(a) + enc(b)`. It does ensure `enc(a + b) = enc(a) + enc(a + - b)[len(enc(a)):]`. Reference: - https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257 - """ - - tokenizer = getattr(self.processing_class, "tokenizer", self.processing_class) - full_tokenized = tokenizer(prompt + answer, add_special_tokens=False) - prompt_input_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"] - - answer_input_ids = full_tokenized["input_ids"][len(prompt_input_ids) :] - answer_attention_mask = full_tokenized["attention_mask"][len(prompt_input_ids) :] - - # Concat tokens to form `enc(a) + enc(a + b)[len(enc(a)):]` - full_concat_input_ids = np.concatenate([prompt_input_ids, answer_input_ids]) - - # Prepare input tokens for token by token comparison - full_input_ids = np.array(full_tokenized["input_ids"]) - - if len(full_input_ids) != len(full_concat_input_ids): - raise ValueError("Prompt input ids and answer input ids should have the same length.") - - # On some tokenizers, like Llama-2 tokenizer, there are occasions where tokens - # can be merged together when tokenizing prompt+answer. This could result - # on the last token from the prompt being different when tokenized on its own - # vs when done as prompt+answer. - response_token_ids_start_idx = len(prompt_input_ids) - - # If tokenized prompt is different than both prompt+answer, then it means the - # last token has changed due to merging. - if prompt_input_ids != full_tokenized["input_ids"][:response_token_ids_start_idx]: - response_token_ids_start_idx -= 1 - - prompt_input_ids = full_tokenized["input_ids"][:response_token_ids_start_idx] - prompt_attention_mask = full_tokenized["attention_mask"][:response_token_ids_start_idx] - - if len(prompt_input_ids) != len(prompt_attention_mask): - raise ValueError("Prompt input ids and attention mask should have the same length.") - - answer_input_ids = full_tokenized["input_ids"][response_token_ids_start_idx:] - answer_attention_mask = full_tokenized["attention_mask"][response_token_ids_start_idx:] - - return dict( - prompt_input_ids=prompt_input_ids, - prompt_attention_mask=prompt_attention_mask, - input_ids=answer_input_ids, - attention_mask=answer_attention_mask, - ) - - def tokenize_row(self, feature, model: PreTrainedModel | nn.Module | None = None) -> dict: - """Tokenize a single row from a CPO specific dataset. - - At this stage, we don't convert to PyTorch tensors yet; we just handle the truncation in case the prompt + - chosen or prompt + rejected responses is/are too long. First we truncate the prompt; if we're still too long, - we truncate the chosen/rejected. - - We also create the labels for the chosen/rejected responses, which are of length equal to the sum of the length - of the prompt and the chosen/rejected response, with `-100` for the prompt tokens. - """ - batch = {} - tokenizer = getattr(self.processing_class, "tokenizer", self.processing_class) - prompt = feature["prompt"] - chosen = feature["chosen"] - rejected = feature["rejected"] - - if not self.is_encoder_decoder: - # Check issues below for more details - # 1. https://github.com/huggingface/trl/issues/907 - # 2. https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257 - # 3. https://github.com/LianjiaTech/BELLE/issues/337 - - if not isinstance(prompt, str): - raise ValueError(f"prompt should be an str but got {type(prompt)}") - prompt_tokens = tokenizer(prompt, add_special_tokens=False) - prompt_tokens = {f"prompt_{k}": v for k, v in prompt_tokens.items()} - - if not isinstance(chosen, str): - raise ValueError(f"chosen should be an str but got {type(chosen)}") - chosen_tokens = self.build_tokenized_answer(prompt, chosen) - - if not isinstance(rejected, str): - raise ValueError(f"rejected should be an str but got {type(rejected)}") - rejected_tokens = self.build_tokenized_answer(prompt, rejected) - - # Last prompt token might get merged by tokenizer and - # it should not be included for generation if that happens - prompt_len_input_ids = len(prompt_tokens["prompt_input_ids"]) - - chosen_prompt_len_input_ids = len(chosen_tokens["prompt_input_ids"]) - rejected_prompt_len_input_ids = len(rejected_tokens["prompt_input_ids"]) - prompt_len_input_ids = min(chosen_prompt_len_input_ids, rejected_prompt_len_input_ids) - - for k, v in prompt_tokens.items(): - prompt_tokens[k] = v[:prompt_len_input_ids] - - # Make sure prompts only have one different token at most an - # and length only differs by 1 at most - num_diff_tokens = sum( - a != b - for a, b in zip(chosen_tokens["prompt_input_ids"], rejected_tokens["prompt_input_ids"], strict=False) - ) - num_diff_len = abs(chosen_prompt_len_input_ids - rejected_prompt_len_input_ids) - if num_diff_tokens > 1 or num_diff_len > 1: - raise ValueError( - "Chosen and rejected prompt_input_ids might only differ on the " - "last token due to tokenizer merge ops." - ) - - # add BOS token to head of prompt. Avoid adding if it's already there - prompt_tokens, chosen_tokens, rejected_tokens = add_bos_token_if_needed( - tokenizer.bos_token_id, - prompt_len_input_ids, - prompt_tokens, - chosen_prompt_len_input_ids, - chosen_tokens, - rejected_prompt_len_input_ids, - rejected_tokens, - ) - - # add EOS token to end of answer. Avoid adding if it's already there - chosen_tokens, rejected_tokens = add_eos_token_if_needed( - tokenizer.eos_token_id, chosen_tokens, rejected_tokens - ) - - # if combined sequence is too long, truncate the response - _unsloth_ul = max(len(chosen_tokens['input_ids']), len(rejected_tokens['input_ids'])) - _unsloth_pl = max(len(chosen_tokens['prompt_input_ids']), len(rejected_tokens['prompt_input_ids'])) - if self.max_length is not None and _unsloth_pl + _unsloth_ul > self.max_length: - _unsloth_keep = min(_unsloth_pl, max(self.max_length - _unsloth_ul, self.max_length // 2)) - _unsloth_keep_start = getattr(self, 'truncation_mode', 'keep_end') == 'keep_start' - for answer_tokens in [chosen_tokens, rejected_tokens, prompt_tokens]: - for k in ['prompt_input_ids', 'prompt_attention_mask']: - answer_tokens[k] = answer_tokens[k][:_unsloth_keep] if _unsloth_keep_start else answer_tokens[k][max(0, len(answer_tokens[k]) - _unsloth_keep):] - for answer_tokens in [chosen_tokens, rejected_tokens]: - for k in ['input_ids', 'attention_mask']: - answer_tokens[k] = answer_tokens[k][: self.max_length - _unsloth_keep] - for answer_tokens in [chosen_tokens, rejected_tokens]: - if len(answer_tokens["prompt_input_ids"]) + len(answer_tokens["input_ids"]) > self.max_length: - for k in ["input_ids", "attention_mask"]: - answer_tokens[k] = answer_tokens[k][ - : max(0, self.max_length - len(answer_tokens["prompt_input_ids"])) - ] - if len(answer_tokens["input_ids"]) == 0: - logger.warning_once( - "The prompt alone fills `max_length`, so truncation leaves an empty completion. Consider " - "increasing `max_length` or filtering out long prompts." - ) - - # Create labels - chosen_sequence_tokens = { - k: chosen_tokens[f"prompt_{k}"] + chosen_tokens[k] for k in ["input_ids", "attention_mask"] - } - rejected_sequence_tokens = { - k: rejected_tokens[f"prompt_{k}"] + rejected_tokens[k] for k in ["input_ids", "attention_mask"] - } - chosen_sequence_tokens["labels"] = chosen_sequence_tokens["input_ids"][:] - chosen_sequence_tokens["labels"][: len(chosen_tokens["prompt_input_ids"])] = [-100] * len( - chosen_tokens["prompt_input_ids"] - ) - rejected_sequence_tokens["labels"] = rejected_sequence_tokens["input_ids"][:] - rejected_sequence_tokens["labels"][: len(rejected_tokens["prompt_input_ids"])] = [-100] * len( - rejected_tokens["prompt_input_ids"] - ) - - for k, toks in { - "chosen_": chosen_sequence_tokens, - "rejected_": rejected_sequence_tokens, - "": prompt_tokens, - }.items(): - for type_key, tokens in toks.items(): - if type_key == "token_type_ids": - continue - batch[f"{k}{type_key}"] = tokens - - else: - chosen_tokens = tokenizer( - chosen, truncation=True, max_length=self.max_completion_length, add_special_tokens=True - ) - rejected_tokens = tokenizer( - rejected, truncation=True, max_length=self.max_completion_length, add_special_tokens=True - ) - prompt_tokens = tokenizer(prompt, add_special_tokens=True) - - batch["chosen_labels"] = chosen_tokens["input_ids"] - batch["rejected_labels"] = rejected_tokens["input_ids"] - batch["prompt_input_ids"] = prompt_tokens["input_ids"] - batch["prompt_attention_mask"] = prompt_tokens["attention_mask"] - - if model is not None and hasattr(model, "prepare_decoder_input_ids_from_labels"): - batch["rejected_decoder_input_ids"] = model.prepare_decoder_input_ids_from_labels( - labels=torch.tensor(batch["rejected_labels"]) - ) - batch["chosen_decoder_input_ids"] = model.prepare_decoder_input_ids_from_labels( - labels=torch.tensor(batch["chosen_labels"]) - ) - - return batch - - @staticmethod - def concatenated_inputs( - batch: dict[str, list | torch.LongTensor], - is_encoder_decoder: bool = False, - padding_value: int = 0, - device: torch.device | None = None, - ) -> dict[str, torch.LongTensor]: - """Concatenate the chosen and rejected inputs into a single tensor. - - Args: - batch: - A batch of data. Must contain the keys 'chosen_input_ids' and 'rejected_input_ids', which are tensors - of shape (batch_size, sequence_length). - is_encoder_decoder: - Whether the model is an encoder-decoder model. - padding_value: - The padding value to use for the concatenated inputs_ids. - device: - The device for the concatenated inputs. - - Returns: - A dictionary containing the concatenated inputs under the key 'concatenated_input_ids'. - """ - concatenated_batch = {} - - if is_encoder_decoder: - max_length = max(batch["chosen_labels"].shape[1], batch["rejected_labels"].shape[1]) - else: - max_length = max(batch["chosen_input_ids"].shape[1], batch["rejected_input_ids"].shape[1]) - - for k in batch: - if k.startswith("chosen") and isinstance(batch[k], torch.Tensor): - if "labels" in k or is_encoder_decoder: - pad_value = -100 - elif k.endswith("_input_ids"): - pad_value = padding_value - elif k.endswith("_attention_mask"): - pad_value = 0 - concatenated_key = k.replace("chosen", "concatenated") - concatenated_batch[concatenated_key] = pad_to_length(batch[k], max_length, pad_value=pad_value) - for k in batch: - if k.startswith("rejected") and isinstance(batch[k], torch.Tensor): - if "labels" in k or is_encoder_decoder: - pad_value = -100 - elif k.endswith("_input_ids"): - pad_value = padding_value - elif k.endswith("_attention_mask"): - pad_value = 0 - concatenated_key = k.replace("rejected", "concatenated") - concatenated_batch[concatenated_key] = torch.cat( - ( - concatenated_batch[concatenated_key], - pad_to_length(batch[k], max_length, pad_value=pad_value), - ), - dim=0, - ).to(device=device) - - if is_encoder_decoder: - concatenated_batch["concatenated_input_ids"] = batch["prompt_input_ids"].repeat(2, 1).to(device=device) - concatenated_batch["concatenated_attention_mask"] = ( - batch["prompt_attention_mask"].repeat(2, 1).to(device=device) - ) - - return concatenated_batch - - def cpo_loss( - self, - policy_chosen_logps: torch.FloatTensor, - policy_rejected_logps: torch.FloatTensor, - ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: - """Compute the CPO loss for a batch of policy and reference model log probabilities. - - Args: - policy_chosen_logps: - Log probabilities of the policy model for the chosen responses. Shape: (batch_size,) - policy_rejected_logps: - Log probabilities of the policy model for the rejected responses. Shape: (batch_size,) - - Returns: - A tuple of three tensors: (losses, chosen_rewards, rejected_rewards). The losses tensor contains the CPO - loss for each example in the batch. The chosen_rewards and rejected_rewards tensors contain the rewards for - the chosen and rejected responses, respectively. - """ - # Apply AlphaPO reward transformation if alpha != 0 - if self.alpha != 0.0: - # Compute probabilities - chosen_probs = torch.exp(policy_chosen_logps) - rejected_probs = torch.exp(policy_rejected_logps) - - # Apply AlphaPO transformation: r = (1 - p^(-alpha)) / alpha - policy_chosen_rewards = (1 - chosen_probs.pow(-self.alpha)) / self.alpha - policy_rejected_rewards = (1 - rejected_probs.pow(-self.alpha)) / self.alpha - - logits = (policy_chosen_rewards - policy_rejected_rewards).to(self.accelerator.device) - else: - # Standard log probability rewards when alpha = 0 - logits = (policy_chosen_logps - policy_rejected_logps).to(self.accelerator.device) - - # The beta is a temperature parameter for the CPO loss, typically something in the range of 0.1 to 0.5. - # We ignore the reference model as beta -> 0. The label_smoothing parameter encodes our uncertainty about the labels and - # calculates a conservative CPO loss. - - if self.loss_type == "simpo": - gamma_logratios = self.simpo_gamma / self.beta - logits = logits - gamma_logratios - # This reduces to Equation 3 from the CPO paper when label_smoothing -> 0. - losses = ( - -F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing) - - F.logsigmoid(-self.beta * logits) * self.label_smoothing - ) - elif self.loss_type == "sigmoid": - # This reduces to Equation 3 from the CPO paper when label_smoothing -> 0. - losses = ( - -F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing) - - F.logsigmoid(-self.beta * logits) * self.label_smoothing - ) - elif self.loss_type == "hinge": - losses = torch.relu(1 - self.beta * logits) - elif self.loss_type == "ipo": - # eqn (17) of the paper where beta is the regularization parameter for the IPO loss, denoted by tau in the paper. - losses = (logits - 1 / (2 * self.beta)) ** 2 - else: - raise ValueError( - f"Unknown loss type: {self.loss_type}. Should be one of ['sigmoid', 'hinge', 'ipo', 'simpo']" - ) - - # Calculate rewards for logging - if self.alpha != 0.0: - # When using AlphaPO transformation, use the transformed rewards - chosen_rewards = self.beta * policy_chosen_rewards.to(self.accelerator.device).detach() - rejected_rewards = self.beta * policy_rejected_rewards.to(self.accelerator.device).detach() - else: - # Standard log probability rewards - chosen_rewards = self.beta * (policy_chosen_logps.to(self.accelerator.device)).detach() - rejected_rewards = self.beta * (policy_rejected_logps.to(self.accelerator.device)).detach() - - return losses, chosen_rewards, rejected_rewards - - @staticmethod - def get_batch_logps( - logits: torch.FloatTensor, - labels: torch.LongTensor, - average_log_prob: bool = False, - is_encoder_decoder: bool = False, - ) -> torch.FloatTensor: - """Compute the log probabilities of the given labels under the given logits. - - Args: - logits: Logits of the model (unnormalized). Shape: (batch_size, sequence_length, vocab_size) - labels: - Labels for which to compute the log probabilities. Label tokens with a value of `-100` are ignored. - Shape: (batch_size, sequence_length) - average_log_prob: - If True, return the average log probability per (non-masked) token. Otherwise, return the sum of the - log probabilities of the (non-masked) tokens. - is_encoder_decoder: Whether the model is an encoder-decoder model. - - Returns: - A tensor of shape (batch_size,) containing the average/sum log probabilities of the given labels under the - given logits. - """ - if logits.shape[:-1] != labels.shape: - raise ValueError("Logits (batch and sequence length dim) and labels must have the same shape.") - - if not is_encoder_decoder: - labels = labels[:, 1:].clone() - logits = logits[:, :-1, :] - loss_mask = labels != -100 - - # dummy token; we'll ignore the losses on these tokens later - labels[labels == -100] = 0 - - per_token_logps = selective_log_softmax(logits, labels) - - if average_log_prob: - return (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1) - else: - return (per_token_logps * loss_mask).sum(-1) - - def concatenated_forward( - self, model: nn.Module, batch: dict[str, list | torch.LongTensor] - ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: - """Run the given model on the given batch of inputs, concatenating the chosen and rejected inputs together. - - We do this to avoid doing two forward passes, because it's faster for FSDP. - """ - concatenated_batch = self.concatenated_inputs( - batch, - is_encoder_decoder=self.is_encoder_decoder, - padding_value=self.pad_token_id, - device=self.accelerator.device, - ) - len_chosen = batch["chosen_labels"].shape[0] - - model_kwargs = ( - { - "decoder_input_ids": self._shift_right(concatenated_batch["concatenated_labels"]), - } - if self.is_encoder_decoder - else {} - ) - - if self.aux_loss_enabled: - model_kwargs["output_router_logits"] = True - - _unsloth_pref_head = _unsloth_pref_chunked_head(self, model) - _unsloth_pref_prior = os.environ.get('UNSLOTH_RETURN_HIDDEN_STATES') - if _unsloth_pref_head is not None: - os.environ['UNSLOTH_RETURN_HIDDEN_STATES'] = '1' - try: - outputs = model( - concatenated_batch["concatenated_input_ids"], - attention_mask=concatenated_batch["concatenated_attention_mask"], - use_cache=False, - **model_kwargs, - ) - finally: - if _unsloth_pref_head is not None: - if _unsloth_pref_prior is None: - os.environ.pop('UNSLOTH_RETURN_HIDDEN_STATES', None) - else: - os.environ['UNSLOTH_RETURN_HIDDEN_STATES'] = _unsloth_pref_prior - all_logits = outputs.logits - if _unsloth_pref_head is not None: - _unsloth_pref_out, all_logits = _unsloth_pref_chunked_forward( - self, model, _unsloth_pref_head, concatenated_batch, len_chosen, all_logits, 'cpo' - ) - if _unsloth_pref_out is not None: - return _unsloth_pref_out - - def cross_entropy_loss(logits, labels): - if not self.is_encoder_decoder: - # Shift so that tokens < n predict n - logits = logits[..., :-1, :].contiguous() - labels = labels[..., 1:].contiguous() - # Flatten the tokens - loss_fct = nn.CrossEntropyLoss() - logits = logits.view(-1, logits.shape[-1]) - labels = labels.view(-1) - # Enable model parallelism - labels = labels.to(logits.device) - loss = loss_fct(logits, labels) - return loss - - labels = concatenated_batch["concatenated_labels"].clone() - - if self.cpo_alpha == 0: - nll_loss = torch.tensor(0.0).to(self.accelerator.device) - else: - nll_loss = cross_entropy_loss(all_logits[:len_chosen], labels[:len_chosen]) - - all_logps = self.get_batch_logps( - all_logits, - concatenated_batch["concatenated_labels"], - average_log_prob=self.loss_type in ["ipo", "simpo"], - is_encoder_decoder=self.is_encoder_decoder, - ) - - chosen_logps = all_logps[:len_chosen] - rejected_logps = all_logps[len_chosen:] - - chosen_logits = all_logits[:len_chosen] - rejected_logits = all_logits[len_chosen:] - - if self.aux_loss_enabled: - return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, nll_loss, outputs.aux_loss) - - return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, nll_loss) - - def get_batch_loss_metrics( - self, - model, - batch: dict[str, list | torch.LongTensor], - train_eval: Literal["train", "eval"] = "train", - ): - """Compute the CPO loss and other metrics for the given batch of inputs for train or test.""" - metrics = {} - - forward_output = self.concatenated_forward(model, batch) - ( - policy_chosen_logps, - policy_rejected_logps, - policy_chosen_logits, - policy_rejected_logits, - policy_nll_loss, - ) = forward_output[:5] - if self.aux_loss_enabled: - aux_loss = forward_output[5] - - losses, chosen_rewards, rejected_rewards = self.cpo_loss( - policy_chosen_logps, - policy_rejected_logps, - ) - - loss = losses.mean() + self.cpo_alpha * policy_nll_loss - reward_accuracies = (chosen_rewards > rejected_rewards).float() - - prefix = "eval_" if train_eval == "eval" else "" - metrics[f"{prefix}rewards/chosen"] = self.accelerator.gather_for_metrics(chosen_rewards).mean().item() - metrics[f"{prefix}rewards/rejected"] = self.accelerator.gather_for_metrics(rejected_rewards).mean().item() - metrics[f"{prefix}rewards/accuracies"] = self.accelerator.gather_for_metrics(reward_accuracies).mean().item() - metrics[f"{prefix}rewards/margins"] = ( - self.accelerator.gather_for_metrics(chosen_rewards - rejected_rewards).mean().item() - ) - metrics[f"{prefix}logps/rejected"] = ( - self.accelerator.gather_for_metrics(policy_rejected_logps).detach().mean().item() - ) - metrics[f"{prefix}logps/chosen"] = ( - self.accelerator.gather_for_metrics(policy_chosen_logps).detach().mean().item() - ) - metrics[f"{prefix}logits/rejected"] = ( - self.accelerator.gather_for_metrics(policy_rejected_logits.detach().mean()).mean().item() - ) - metrics[f"{prefix}logits/chosen"] = ( - self.accelerator.gather_for_metrics(policy_chosen_logits.detach().mean()).mean().item() - ) - metrics[f"{prefix}nll_loss"] = self.accelerator.gather_for_metrics(policy_nll_loss).detach().mean().item() - - if self.aux_loss_enabled: - loss += self.aux_loss_coef * aux_loss - - return loss, metrics - - def compute_loss( - self, - model: PreTrainedModel | nn.Module, - inputs: dict[str, torch.Tensor | Any], - return_outputs=False, - num_items_in_batch=None, - ) -> torch.Tensor | tuple[torch.Tensor, dict[str, torch.Tensor]]: - compute_loss_context_manager = ( - autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() - ) - - with compute_loss_context_manager: - loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval="train") - - # force log the metrics - self.store_metrics(metrics, train_eval="train") - - if return_outputs: - return (loss, metrics) - return loss - - def generate_from_model(self, model, batch: dict[str, torch.LongTensor]) -> str: - """Generate samples from the model and reference model for the given batch of inputs.""" - - # If one uses `generate_during_eval` with peft + bf16, we need to explicitly call generate with - # the torch amp context manager as some hidden states are silently casted to full precision. - generate_context_manager = ( - autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() - ) - - with generate_context_manager: - policy_output = model.generate( - input_ids=batch["prompt_input_ids"], - attention_mask=batch["prompt_attention_mask"], - max_length=self.max_length, - do_sample=True, - pad_token_id=self.processing_class.pad_token_id, - ) - - policy_output = pad_to_length(policy_output, self.max_length, self.processing_class.pad_token_id) - policy_output_decoded = self.processing_class.batch_decode(policy_output, skip_special_tokens=True) - - return policy_output_decoded - - def prediction_step( - self, - model: PreTrainedModel | nn.Module, - inputs: dict[str, torch.Tensor | Any], - prediction_loss_only: bool, - ignore_keys: list[str] | None = None, - ): - if ignore_keys is None: - if hasattr(model, "config"): - ignore_keys = getattr(model.config, "keys_to_ignore_at_inference", []) - else: - ignore_keys = [] - - prediction_context_manager = ( - autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() - ) - - with torch.no_grad(), prediction_context_manager: - loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval="eval") - - # force log the metrics - self.store_metrics(metrics, train_eval="eval") - - if prediction_loss_only: - return (loss.detach(), None, None) - - # logits for the chosen and rejected samples from model - logits_dict = { - "eval_logits/chosen": metrics["eval_logits/chosen"], - "eval_logits/rejected": metrics["eval_logits/rejected"], - } - logits = [v for k, v in logits_dict.items() if k not in ignore_keys] - logits = torch.tensor(logits, device=self.accelerator.device) - labels = torch.zeros(logits.shape[0], device=self.accelerator.device) - - return (loss.detach(), logits, labels) - - def store_metrics(self, metrics: dict[str, float], train_eval: Literal["train", "eval"] = "train") -> None: - for key, value in metrics.items(): - self._stored_metrics[train_eval][key].append(value) - - def evaluation_loop( - self, - dataloader: DataLoader, - description: str, - prediction_loss_only: bool | None = None, - ignore_keys: list[str] | None = None, - metric_key_prefix: str = "eval", - ) -> EvalLoopOutput: - """ - Overriding built-in evaluation loop to store metrics for each batch. Prediction/evaluation loop, shared by - `Trainer.evaluate()` and `Trainer.predict()`. - - Works both with or without labels. - """ - - # Sample and save to game log if requested (for one batch to save time) - if self.generate_during_eval: - # Generate random indices within the range of the total number of samples - num_samples = len(dataloader.dataset) - random_indices = random.sample(range(num_samples), k=self.args.eval_batch_size) - - # Use dataloader.dataset.select to get the random batch without iterating over the DataLoader - random_batch_dataset = dataloader.dataset.select(random_indices) - random_batch = self.data_collator(random_batch_dataset) - random_batch = self._prepare_inputs(random_batch) - - policy_output_decoded = self.generate_from_model(self.model, random_batch) - - table = pd.DataFrame( - columns=["Prompt", "Policy"], - data=[ - [prompt, pol[len(prompt) :]] - for prompt, pol in zip(random_batch["prompt"], policy_output_decoded, strict=True) - ], - ) - if "wandb" in self.args.report_to: - wandb.log({"game_log": wandb.Table(data=table)}) - - if "comet_ml" in self.args.report_to: - log_table_to_comet_experiment( - name="game_log.csv", - table=table, - ) - - # Base evaluation - initial_output = super().evaluation_loop( - dataloader, description, prediction_loss_only, ignore_keys, metric_key_prefix - ) - - return initial_output - - def log(self, logs: dict[str, float], start_time: float | None = None) -> None: - """ - Log `logs` on the various objects watching training, including stored metrics. - - Args: - logs (`dict[str, float]`): - The values to log. - start_time (`float`, *optional*): - Start time of the training. - """ - # logs either has 'loss' or 'eval_loss' - train_eval = "train" if "loss" in logs else "eval" - # Add averaged stored metrics to logs - for key, metrics in self._stored_metrics[train_eval].items(): - logs[key] = torch.tensor(metrics).mean().item() - del self._stored_metrics[train_eval] - return super().log(logs, start_time) - - def _shift_right(self, input_ids): - if self.decoder_start_token_id is None: - raise ValueError( - "model.config.decoder_start_token_id has to be defined. It is usually set to the pad_token_id." - ) - - # shift inputs to the right - if is_torch_fx_proxy(input_ids): - # Item assignment is not supported natively for proxies. - shifted_input_ids = torch.full(input_ids.shape[:-1] + (1,), self.decoder_start_token_id) - shifted_input_ids = torch.cat([shifted_input_ids, input_ids[..., :-1]], dim=-1) - else: - shifted_input_ids = input_ids.new_zeros(input_ids.shape) - shifted_input_ids[..., 1:] = input_ids[..., :-1].clone() - shifted_input_ids[..., 0] = self.decoder_start_token_id - - if self.pad_token_id is None: - raise ValueError("model.config.pad_token_id has to be defined.") - # replace possible -100 values in labels by `pad_token_id` - shifted_input_ids.masked_fill_(shifted_input_ids == -100, self.pad_token_id) - - return shifted_input_ids - - # Ensure the model card is saved along with the checkpoint - def _save_checkpoint(self, model, trial): - if self.args.hub_model_id is None: - model_name = Path(self.args.output_dir).name - else: - model_name = self.args.hub_model_id.split("/")[-1] - self.create_model_card(model_name=model_name) - super()._save_checkpoint(model, trial) -class UnslothCPOTrainer(_UnslothCPOTrainer): - """ - - Initialize CPOTrainer. - - Args: - model ([`~transformers.PreTrainedModel`]): - The model to train, preferably an [`~transformers.AutoModelForSequenceClassification`]. - args ([`experimental.cpo.CPOConfig`]): - The CPO config arguments to use for training. - data_collator ([`~transformers.DataCollator`]): - The data collator to use for training. If None is specified, the default data collator - ([`experimental.utils.DPODataCollatorWithPadding`]) will be used which will pad the sequences to the - maximum length of the sequences in the batch, given a dataset of paired sequences. - train_dataset ([`~datasets.Dataset`]): - The dataset to use for training. - eval_dataset ([`~datasets.Dataset`]): - The dataset to use for evaluation. - processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.BaseImageProcessor`], [`~transformers.FeatureExtractionMixin`] or [`~transformers.ProcessorMixin`], *optional*): - Processing class used to process the data. If provided, will be used to automatically process the inputs - for the model, and it will be saved along the model to make it easier to rerun an interrupted training or - reuse the fine-tuned model. - model_init (`Callable[[], transformers.PreTrainedModel]`): - The model initializer to use for training. If None is specified, the default model initializer will be - used. - callbacks (`list[transformers.TrainerCallback]`): - The callbacks to use for training. - optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`): - The optimizer and scheduler to use for training. - preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): - The function to use to preprocess the logits before computing the metrics. - peft_config ([`~peft.PeftConfig`], *optional*): - The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped in - a PEFT model. - compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): - The function to use to compute the metrics. Must take a `EvalPrediction` and return a dictionary string to - metric values. - - """ - def __init__( - self, - model = None, - args = None, - data_collator = None, - train_dataset = None, - eval_dataset = None, - processing_class = None, - model_init = None, - callbacks = None, - preprocess_logits_for_metrics = None, - peft_config = None, - compute_metrics = None, - **kwargs - ): - if args is None: args = UnslothCPOConfig() - _unsloth_precision_requested = tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) - use_bf16 = getattr(args, 'bf16', False) - if type(use_bf16) is not bool: use_bf16 = False - use_fp16 = getattr(args, 'fp16', False) - if type(use_fp16) is not bool: use_fp16 = False - force_float32 = False - try: - from unsloth_zoo.device_type import device_is_bf16_supported as _bf16_supported - except Exception: - _bf16_supported = torch.cuda.is_bf16_supported - try: - from unsloth_zoo.device_type import DEVICE_TYPE as _unsloth_device_type - except Exception: - _unsloth_device_type = None - full_finetuning = getattr(model, '_unsloth_full_finetuning', None) - if full_finetuning is None: full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1' - model_forced_float32 = getattr(model, '_unsloth_forced_float32', None) - if model_forced_float32 is None: model_forced_float32 = os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1' - if model_forced_float32 and not (full_finetuning and _bf16_supported()): - print('Unsloth: Switching to float32 training since model cannot work with float16') - force_float32 = True - mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') - dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None) - if dtype is None: dtype = model.get_input_embeddings().weight.dtype - from unsloth_zoo.utils import _get_dtype - dtype = _get_dtype(dtype) - float16 = dtype == torch.float16 - bfloat16 = dtype == torch.bfloat16 - float32 = dtype == torch.float32 - user_float32 = bool(getattr(model, '_unsloth_user_float32', False)) - if full_finetuning: - if bfloat16 and use_fp16: use_fp16 = False - if float16 and use_bf16: use_bf16 = False - if not force_float32 and bfloat16 and use_fp16 and _unsloth_device_type == 'xpu' and not _bf16_supported(): - print('Unsloth: The model was loaded in bfloat16 on purpose, so training uses bfloat16 mixed precision instead of float16.') - use_fp16 = False - use_bf16 = True - args.fp16 = False - args.bf16 = True - if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`') - if not force_float32 and (bfloat16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`') - if force_float32: - # Forced float32 training - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32' and float32 and user_float32 and not _bf16_supported(): - print('Unsloth: Model is in float32 and this GPU has no bfloat16 support, so training stays in float32. Pass fp16 = True to force float16 mixed precision instead.') - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32': - # Mixed precision training. bf16 only if the GPU supports it; V100/T4 use fp16. - use_bf16_amp = (not float16) and (_bf16_supported() or (bfloat16 and _unsloth_device_type == 'xpu')) - args.fp16 = not use_bf16_amp - args.bf16 = use_bf16_amp - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16_amp else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16_amp else 'fp16' - # args.mixed_precision is a new argument which needs to be set now - elif mixed_precision_dtype == 'bfloat16': - # Both False since bfloat16 full finetuning doesn't do any autocasting. - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif use_bf16 or use_fp16: - # transformers <5 exported this itself from fp16/bf16; 5.x dropped the write, so an - # explicit flag left it unset and GRPO readers defaulted to 'fp16', wrapping a - # bfloat16 model in a float16 autocast. See unslothai/unsloth#4891. - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16 else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16 else 'fp16' - - if getattr(args, 'eval_dataset', None) is not None and getattr(args, 'eval_strategy', 'no') == 'no': - args.eval_strategy = 'steps' - if getattr(args, 'eval_steps', None) is None: args.eval_steps = 0.1 - ga_steps = getattr(args, 'gradient_accumulation_steps', None) - if ga_steps is not None and ga_steps > 1: - from transformers import __version__ as transformers_version - if Version(transformers_version) <= Version('4.45.2'): - print('**** Unsloth: Please use our fixed gradient_accumulation_steps by updating transformers, TRL and Unsloth!\n' - '`pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth transformers trl unsloth_zoo`') - if getattr(args, 'eval_strategy', 'no') != 'no': - eval_bsz = getattr(args, 'per_device_eval_batch_size', 8) - eval_generations = getattr(args, 'num_generations_eval', None) or getattr(args, 'num_generations', None) or 1 - if eval_bsz == 8 and args.per_device_train_batch_size < eval_bsz and (args.per_device_train_batch_size * args.world_size) % eval_generations == 0: args.per_device_eval_batch_size = args.per_device_train_batch_size - if getattr(args, 'eval_accumulation_steps', None) is None and ga_steps is not None: args.eval_accumulation_steps = ga_steps - fp16_full_eval = getattr(args, 'fp16_full_eval', False) - if type(fp16_full_eval) is not bool: fp16_full_eval = False - bf16_full_eval = getattr(args, 'bf16_full_eval', False) - if type(bf16_full_eval) is not bool: bf16_full_eval = False - if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True - if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False - if force_float32: - args.bf16_full_eval = False - args.fp16_full_eval = False - elif os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16': - args.bf16_full_eval = True - args.fp16_full_eval = False - elif not bf16_full_eval and not fp16_full_eval: - args.bf16_full_eval = args.bf16 - args.fp16_full_eval = args.fp16 - if tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) != _unsloth_precision_requested: - from unsloth.models.loader_utils import sync_deepspeed_mixed_precision - sync_deepspeed_mixed_precision(args) - if model is not None: - from unsloth.models.loader_utils import apply_deepspeed_zero3_model_setup - apply_deepspeed_zero3_model_setup(model) - if model is not None: - _warnings_issued = getattr(model, 'warnings_issued', None) - if _warnings_issued is None: - model.warnings_issued = {} - elif not isinstance(_warnings_issued, dict): - try: - model.warnings_issued = dict(_warnings_issued) - except Exception: - model.warnings_issued = {} - if 'max_seq_length' not in locals() and not hasattr(args, 'max_seq_length'): - pass - else: - model_max_seq_length = getattr(model, 'max_seq_length', None) - args_max_seq_length = getattr(args, 'max_seq_length', None) - if args_max_seq_length is None and model_max_seq_length is not None: - max_seq_length = model.max_seq_length - if hasattr(args, 'max_seq_length'): args.max_seq_length = max_seq_length - elif args_max_seq_length is not None and model_max_seq_length is not None: - if args_max_seq_length > model_max_seq_length: - print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but ' - 'the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.') - args.max_seq_length = model_max_seq_length - _unsloth_model_msl = getattr(model, 'max_seq_length', None) - if isinstance(_unsloth_model_msl, int) and _unsloth_model_msl > 0 and hasattr(args, 'max_length'): - if args.max_length is None or args.max_length > _unsloth_model_msl: - print('Unsloth: `max_length = ' + str(args.max_length) + '` exceeds the model max_seq_length of ' + str(_unsloth_model_msl) + ', so it is reduced to ' + str(_unsloth_model_msl) + '.') - args.max_length = _unsloth_model_msl - _unsloth_mpl = getattr(args, 'max_prompt_length', 0) - if (_unsloth_mpl is None and args.max_length <= 128) or (_unsloth_mpl is not None and _unsloth_mpl >= args.max_length): - args.max_prompt_length = args.max_length // 2 - if model is not None and hasattr(model, 'for_training'): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right' - if 'processing_class' in locals(): - if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right' - if hasattr(processing_class, 'tokenizer') and hasattr(processing_class.tokenizer, 'padding_side'): processing_class.tokenizer.padding_side = 'right' - __tokenizer = processing_class if 'processing_class' in locals() else tokenizer - from unsloth_zoo.vision_utils import UnslothVisionDataCollator - if not isinstance(data_collator, UnslothVisionDataCollator): - if isinstance(data_collator, DataCollatorForSeq2Seq) and 'labels' not in _unsloth_dataset_column_names(train_dataset): - data_collator = TransformersDataCollatorForLanguageModeling( - __tokenizer, - mlm = False, - mlm_probability = 0.0, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling) and 'labels' in _unsloth_dataset_column_names(train_dataset): - data_collator = DataCollatorForSeq2Seq( - __tokenizer, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - else: - if hasattr(args, 'remove_unused_columns'): args.remove_unused_columns = False - if hasattr(args, 'dataset_text_field'): args.dataset_text_field = '' - if hasattr(args, 'dataset_kwargs'): args.dataset_kwargs = {'skip_prepare_dataset': True} - if not isinstance(data_collator, UnslothVisionDataCollator): - if not hasattr(__tokenizer, 'pad') and hasattr(__tokenizer, 'tokenizer'): - if isinstance(data_collator, DataCollatorForSeq2Seq): - data_collator = DataCollatorForSeq2Seq( - __tokenizer.tokenizer, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling): - data_collator = TransformersDataCollatorForLanguageModeling( - __tokenizer.tokenizer, - mlm = False, - mlm_probability = 0.0, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - other_metrics = [] - - from unsloth_zoo.logging_utils import PatchRLStatistics - PatchRLStatistics('cpo_trainer', other_metrics) - - # [TODO] Fix up DataParallel multiplying batch sizes - # [TODO] DDP works, but DP seems to not work? [TODO] - if getattr(args, "parallel_mode", None) == ParallelMode.NOT_DISTRIBUTED and args.n_gpu > 1: - if getattr(args, "_n_gpu", 1) != 1: - args._n_gpu = 1 - if "model" in locals(): - try: - from unsloth.models._gc_config_default import honour_gradient_checkpointing_off - honour_gradient_checkpointing_off(args, model) - except Exception: - pass - if "model" in locals() and hasattr(model, "for_training"): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if "model" in locals(): - try: - from unsloth.models._ddp_defaults import unsloth_default_ddp_find_unused_parameters - unsloth_default_ddp_find_unused_parameters(args, model) - except Exception: - pass - super().__init__( - model = model, - args = args, - data_collator = data_collator, - train_dataset = train_dataset, - eval_dataset = eval_dataset, - processing_class = processing_class, - model_init = model_init, - callbacks = callbacks, - preprocess_logits_for_metrics = preprocess_logits_for_metrics, - peft_config = peft_config, - compute_metrics = compute_metrics,**kwargs) - # A tensor parallel model (tp_native.tp_trainer_setup) stays on its training path: for_inference would - # route the decoder layers onto decode paths outside the tensor parallel regions. - if "model" in locals() and hasattr(model, "for_inference") and not getattr(self, "_unsloth_tensor_parallel", False): - model.for_inference() - if hasattr(self, 'neftune_hook_handle'): - self.neftune_hook_handle.remove() - if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle - if getattr(args, 'neftune_noise_alpha', None) is not None: - model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha - pass - if hasattr(self, 'accelerator'): - scaler = self.accelerator.scaler - current_model = model - while hasattr(current_model, 'model'): - current_model.accelerator_scaler = scaler - current_model = current_model.model - current_model.accelerator_scaler = scaler - pass - if hasattr(self, 'train'): - self.train = MethodType(prepare_for_training_mode(self.__class__.train), self) - pass - if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'): - _vllm_tok = self.llm.get_tokenizer() - _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None) - if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None: - _vllm_tok.chat_template = _pc.chat_template - pass - if getattr(self, 'aux_loss_enabled', False) and hasattr(getattr(self, 'model', None), 'config'): - _text_config = self.model.config - if hasattr(_text_config, 'get_text_config'): _text_config = _text_config.get_text_config() - _n_experts = [getattr(_text_config, _k) for _k in ('num_local_experts', 'num_experts', 'n_routed_experts', 'moe_num_experts') if isinstance(getattr(_text_config, _k, None), int)] - if _n_experts and all(_n <= 1 for _n in _n_experts): - self.aux_loss_enabled = False - _text_config.output_router_logits = False - pass - -pass - - -if hasattr(logger, "addFilter"): - import logging - class HideLoggingMessage(logging.Filter): - def __init__(self, text): self.text = text - def filter(self, x): return not (self.text in x.getMessage()) - pass - logger.addFilter(HideLoggingMessage("`use_cache=True`")) - diff --git a/unsloth_compiled_cache/UnslothDPOTrainer.py b/unsloth_compiled_cache/UnslothDPOTrainer.py deleted file mode 100644 index b9495f780db4e23c9cd9465b01507d55dc03fdd3..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/UnslothDPOTrainer.py +++ /dev/null @@ -1,3233 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from unsloth_zoo.temporary_patches.common import torch_compile -from unsloth_zoo.temporary_patches.common import _maybe_compile -import functools -import transformers -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from trl.trainer.dpo_trainer import (Any, AutoProcessor, BaseTunerLayer, BitsAndBytesConfig, Callable, DPOConfig, DPOTrainer, DataCollator, DataCollatorForPreference, DataCollatorForVisionPreference, DataLoader, Dataset, DatasetDict, EvalPrediction, F, Hasher, IterableDataset, IterableDatasetDict, LoraConfig, Path, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, SyncRefModelCallback, TrainerCallback, Version, _BaseTrainer, _tokenize, accelerate, add_fused_lm_head, apply_chat_template, broadcast_object_list, common_prefix_length, concatenate_datasets, contextlib, create_model_from_path, dataclass, defaultdict, disable_dropout_in_model, disable_gradient_checkpointing, extract_prompt, get_act_offloading_ctx_manager, get_config_model_id, get_peft_model, global_then_local_main_first, hash_module, is_conversational, is_peft_available, is_peft_model, json, logger, math, os, pad, peft, prepare_deepspeed, prepare_fsdp, set_seed, textwrap, torch, tqdm, transformers, use_adapter, warnings, AutoProcessor, BaseTunerLayer, BitsAndBytesConfig, Callable, DPOConfig, DPOTrainer, DataCollator, DataCollatorForPreference, DataCollatorForVisionPreference, Dataset, DatasetDict, EvalPrediction, F, IterableDataset, IterableDatasetDict, LoraConfig, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, SyncRefModelCallback, TrainerCallback, Version, accelerate, add_fused_lm_head, contextlib, create_model_from_path, defaultdict, disable_dropout_in_model, get_act_offloading_ctx_manager, get_config_model_id, get_peft_model, hash_module, is_peft_available, is_peft_model, logger, os, pad, peft, prepare_deepspeed, prepare_fsdp, set_seed, torch, transformers, warnings, F, PeftModel, PreTrainedModel, is_peft_available, logger, os, peft, torch) - - -import os -import math -import logging -from typing import * -from dataclasses import dataclass, field -from packaging.version import Version -import torch -import numpy as np -from contextlib import nullcontext -from torch.nn import functional as F -import inspect -from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling -from transformers.training_args import ParallelMode -from unsloth_zoo.device_type import DEVICE_TYPE, DEVICE_TYPE_TORCH, device_synchronize - -# Wrap trainer with padding to right and enable training mode -import functools -from types import MethodType -try: - from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers -except: - def reset_unsloth_gradient_checkpointing_buffers(): pass -# Canonical reset lives in unsloth.models._utils so the SFT auto-packing wrapper and the plain -# Trainer loop can import the same helper; fall back to a no-op only if it can't be imported. -try: - from unsloth.models._utils import _unsloth_reset_stray_compile_cache -except Exception: - def _unsloth_reset_stray_compile_cache(self): pass -try: - from unsloth.models._utils import _unsloth_dataset_column_names -except Exception: - def _unsloth_dataset_column_names(dataset): return dataset.column_names -# Drops/renames config arguments the installed TRL no longer accepts, so a -# script pinned to an older TRL keeps working after an upgrade. Falls back to -# the historical raw passthrough so this can never break trainer construction. -try: - from unsloth.models.rl_config_compat import filter_config_init_kwargs as _unsloth_filter_config_init_kwargs - # A cache file generated here can be imported by an older Unsloth whose filter - # predates `mirrored_from`, so drop the argument rather than raise TypeError. - if "mirrored_from" not in inspect.signature(_unsloth_filter_config_init_kwargs).parameters: - _unsloth_filter_config_init_kwargs_old = _unsloth_filter_config_init_kwargs - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): - return _unsloth_filter_config_init_kwargs_old(config_class, kwargs) -except Exception: - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): return kwargs -def prepare_for_training_mode(f): - @functools.wraps(f) - def wrapper(self, *args, **kwargs): - # Drop any torch.compile graph cache poisoned by a stray pre-train forward. - try: - _unsloth_reset_stray_compile_cache(self) - except Exception: - pass - # Finish the previous W&B run if this is a subsequent train() call. - # We do this at the START of train() (not the end) so that - # evaluate() / log() still work after train() completes. - # HF's WandbCallback.setup() will call wandb.init() for the new run. - # See: https://github.com/unslothai/unsloth/issues/3954 - if getattr(self, '_unsloth_training_completed', False): - try: - import wandb - if wandb.run is not None: - wandb.finish() - # Reset HF's WandbCallback so it calls wandb.init() for the new run - for cb in self.callback_handler.callbacks: - if type(cb).__name__ == 'WandbCallback': - cb._initialized = False - break - except: - pass - # Enable training mode - _was_training = None - # Restore the GC mode the model was configured with at setup; fall back to - # the training args only when it wasn't recorded (issue #4735). Use hasattr, - # not a None sentinel, so a deliberately-recorded None is restored verbatim. - _model = getattr(self, 'model', None) - if hasattr(_model, '_unsloth_gradient_checkpointing'): - use_gc = _model._unsloth_gradient_checkpointing - else: - use_gc = getattr(self.args, 'gradient_checkpointing', True) - if hasattr(self, 'model') and hasattr(self.model, "training"): - _was_training = self.model.training - if hasattr(self, 'model') and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - output = f(self, *args, **kwargs) - # Restore previous mode when possible - if hasattr(self, 'model') and hasattr(self.model, "for_inference"): - if _was_training is False: - self.model.for_inference() - elif _was_training is True and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - # Reset gradient checkpointing buffers to free memory while staying ready for next run - try: - reset_unsloth_gradient_checkpointing_buffers() - except: - pass - # Mark that training completed so the next train() call can - # finish this W&B run before starting a new one - self._unsloth_training_completed = True - return output - return wrapper -pass - -torch_compile_options = { - "epilogue_fusion" : True, - "max_autotune" : False, - "shape_padding" : True, - "trace.enabled" : False, - "triton.cudagraphs" : False, -} - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_hidden_states_selective_log_softmax( - hidden_states: torch.Tensor, - lm_head: torch.Tensor, - index: torch.Tensor, - chunks: int = 4, - logit_scale_multiply: float = 0.0, - logit_scale_divide: float = 0.0, - logit_softcapping: float = 0.0, - temperature: float = 1.0, - # Rows per chunk cap. Read HERE, in the default, not in the body: the body - # is traced with fullgraph = True, and `os.environ` is an unsupported op - # there on torch 2.4 -- Dynamo raises - # torch._dynamo.exc.Unsupported: const method call bytes.decode - # from os._Environ.__getitem__, which the eager fallback does not catch - # (it only catches recompile-limit and disabled-hook breaks), so the very - # first call would die even with the variable unset. A default is evaluated - # once when the def runs, which is import time, outside any traced region. - # It is also a plain int argument, so Dynamo guards on it instead of - # constant-folding an unguarded read (2.7+ never notice a later change). - # A non-numeric value is ignored rather than raised on, so a typo cannot - # break the import. 0 keeps the previous chunk boundaries exactly. - max_rows_per_chunk: int = ( - int(os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip()) - if os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip().isdigit() - else 0 - ), - lm_head_bias: Optional[torch.Tensor] = None, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, - # Read in the default for the same Dynamo reason as max_rows_per_chunk. Pairs with vLLM's - # head_dtype=float32 that load_vllm requests under the same flag. - fp32_lm_head: bool = ( - os.environ.get("UNSLOTH_VLLM_FP32_LM_HEAD", "0").strip().lower() in ("1", "true", "on", "yes") - ), - # Older Unsloth scores old / ref log-probs here and the policy with the fused kernel: one scorer keeps ratio at 1. - fused_linear: bool = ( - os.environ.get("UNSLOTH_FUSED_LINEAR", "1") != "0" - and os.environ.get("UNSLOTH_GRPO_CHUNKED_LOGPROBS", "1") != "0" - ), -) -> torch.Tensor: - # All Unsloth Zoo code licensed under AGPL3 - from unsloth_zoo.fused_losses import chunked_logprob as _chunked_logprob - if (fused_linear and _chunked_logprob._FUSED_LINEAR and not fp32_lm_head - and support_top_k == 0 and support_top_p == 1.0 and support_min_p == 0.0 - and support_excluded_ids is None): - return _chunked_logprob.unsloth_chunked_logprobs( - hidden_states, lm_head, index, - temperature = temperature, - logit_scale_multiply = logit_scale_multiply, - logit_scale_divide = logit_scale_divide, - logit_softcapping = logit_softcapping, - lm_head_bias = lm_head_bias, - max_rows_per_chunk = max_rows_per_chunk if max_rows_per_chunk > 0 else None, - ).reshape((hidden_states.shape[0], hidden_states.shape[1])) - # Reshape on this tensor's own last dim: a no-op, so a wrong-width caller - # cannot have its row count silently rewritten and instead fails at the - # matmul below, which prints both operands. Do not swap in a bare - # torch._check: it reports only "Expected cond to be True", naming neither - # operand, and Dynamo rejects a message-carrying one. Callers dispatch on - # the width first -- see `compute_logprobs_chunk`, the packed path and - # `_pg_grad_forward`. - flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) - flat_index = index.reshape(-1) - - # Each chunk materialises rows x vocab logits and then a float32 copy of - # them, all on the device holding the output head. With a large vocabulary - # and a fixed chunk count that grows with the batch, so the peak scales with - # the batch rather than staying bounded. max_rows_per_chunk caps the rows - # per chunk instead, which is pure loop splitting: more, smaller chunks, - # same concatenated result. 0 (the default) keeps the previous chunk - # boundaries exactly. - if max_rows_per_chunk > 0: - n_rows = flat_hidden_states.shape[0] - chunks = max(chunks, -(-n_rows // max_rows_per_chunk)) - chunks = min(chunks, max(n_rows, 1)) - - # A one-row chunk's head gradient is a K=1 matmul Inductor mis-lowers under dynamic shapes (illegal - # memory access or NaN head gradients, torch 2.13), so keep chunks at 2+ rows when the head trains. - # Exactly `chunks` row ranges, so the unrolled loop length never depends on T: torch.chunk - # returns ceil(T / ceil(T / chunks)) pieces and recompiled per new count; tensor_split guards - # on T % chunks. At most rows // 2 chunks when the head trains keeps every piece at 2+ rows. - lm_head_grad = torch.is_grad_enabled() and lm_head.requires_grad - if lm_head_grad and flat_hidden_states.shape[0] < 2 * chunks: - chunks = max(flat_hidden_states.shape[0] // 2, 1) - - n_rows = flat_hidden_states.shape[0] - chunks = max(1, min(chunks, n_rows)) - bounds = [(i * n_rows) // chunks for i in range(chunks + 1)] - chunked_hidden_states = [flat_hidden_states[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - - # Cast once, outside the loop, so every chunk shares one fp32 head copy (also for autograd). - if fp32_lm_head: - lm_head = lm_head.to(torch.float32) - if lm_head_bias is not None: - lm_head_bias = lm_head_bias.to(torch.float32) - - all_per_token_logps = [] - - for chunk_hidden_states, chunk_index in zip(chunked_hidden_states, chunked_index): - # When the model is dispatched over several devices, the output head can - # sit on a different one from the hidden states, because accelerate - # places the tail of the model on the last device it fills. Co-locate on - # the head's device before the matmul, otherwise this raises - # Unhandled FakeTensor Device Propagation for aten.mm.default, - # found two different devices cuda:0, cuda:1 - # On a single device every .to() here is a no-op and the result is - # bit-identical to before. - chunk_hidden_states = chunk_hidden_states.to(device = lm_head.device, dtype = lm_head.dtype) - chunk_index = chunk_index.to(lm_head.device) - if fp32_lm_head: - # The GRPO loss runs this under bf16 / fp16 autocast, which would downcast the matmul. - with torch.autocast(device_type = chunk_hidden_states.device.type, enabled = False): - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - elif lm_head_bias is None: - chunk_logits = chunk_hidden_states @ lm_head.t() - else: - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - - if logit_scale_multiply != 0.0: - chunk_logits = chunk_logits * logit_scale_multiply - if logit_scale_divide != 0.0: - chunk_logits = chunk_logits / logit_scale_divide - if logit_softcapping != 0.0: - chunk_logits = logit_softcapping * torch.tanh(chunk_logits / logit_softcapping) - - chunk_logits = chunk_logits.to(torch.float32) - - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - per_token_logps = _truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - ) - else: - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - # Return to the caller's device so the concatenation below and every - # downstream consumer see the device they started on. - all_per_token_logps.append(per_token_logps.to(hidden_states.device)) - - all_per_token_logps = torch.concat(all_per_token_logps) - - all_per_token_logps = all_per_token_logps.reshape((hidden_states.shape[0], hidden_states.shape[1])) - return all_per_token_logps - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_selective_log_softmax( - logits, - index, - temperature: float = 1.0, - chunks: int = 4, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, -): - # Fixed `chunks` row ranges; see chunked_hidden_states_selective_log_softmax. - flat_logits = logits.reshape(-1, logits.shape[-1]) - flat_index = index.reshape(-1) - # Fewer rows than chunks: one row per range, as torch.chunk did, never empty ranges. - chunks = max(1, min(chunks, flat_logits.shape[0])) - bounds = [(i * flat_logits.shape[0]) // chunks for i in range(chunks + 1)] - chunked_logits = [flat_logits[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - all_per_token_logps = [] - # Per-chunk selective_log_softmax. - for chunk_logits, chunk_index in zip(chunked_logits, chunked_index): - chunk_logits = chunk_logits.to(torch.float32) - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - all_per_token_logps.append(_truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - )) - continue - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - all_per_token_logps.append(per_token_logps) - pass - all_per_token_logps = torch.concat(all_per_token_logps) - all_per_token_logps = all_per_token_logps.reshape((logits.shape[0], logits.shape[1])) - return all_per_token_logps - -def calculate_pad_tokens_in_prompt( - input_ids: torch.Tensor, - logits_to_keep: int, - pad_token_id: int -) -> torch.Tensor: - """Count left-padded tokens per sequence, e.g. [pad, pad, pad, cat] -> 3.""" - if logits_to_keep >= input_ids.shape[1]: - raise ValueError("logits_to_keep must be smaller than the sequence length.") - - prompt_section = input_ids[:, :-logits_to_keep] - - padding_mask = (prompt_section == pad_token_id) - - pad_token_counts = padding_mask.sum(dim=1) - - return pad_token_counts - -def create_completion_attention_mask( - completion_input_ids: torch.Tensor, - left_pad_tokens_per_prompt: torch.Tensor, - max_left_pad: int, - pad_token_id: int -) -> torch.Tensor: - """Build a completion mask that zeros leading prompt and trailing pad tokens. - - For [p,p,p,c,c,c,pad,pad,pad] (p=sliced prompt, c=completion, pad=padding) - this returns [0,0,0,1,1,1,0,0,0]. - """ - batch_size, completion_len = completion_input_ids.shape - device = completion_input_ids.device - - num_tokens_to_mask = max_left_pad - left_pad_tokens_per_prompt - - indices = torch.arange(completion_len, device=device).unsqueeze(0) - shift_mask = indices >= num_tokens_to_mask.unsqueeze(1) - - non_padding_mask = (completion_input_ids != pad_token_id) - - final_mask = shift_mask & non_padding_mask - - return final_mask - -def left_pack_padding(tensor: torch.Tensor, pad_id: int) -> torch.Tensor: - """Move all padding tokens in each sequence to the right.""" - mask = (tensor != pad_id) - # stable=True since the binary mask is unordered. - sorted_indices = torch.argsort(mask, dim=1, descending=True, stable=True) - packed_tensor = torch.gather(tensor, 1, sorted_indices) - return packed_tensor - -def align_logprobs_with_mask( - logprob_tensor: torch.Tensor, - attention_mask: torch.Tensor, - pad_value: float = 0.0 -) -> torch.Tensor: - """Align a log probability tensor with a given attention mask.""" - - device = logprob_tensor.device - batch_size, logprob_seq_len = logprob_tensor.shape - mask_seq_len = attention_mask.shape[1] - - padded_logprobs = torch.full( - attention_mask.shape, - fill_value=pad_value, - dtype=logprob_tensor.dtype, - device=device - ) - - left_pad_counts = torch.argmax(attention_mask, dim=1) - - cols = torch.arange(logprob_seq_len, device=device) - dest_indices = left_pad_counts.unsqueeze(1) + cols - - # Destination row indices, shape [batch_size, logprob_seq_len]. - row_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand_as(dest_indices) - - # Keep only in-bounds destinations, then scatter via advanced indexing. - valid_mask = dest_indices < mask_seq_len - valid_rows = row_indices[valid_mask] - valid_cols = dest_indices[valid_mask] - valid_vals = logprob_tensor[valid_mask] - padded_logprobs[valid_rows, valid_cols] = valid_vals - - return padded_logprobs - -def align_completion_tool_mask( - tool_mask: torch.Tensor, - completion_mask: torch.Tensor, -) -> torch.Tensor: - """Align a raw completion-length tool/env mask with Unsloth's repacked loss mask.""" - if tool_mask is None: - return completion_mask - if tool_mask.shape[0] != completion_mask.shape[0]: - raise ValueError("tool_mask batch size must match completion_mask batch size.") - - tool_mask = tool_mask.to(device=completion_mask.device) - if tool_mask.shape == completion_mask.shape: - aligned_tool_mask = tool_mask - else: - aligned_tool_mask = align_logprobs_with_mask( - tool_mask, - completion_mask, - pad_value=0, - ) - return completion_mask * aligned_tool_mask.to(dtype=completion_mask.dtype) - -def autotune_batch_and_chunks( - total_input_rows, - seq_len, - hidden_size, - vocab_size, - dtype_bytes=16, - multiplier=None -): - if multiplier is None: - final_m = max(4, seq_len // 4096) - else: - final_m = multiplier - - if torch.cuda.is_available(): - free_bytes, _ = torch.cuda.mem_get_info() - limit_gb = (free_bytes / (1024**3))*.80 - elif hasattr(torch, "xpu") and torch.xpu.is_available(): - # XPU: estimate free memory as total - reserved. - total_mem = torch.xpu.get_device_properties(0).total_memory - reserved_mem = torch.xpu.memory_reserved() - free_bytes = total_mem - reserved_mem - limit_gb = (free_bytes / (1024**3)) * 0.80 - else: - # Fallback: assume 8GB available. - limit_gb = 8.0 - - bytes_to_gb = 1024**3 - - b_vals = torch.arange(total_input_rows, 0, -1, device='cpu', dtype=torch.float32) - - hidden_gb = (b_vals * seq_len * hidden_size * dtype_bytes) / bytes_to_gb - - base_logits = ((b_vals/total_input_rows) * b_vals * seq_len * vocab_size * dtype_bytes) / bytes_to_gb - logits_gb = base_logits / final_m - - total_mem_gb = hidden_gb + logits_gb - - valid_mask = total_mem_gb <= limit_gb - valid_indices = torch.nonzero(valid_mask, as_tuple=False) - - if valid_indices.shape[0] == 0: - #This means your GPU will OOM - # Capped at the row count: unsloth's no-grad pass divides rows by this without max(1, ...). - return max(1, min(4, total_input_rows)), final_m - - best_idx = valid_indices[0].item() - final_b = int(b_vals[best_idx].item()) - - return final_b, final_m - -def sanitize_logprob(logprob): - """Local port of trl.scripts.vllm_serve.sanitize_logprob. - Filters NaN logprobs from vLLM outputs.""" - value = logprob.logprob - if math.isnan(value): - logging.getLogger(__name__).warning( - f"Generated NaN logprob, token logprob '{logprob}' will be ignored" - ) - return None - return value -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 _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_text_configs(config): - """``config`` and its text sub-config, the two places a transform can live: Gemma-4-style configs keep it on ``config.text_config``, T5Gemma-style ones only reach it via ``config.get_text_config()``.""" - if config is None: - return [] - holders = [config] - 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: - holders.append(text_cfg) - return holders - -def _unsloth_get_final_logit_softcapping(model): - """The soft cap the loss applies, under any of its three spellings: ``final_logit_softcapping`` (Gemma), ``logits_soft_cap`` (RecurrentGemma), ``output_logit_soft_cap`` (xLSTM). Returns 0 if unset.""" - config = _unsloth_get_model_config(model) - if config is None: - return 0 - for holder in _unsloth_text_configs(config): - for name in ("final_logit_softcapping", "logits_soft_cap", "output_logit_soft_cap"): - softcap = getattr(holder, name, None) - if softcap: - return softcap - return 0 - -def _unsloth_resolve_logit_scales(model_config): - """``(multiply, divide)`` for the logits, read with the same field table and per-family overrides ``detect_logit_transforms`` uses. - - Only reached when the installed unsloth_zoo predates that helper. Must stay in step with ``resolve_logit_transforms`` in unsloth/models/llama.py: if the forward applies a transform GRPO does not, the policy log-probabilities come from different logits than the ones generated, which silently shifts every importance ratio. - """ - # Granite divides, HyperCLOVA X multiplies, MiniCPM3 scales hidden states (not a logit transform). - overrides = { - ("logits_scaling", "hyperclovax"): "multiply", - ("logits_scaling", "minicpm3"): None, - } - found = {"multiply": 0, "divide": 0} - for holder in _unsloth_text_configs(model_config): - model_type = getattr(holder, "model_type", "") or "" - for bucket, names in ( - ("multiply", ("logit_scale", "lm_head_multiplier", "output_multiplier")), - ("divide", ("logits_scaling",)), - ): - for name in names: - target = overrides.get((name, model_type), bucket) - if target is None or found[target]: - continue - value = getattr(holder, name, 0) or 0 - if value: - found[target] = value - break - return found["multiply"], found["divide"] - -def _unsloth_grpo_returns_hidden_states(model, tensor, lm_head): - """Does ``tensor`` (a forward's ``.logits``) carry hidden states or real logits? - - ``_get_per_token_logps_and_entropies`` sets ``UNSLOTH_RETURN_HIDDEN_STATES=1``, but only a forward that honours the name hands hidden states back as ``.logits``; any other forward returns a real ``[.., vocab]`` tensor that must not reach the ``lm_head`` matmul. - - The primary test is an explicit signal that the forward honours the flag, both set outside this file: ``__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__`` on the generated class, written by ``unsloth_zoo.compiler.create_standalone_class`` exactly when ``apply_fused_lm_head`` gave that forward its own branch; and ``_unsloth_grpo_hidden_states_forward_wrapped``, set by ``_install_grpo_hidden_states_forward_wrapper`` in ``unsloth/models/rl.py`` for models the compiler did not rewrite. That wrapper degrades to real logits when the model cannot produce hidden states and records whether it did so in ``_unsloth_grpo_hidden_states_degraded`` before returning, so reading the pair after a forward describes the call that just finished; degradation is per call, not per model. - - The width comparison stays as the fallback, for an ``unsloth_zoo`` old enough that it never writes the marker. It is decisive whenever ``vocab_size != hidden_size``, and the signal may only overrule it when it is not, which is the one case the shape cannot answer. - """ - if tensor.shape[-1] != lm_head.shape[1]: - return False # vocab-wide: real logits, whatever any signal claims - if lm_head.shape[0] != lm_head.shape[1]: - return True # hidden-wide and vocab_size != hidden_size: hidden states - return _unsloth_grpo_hidden_states_signal(model) is not False - -def _unsloth_grpo_hidden_states_signal(model): - """``True``/``False`` if the forward honours ``UNSLOTH_RETURN_HIDDEN_STATES``, ``None`` when neither marker is present. See ``_unsloth_grpo_returns_hidden_states`` for where each marker is set. Walks the wrapper chain because the markers are set on whichever object the trainer saw, which may be the DDP module or the PEFT base model rather than the object handed to the logprob loop.""" - candidates = [] - pending = [model] - while pending and len(candidates) < 8: - candidate = pending.pop(0) - if candidate is None or any(candidate is seen for seen in candidates): - continue - candidates.append(candidate) - get_base_model = getattr(candidate, "get_base_model", None) - if callable(get_base_model): - try: - pending.append(get_base_model()) - except Exception: - pass - for _attr in ("module", "base_model", "model"): - child = getattr(candidate, _attr, None) - if child is not None and hasattr(child, "forward"): - pending.append(child) - for candidate in candidates: - if getattr(candidate, "__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__", False): - return True - if getattr(type(candidate), "__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__", False): - return True - if any( - getattr(candidate, "_unsloth_grpo_hidden_states_forward_wrapped", False) - for candidate in candidates - ): - # The wrapper honours the flag unless it recorded that this call could not. - if any( - hasattr(candidate, "_unsloth_grpo_hidden_states_degraded") for candidate in candidates - ): - return not any( - getattr(candidate, "_unsloth_grpo_hidden_states_degraded", False) - for candidate in candidates - ) - # An unsloth/models/rl.py predating the per-call attribute set only the warn-once flag; it is the best signal such a wrapper offers. - return not any( - getattr(candidate, "_unsloth_grpo_hidden_states_warning_issued", False) - for candidate in candidates - ) - return None - -def _unsloth_pref_returns_hidden_states(model, tensor, lm_head): - """``_unsloth_grpo_returns_hidden_states`` for the preference routes, which need a zoo that writes the hidden-state - markers: a square head with no signal is a forward that ignored the flag, so its output is real logits.""" - if tensor.shape[-1] != lm_head.shape[1]: - return False - if lm_head.shape[0] != lm_head.shape[1]: - return True - return _unsloth_grpo_hidden_states_signal(model) is True - -def _unsloth_gkd_dense_head(model): - """The output head when it is a plain dense ``[vocab, hidden]`` projection the chunked loss can read directly. - Rejects bnb heads (``nn.Linear`` subclasses with packed weights), PEFT heads and ZeRO-3 partitioned weights. - """ - get_output_embeddings = getattr(model, "get_output_embeddings", None) - if not callable(get_output_embeddings): - return None - try: - head = get_output_embeddings() - except Exception: - return None - if type(head) is not torch.nn.Linear: - return None - # An lm_head that is not the output embeddings runs before them (ModernBERT decoder, RoBERTa-style heads). - lm_head = getattr(model, "lm_head", None) - if isinstance(lm_head, torch.nn.Module) and lm_head is not head: - return None - weight = getattr(head, "weight", None) - if not isinstance(weight, torch.Tensor) or weight.dim() != 2 or weight.numel() == 0: - return None - if weight.dtype not in (torch.float16, torch.bfloat16, torch.float32, torch.float64): - return None - if type(weight) not in (torch.Tensor, torch.nn.Parameter): - return None - bias = getattr(head, "bias", None) - if bias is not None and (bias.dim() != 1 or bias.shape[0] != weight.shape[0]): - return None - if weight.is_meta or (bias is not None and bias.is_meta): - return None - return head - -def _unsloth_gkd_logit_transforms(model): - """``(scale, softcap)`` exactly as the model's own forward applies them to its logits.""" - model_config = _unsloth_get_model_config(model) - if detect_logit_transforms is not None: - _transforms = detect_logit_transforms(model_config) - logit_softcapping = _transforms["logit_softcapping"] - logit_scale_multiply = _transforms["logit_scale_multiply"] - logit_scale_divide = _transforms["logit_scale_divide"] - else: - logit_softcapping = _unsloth_get_final_logit_softcapping(model) - logit_scale_multiply, logit_scale_divide = _unsloth_resolve_logit_scales(model_config) - scale = 1.0 - if logit_scale_multiply: - scale = scale * float(logit_scale_multiply) - if logit_scale_divide: - scale = scale / float(logit_scale_divide) - return scale, float(logit_softcapping or 0.0) - -def _unsloth_pref_logit_transforms(model): - """``(multiply, divide, softcap)`` for the chunked kernel, from the GKD resolver (one scale, 0.0 = off).""" - scale, softcap = _unsloth_gkd_logit_transforms(model) - return (0.0 if scale == 1.0 else scale), 0.0, softcap - -def _unsloth_stock_logps_hook( - self, - method, - hooks = ("get_batch_logps",), -): - """False when a subclass overrides one of ``hooks`` below the class whose ``method`` is the rewritten one: TRL's - body would call the override (or it reads the full logits the chunked path replaces with stand-ins).""" - owner = None - for klass in type(self).__mro__: - # The rewritten body calls a chunked or fused-pass helper; a subclass method wrapping it via super() does not. - names = getattr(getattr(vars(klass).get(method), "__code__", None), "co_names", ()) - if any( - name.startswith("_unsloth_") and ("chunked" in name or "fused_pass" in name) - for name in names - ): - owner = klass - break - if owner is None: - return True - for name in hooks: - hook = next((klass for klass in type(self).__mro__ if name in vars(klass)), None) - if hook is not None and hook not in owner.__mro__: - return False - return True - -def _unsloth_dpo_chunked_head( - self, - model, - concatenated_batch, - method = "concatenated_forward", -): - """``(unwrapped_model, head, (multiply, divide, softcap))`` when concatenated_forward can score from hidden - states, else ``None`` (TRL's own logits path runs). Decided before the forward, which then skips the logits.""" - reason = None - if _unsloth_masked_logprobs is None: - reason = "unsloth_zoo has no chunked_logprob" - elif os.environ.get("UNSLOTH_FUSED_LINEAR", "1") == "0": - reason = "UNSLOTH_FUSED_LINEAR=0" - elif os.environ.get("UNSLOTH_DPO_CHUNKED", "1").lower() in ("0", "false", "no", "off"): - reason = "UNSLOTH_DPO_CHUNKED=0" - elif getattr(self, "is_encoder_decoder", False): - reason = "encoder-decoder" - elif getattr(self, "aux_loss_enabled", False): - reason = "router aux loss needs the stock forward" - elif not _unsloth_stock_logps_hook( - self, - method, - ("get_batch_logps",) - if method == "concatenated_forward" - else ("get_batch_logps", "get_batch_loss_metrics", method), - ): - # KTO hands its callers logits stand-ins; DPO's outputs are exact, so a DPO wrapper is fine. - reason = "get_batch_logps / get_batch_loss_metrics overridden" - elif getattr(self, "padding_free", False): - reason = "padding_free" - elif getattr(self, "use_weighting", False): - reason = "use_weighting needs full-vocabulary log-probs" - elif getattr(self, "is_fsdp_enabled", False) or getattr(self, "is_deepspeed_enabled", False): - reason = "FSDP / DeepSpeed" - elif any( - k in concatenated_batch - for k in ("pixel_values", "pixel_attention_mask", "image_sizes", "token_type_ids") - + tuple(_DPO_VISION_KEYS) - ): - reason = "vision inputs" - if reason is None: - try: - unwrapped = self.accelerator.unwrap_model(model) - except Exception: - unwrapped = None - reason = "unwrap_model" - if reason is None: - head = _unsloth_gkd_dense_head(unwrapped) - if head is None: - reason = "output head is not a dense nn.Linear" - # The kernel reads weight / bias directly: a hooked head (accelerate, forward / backward hooks) is not - # F.linear(h, W, b), and its backward hooks would never see the gradient. - elif ( - hasattr(head, "_hf_hook") - or "forward" in vars(head) - or head._forward_hooks - or head._forward_pre_hooks - or head._backward_hooks - or getattr(head, "_backward_pre_hooks", None) - ): - reason = "output head has hooks or a patched forward" - if reason is None: - for config in _unsloth_text_configs(_unsloth_get_model_config(unwrapped)): - if getattr(config, "model_type", None) == "minicpm3": - reason = "minicpm3 scales hidden states before the head" - if reason is None and getattr(model, "find_unused_parameters", False): - if any(p is not None and p.requires_grad for p in (head.weight, head.bias)): - reason = "DDP find_unused_parameters with a trainable head" - if reason is not None: - try: - fallbacks = self.__dict__.setdefault("_unsloth_dpo_chunked_fallbacks", {}) - fallbacks[reason] = fallbacks.get(reason, 0) + 1 - except Exception: - pass - return None - return unwrapped, head, _unsloth_pref_logit_transforms(unwrapped) - -def _unsloth_dpo_model_forward(model, input_ids, model_kwargs, want_hidden_states): - if not want_hidden_states: - return model(input_ids, **model_kwargs) - prior = os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES") - os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" - try: - return model(input_ids, **model_kwargs) - finally: - if prior is None: - os.environ.pop("UNSLOTH_RETURN_HIDDEN_STATES", None) - else: - os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = prior - -def _unsloth_dpo_returns_hidden_states(self, chunked, logits): - if chunked is None: - return False - unwrapped, head, _ = chunked - hidden = _unsloth_pref_returns_hidden_states(unwrapped, logits, head.weight) - if not hidden: - fallbacks = self.__dict__.setdefault("_unsloth_dpo_chunked_fallbacks", {}) - reason = "forward returned logits, not hidden states" - fallbacks[reason] = fallbacks.get(reason, 0) + 1 - return hidden - -def _unsloth_dpo_hidden_logps(self, chunked, hidden_states, labels, loss_mask): - """``(per_token_logps [B, S], mean_logits [B, S, 1])`` where TRL computes ``selective_log_softmax(logits, labels)`` - and later reads ``logits[mask].mean()``. Rows outside ``loss_mask`` are 0 and never reach the head; the mean over - a masked row's vocab equals TRL's element mean, so the stand-in keeps the ``logits/*`` metrics exact.""" - _, head, (multiply, divide, softcap) = chunked - per_token_logps, mean_logits = _unsloth_masked_logprobs( - hidden_states.to(head.weight.device), - head.weight, - labels, - mask = loss_mask, - logit_scale_multiply = multiply, - logit_scale_divide = divide, - logit_softcapping = softcap, - lm_head_bias = head.bias, - return_mean_logits = True, - ) - try: - self._unsloth_dpo_chunked_calls = getattr(self, "_unsloth_dpo_chunked_calls", 0) + 1 - except Exception: - pass - return per_token_logps.to(labels.device), mean_logits.unsqueeze(-1).to(labels.device) - -def _unsloth_dpo_nll(chosen_logps, chosen_labels): - """``F.cross_entropy(chosen_logits, chosen_labels, ignore_index=0)`` from the per-token log-probs at the same - positions (TRL zeroes masked labels, so 0 is its ignore index).""" - keep = chosen_labels != 0 - return -(chosen_logps * keep).sum() / keep.sum() - -def _unsloth_dpo_forward(model, model_kwargs, want_hidden_states): - if not want_hidden_states: - return model(**model_kwargs) - prior = os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES") - os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" - try: - return model(**model_kwargs) - finally: - if prior is None: - os.environ.pop("UNSLOTH_RETURN_HIDDEN_STATES", None) - else: - os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = prior - -def _unsloth_dpo_loss_head( - self, - model, - inputs, - return_outputs = False, -): - """``_unsloth_dpo_chunked_head`` for TRL 0.29 - 1.14's ``_compute_loss``, which also needs the fused linear op - (entropy and argmax match) and must not hand hidden states back as ``outputs.logits``.""" - reason = None - if _unsloth_fused_linear_logprobs is None: - reason = "unsloth_zoo has no fused_linear_logprobs" - elif os.environ.get("UNSLOTH_FUSED_LINEAR", "1") == "0": - reason = "UNSLOTH_FUSED_LINEAR=0" - elif return_outputs: - reason = "return_outputs needs the logits" - if reason is not None: - try: - fallbacks = self.__dict__.setdefault("_unsloth_dpo_chunked_fallbacks", {}) - fallbacks[reason] = fallbacks.get(reason, 0) + 1 - except Exception: - pass - return None - return _unsloth_dpo_chunked_head(self, model, inputs, "_compute_loss") - -def _unsloth_dpo_loss_stats( - self, - chunked, - hidden_states, - labels, - mask, - ref = False, -): - """TRL's per-token statistics of ``logits = head(hidden_states)`` at completion positions (``mask``), the rest 0: - ``(log_probs, entropy, mean_logits [.., 1], is_top1)``, or the log-probs alone for ``ref``. ``mean_logits`` stands - in for the logits: ``x[mask].mean(-1)`` on it is TRL's vocab mean. No temperature: TRL's DPO scores at 1.0.""" - _, head, (multiply, divide, softcap) = chunked - # TRL projects the reference logits inside its no_grad forward: a trainable (or shared) head gets no gradient. - with torch.set_grad_enabled(torch.is_grad_enabled() and not ref): - out = _unsloth_fused_linear_logprobs( - hidden_states, - head.weight, - labels, - bias = head.bias, - mask = mask, - logit_scale_multiply = multiply, - logit_scale_divide = divide, - logit_softcapping = softcap, - return_entropy = not ref, - return_mean_logits = not ref, - return_is_top1 = not ref, - entropy_requires_grad = False, - ) - try: - self._unsloth_dpo_chunked_calls = getattr(self, "_unsloth_dpo_chunked_calls", 0) + 1 - except Exception: - pass - if ref: - return out - logps, entropy, mean_logits, is_top1 = out - return logps, entropy, mean_logits.unsqueeze(-1), is_top1 - -try: - from unsloth_zoo.device_map_planner import detect_logit_transforms -except Exception: - detect_logit_transforms = None -try: - from unsloth_zoo.fused_losses.chunked_logprob import ( - unsloth_masked_logprobs as _unsloth_masked_logprobs, - ) -except Exception: - _unsloth_masked_logprobs = None -try: - from unsloth_zoo.fused_losses.fused_linear import ( - fused_linear_logprobs as _unsloth_fused_linear_logprobs, - ) -except Exception: - _unsloth_fused_linear_logprobs = None -_DPO_VISION_KEYS = ('pixel_position_ids', 'image_position_ids', 'mm_token_type_ids') - -def _unsloth_fused_pass_head(self, model, inputs, return_outputs, method): - """``_unsloth_dpo_chunked_head`` for a rewritten pass that also needs the fused linear op, keeps TRL's path for - the kill switch, ``return_outputs`` (the caller reads ``outputs.logits``) and TRL's own liger / chunked head.""" - reason = None - if _unsloth_fused_linear_logprobs is None: - reason = "unsloth_zoo has no fused_linear_logprobs" - elif os.environ.get("UNSLOTH_FUSED_LINEAR", "1") == "0": - reason = "UNSLOTH_FUSED_LINEAR=0" - elif return_outputs: - reason = "return_outputs needs the logits" - elif getattr(self, "use_liger_kernel", False): - reason = "use_liger_kernel runs TRL's own chunked head" - if reason is not None: - try: - fallbacks = self.__dict__.setdefault("_unsloth_dpo_chunked_fallbacks", {}) - fallbacks[reason] = fallbacks.get(reason, 0) + 1 - except Exception: - pass - return None - return _unsloth_dpo_chunked_head(self, model, inputs, method) - -def _unsloth_fused_pass_logps( - self, - chunked, - hidden_states, - labels, - mask, - stats = False, - grad = False, -): - """TRL's ``selective_log_softmax(head(hidden_states), labels)`` with rows outside ``mask`` zeroed (TRL's - ``row_mask`` or its ``logps[mask == 0] = 0.0``). ``stats`` also returns the entropy (no grad, as TRL computes it - from detached logits) and the per-position vocab mean ``[.., 1]``: ``x[mask].mean(-1)`` on it is TRL's mean of - the logits row. A shorter ``hidden_states`` (the model truncated input_ids) is scored over the common length and - zero padded. Only the policy pass (``grad``) keeps a graph: TRL scores the reference / KL logits it built under - no_grad, so a trainable head must get no gradient from them.""" - _, head, (multiply, divide, softcap) = chunked - width = labels.shape[1] - n = min(hidden_states.shape[1], width) - with torch.set_grad_enabled(grad and torch.is_grad_enabled()): - out = _unsloth_fused_linear_logprobs( - hidden_states[:, :n], - head.weight, - labels[:, :n], - bias = head.bias, - mask = mask[:, :n].bool(), - logit_scale_multiply = multiply, - logit_scale_divide = divide, - logit_softcapping = softcap, - return_entropy = stats, - return_mean_logits = stats, - entropy_requires_grad = False, - ) - try: - self._unsloth_dpo_chunked_calls = getattr(self, "_unsloth_dpo_chunked_calls", 0) + 1 - except Exception: - pass - out = out if isinstance(out, tuple) else (out,) - out = tuple(torch.nn.functional.pad(x, (0, width - n)) if n != width else x for x in out) - if not stats: - return out[0].to(labels.device) - logps, entropy, mean_logits = out - return ( - logps.to(labels.device), - entropy.to(labels.device), - mean_logits.unsqueeze(-1).to(labels.device), - ) - -def _unsloth_fused_pass_vocab_sums(self, chunked, hidden_states): - """``[B, S, 1]`` per-position vocab sums of ``head(hidden_states)`` (no grad): TRL 1.6 - 1.8 KTO logs - ``logits.index_select(0, idx).nansum()``, a sum over every position and vocab entry, which this keeps exact.""" - _, head, (multiply, divide, softcap) = chunked - with torch.no_grad(): - hidden_states = hidden_states.detach() - _, mean_logits = _unsloth_fused_linear_logprobs( - hidden_states, - head.weight, - torch.zeros(hidden_states.shape[:-1], dtype = torch.long, device = hidden_states.device), - bias = head.bias, - logit_scale_multiply = multiply, - logit_scale_divide = divide, - logit_softcapping = softcap, - return_mean_logits = True, - ) - return (mean_logits * head.weight.shape[0]).unsqueeze(-1).to(hidden_states.device) -@dataclass -class UnslothDPOConfig(DPOConfig): - """ - - Configuration class for the [`DPOTrainer`]. - - This class includes only the parameters that are specific to DPO training. For a full list of training arguments, - please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may - differ from those in [`~transformers.TrainingArguments`]. - - Using [`~transformers.HfArgumentParser`] we can turn this class into - [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the - command line. - - Parameters: - > Parameters that control the model - - model_init_kwargs (`dict[str, Any]`, *optional*): - Keyword arguments for [`~transformers.AutoModelForCausalLM.from_pretrained`], used when the `model` - argument of the [`DPOTrainer`] is provided as a string. The `revision` value is also used when loading the - processing class. - trust_remote_code (`bool`, *optional*, defaults to `False`): - Whether to allow loading models and tokenizers that ship custom Python code from the Hub. Forwarded to - [`~transformers.AutoModelForCausalLM.from_pretrained`] and [`~transformers.AutoProcessor.from_pretrained`]. - router_aux_loss_coef (`float`, *optional*): - Coefficient of the load-balancing auxiliary loss for Mixture-of-Experts (MoE) models, added to the training - loss with this weight. When not set, the value declared by the model config is used. Set to `0.0` to - disable it. Fails when used with a non-MoE model. - disable_dropout (`bool`, *optional*, defaults to `True`): - Whether to disable dropout in the model and reference model. - - > Parameters that control the data preprocessing - - dataset_num_proc (`int`, *optional*): - Number of processes to use for processing the dataset. - max_length (`int` or `None`, *optional*, defaults to `1024`): - Maximum length of the tokenized sequence. Sequences longer than `max_length` are truncated from the left or - right depending on the `truncation_mode`. If `None`, no truncation is applied. - truncation_mode (`str`, *optional*, defaults to `"keep_start"`): - Truncation mode to use when the sequence exceeds `max_length`. The only supported value is `"keep_start"`. - The `"keep_end"` value is deprecated and will be removed in v2.0.0. - padding_free (`bool`, *optional*, defaults to `False`): - Whether to perform forward passes without padding by flattening all sequences in the batch into a single - continuous sequence. This reduces memory usage by eliminating padding overhead. Currently, this is only - supported with the FlashAttention 2 or 3, which can efficiently handle the flattened batch structure. - Temporarily unavailable: since the DPO refactor, setting it to `True` warns and falls back to standard - padding. It is planned to return in a future update. - pad_to_multiple_of (`int`, *optional*): - If set, the sequences will be padded to a multiple of this value. - precompute_ref_log_probs (`bool`, *optional*, defaults to `False`): - Whether to precompute the reference model log probabilities for the entire training dataset before - training. This allows to save memory during training, as the reference model does not need to be kept in - memory. - precompute_ref_batch_size (`int`, *optional*): - Batch size to use when precomputing reference model log probabilities. This can be set higher than the - training batch size to speed up preprocessing. If `None`, defaults to `per_device_train_batch_size` for - training and `per_device_eval_batch_size` for evaluation. - - > Parameters that control the training - - loss_type (`list[str]`, *optional*, defaults to `["sigmoid"]`): - Type of loss to use. Possible values are: `'sigmoid'`, `'hinge'`, `'ipo'`, `'exo_pair'`, `'nca_pair'`, - `'robust'`, `'bco_pair'`, `'sppo_hard'`, `'aot'`, `'aot_unpaired'`, `'apo_zero'`, `'apo_down'`, - `'discopop'`, `'sft'`, `'sigmoid_norm'`. If multiple loss types are provided, they will be combined using - the weights specified in `loss_weights`. - loss_weights (`list[float]`, *optional*): - List of loss weights for multi-loss combinations. Used when combining multiple loss types. Example: `[0.8, - 0.2, 1.0]` for MPO. If not provided, defaults to equal weights (`1.0`) for all loss types. - ld_alpha (`float`, *optional*): - α parameter from the LD-DPO paper, which controls the weighting of the verbose token log-probabilities in - responses. If `None`, no weighting is applied to the verbose part, and the loss is equivalent to the - standard DPO loss. Must be in [0.0, 1.0]: `ld_alpha=1.0` applies no weighting, and `ld_alpha=0.0` masks - tokens beyond shared lengths. - f_divergence_type (`str`, *optional*, defaults to `"reverse_kl"`): - f-divergence regularizer between policy and reference (f-DPO paper). Possible values are: `reverse_kl` - (default), `forward_kl`, `js_divergence`, `alpha_divergence`. Only the loss types built on the - chosen-rejected reward difference support a non-default value: `'sigmoid'`, `'sigmoid_norm'`, `'hinge'`, - `'ipo'`, `'exo_pair'`, `'robust'`, `'discopop'`, `'sft'`. The other loss types have no valid f-divergence - generalization and are rejected with a non-default value. - f_alpha_divergence_coef (`float`, *optional*, defaults to `0.5`): - α coefficient for the α-divergence u^-α regularizer, used only when `f_divergence_type='alpha_divergence'`. - label_smoothing (`float`, *optional*, defaults to `0.0`): - Label smoothing parameter used in Robust DPO and EXO. In Robust DPO, it is interpreted as the probability - that a preference label is flipped and must lie in [0.0, 0.5); a typical value recommended by the Robust - DPO paper is 0.1. In EXO, it corresponds to the ε label smoothing parameter, for which the paper recommends - a typical value of 1e-3. - beta (`float`, *optional*, defaults to `0.1`): - Parameter controlling the deviation from the reference model. Higher β means less deviation from the - reference model. For the IPO loss (`loss_type='ipo'`), this value is the regularization parameter denoted - by τ in the [paper](https://huggingface.co/papers/2310.12036). - use_weighting (`bool`, *optional*, defaults to `False`): - Whether to apply [WPO](https://huggingface.co/papers/2406.11827)-style weighting to preference pairs using - the policy's length-normalized sequence probabilities. - discopop_tau (`float`, *optional*, defaults to `0.05`): - τ/temperature parameter from the DiscoPOP paper, which controls the shape of the log-ratio modulated loss - when using `loss_type='discopop'`. The paper recommends the default value `discopop_tau=0.05`. - activation_offloading (`bool`, *optional*, defaults to `False`): - Whether to offload the activations to the CPU. - sync_ref_model (`bool`, *optional*, defaults to `False`): - Whether to synchronize the reference model with the active model every `ref_model_sync_steps` steps, using - the `ref_model_mixup_alpha` parameter. This synchronization originates from the - [TR-DPO](https://huggingface.co/papers/2404.09656) paper. `sync_ref_model=True` is not yet compatible with - PEFT or `precompute_ref_log_probs=True`. - ref_model_mixup_alpha (`float`, *optional*, defaults to `0.6`): - α parameter from the TR-DPO paper, which controls the mix between the current policy and the previous - reference policy during updates. The reference policy is updated according to the equation: `π_ref = α * - π_θ + (1 - α) * π_ref_prev`. To use this parameter, you must set `sync_ref_model=True`. - ref_model_sync_steps (`int`, *optional*, defaults to `512`): - τ parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which determines how - frequently the current policy is synchronized with the reference policy. To use this parameter, you must - set `sync_ref_model=True`. - - > Deprecated parameters - - pad_token: - - - - Parameter `pad_token` is deprecated and will be removed in version v2.0.0. Set `tokenizer.pad_token` - directly and pass it as `processing_class` to the trainer instead. - - - - > [!NOTE] - > These parameters have default values different from [`~transformers.TrainingArguments`]: - > - `logging_steps`: Defaults to `10` instead of `500`. - > - `gradient_checkpointing`: Defaults to `True` instead of `False`. - > - `bf16`: Defaults to `True` if `fp16` is not set, instead of `False`. - > - `learning_rate`: Defaults to `1e-6` instead of `5e-5`. - - """ - _unsloth_accepts_adamw_flash = True - vllm_sampling_params: Optional[Any] = field( - default = None, - metadata = {'help': 'vLLM SamplingParams'}, - ) - unsloth_num_chunks : Optional[int] = field( - default = -1, - metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'}, - ) - unsloth_logit_chunk_multiplier : Optional[int] = field( - default = None, - metadata = {'help': 'Multiplier for chunked logit computations.'}, - ) - unsloth_grpo_mini_batch : Optional[int] = field( - default = None, - metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'}, - ) - max_seq_length : Optional[int] = field( - default = None, - metadata = {'help': 'Maximum sequence length to truncate to.'}, - ) - context_parallel_size : Optional[int] = field( - default = 1, - metadata = {'help': 'Ranks per context parallel group (training forwards only). 1 disables it.'}, - ) - context_parallel_backend : Optional[str] = field( - default = 'ulysses', - metadata = {'help': 'Context parallel attention backend; preference and RL trainers support ulysses.'}, - ) - def __init__( - self, - output_dir = None, - per_device_train_batch_size = 4, - num_train_epochs = 3.0, - max_steps = -1, - learning_rate = 5e-05, - lr_scheduler_type = 'linear', - lr_scheduler_kwargs = None, - warmup_steps = 0.1, - optim = 'adamw_8bit', - optim_args = None, - weight_decay = 0.001, - adam_beta1 = 0.9, - adam_beta2 = 0.999, - adam_epsilon = 1e-08, - optim_target_modules = None, - gradient_accumulation_steps = 2, - average_tokens_across_devices = True, - max_grad_norm = 1.0, - label_smoothing_factor = 0.0, - bf16 = False, - fp16 = False, - bf16_full_eval = False, - fp16_full_eval = False, - tf32 = None, - gradient_checkpointing = True, - gradient_checkpointing_kwargs = None, - torch_compile = False, - torch_compile_backend = None, - torch_compile_mode = None, - use_liger_kernel = False, - liger_kernel_config = None, - use_cache = False, - neftune_noise_alpha = None, - torch_empty_cache_steps = 250, - auto_find_batch_size = False, - logging_strategy = 'steps', - logging_steps = 1, - logging_first_step = False, - log_on_each_node = True, - logging_nan_inf_filter = False, - include_num_input_tokens_seen = False, - log_level = 'passive', - log_level_replica = 'warning', - disable_tqdm = None, - report_to = 'none', - run_name = None, - project = 'huggingface', - trackio_space_id = 'trackio', - eval_strategy = 'no', - eval_steps = None, - eval_delay = 0, - per_device_eval_batch_size = 4, - prediction_loss_only = False, - eval_on_start = False, - eval_do_concat_batches = True, - eval_use_gather_object = False, - eval_accumulation_steps = 2, - batch_eval_metrics = False, - save_only_model = False, - save_strategy = 'steps', - save_steps = 500, - save_on_each_node = False, - save_total_limit = None, - enable_jit_checkpoint = False, - push_to_hub = False, - hub_token = None, - hub_private_repo = None, - hub_model_id = None, - hub_strategy = 'every_save', - hub_always_push = False, - hub_revision = None, - load_best_model_at_end = False, - metric_for_best_model = None, - greater_is_better = None, - ignore_data_skip = False, - restore_callback_states_from_checkpoint = False, - full_determinism = False, - seed = 3407, - data_seed = None, - use_cpu = False, - accelerator_config = None, - parallelism_config = None, - dataloader_drop_last = False, - dataloader_num_workers = 0, - dataloader_pin_memory = True, - dataloader_persistent_workers = False, - dataloader_prefetch_factor = None, - remove_unused_columns = True, - label_names = None, - train_sampling_strategy = 'random', - length_column_name = 'length', - ddp_find_unused_parameters = None, - ddp_bucket_cap_mb = None, - ddp_broadcast_buffers = None, - ddp_backend = None, - ddp_timeout = 1800, - fsdp = None, - fsdp_config = None, - deepspeed = None, - debug = '', - skip_memory_metrics = True, - do_train = False, - do_eval = False, - do_predict = False, - resume_from_checkpoint = None, - warmup_ratio = None, - logging_dir = None, - local_rank = -1, - model_init_kwargs = None, - trust_remote_code = False, - router_aux_loss_coef = None, - disable_dropout = True, - dataset_num_proc = None, - max_length = 1024, - truncation_mode = 'keep_start', - padding_free = None, - pad_to_multiple_of = None, - precompute_ref_log_probs = False, - precompute_ref_batch_size = None, - loss_weights = None, - ld_alpha = None, - f_divergence_type = 'reverse_kl', - f_alpha_divergence_coef = 0.5, - label_smoothing = 0.0, - beta = 0.1, - use_weighting = False, - discopop_tau = 0.05, - activation_offloading = False, - sync_ref_model = False, - ref_model_mixup_alpha = 0.6, - ref_model_sync_steps = 512, - pad_token = None, - vllm_sampling_params = None, - unsloth_num_chunks = -1, - unsloth_logit_chunk_multiplier = None, - unsloth_grpo_mini_batch = None, - max_seq_length = None,context_parallel_size = 1,context_parallel_backend = 'ulysses', - **kwargs, - ): - if learning_rate < 1e-7: print(f'Unsloth: Your learning rate of `{learning_rate}` is too small and less than 1e-7! Consider increasing it, otherwise gradient updates will be close to 0!') - if learning_rate > 1: print(f'Unsloth: Your learning rate of `{learning_rate}` is way too larger > 1! Consider decreasing it to 1e-1, otherwise gradient updates will explode!') - _unsloth_fsdp2 = os.environ.get('FSDP_VERSION', '').strip() == '2' - _unsloth_fsdp_cfg = fsdp_config - if fsdp and isinstance(_unsloth_fsdp_cfg, str) and os.path.isfile(_unsloth_fsdp_cfg): - try: - import json as _unsloth_json - with open(_unsloth_fsdp_cfg, encoding = 'utf-8') as _unsloth_f: - _unsloth_fsdp_cfg = _unsloth_json.load(_unsloth_f) - except Exception: - _unsloth_fsdp_cfg = None - if fsdp and _unsloth_fsdp_cfg is None: - _unsloth_fsdp_cfg = {} - if fsdp and isinstance(_unsloth_fsdp_cfg, dict): - _unsloth_fsdp2 = _unsloth_fsdp2 or str(_unsloth_fsdp_cfg.get('fsdp_version', _unsloth_fsdp_cfg.get('version', 1))) == '2' - _unsloth_flash_adamw = str(getattr(optim, 'value', optim)).strip().lower() == 'adamw_flash' - if _unsloth_flash_adamw: optim = 'adamw_torch' - if _unsloth_fsdp2 and str(getattr(optim, 'value', optim)) == 'adamw_8bit': - print("Unsloth: FSDP2 shards parameters as DTensors, which bitsandbytes' 8-bit optimizers cannot step. Switching optim = 'adamw_8bit' to 'adamw_torch_fused'.") - optim = 'adamw_torch_fused' - if num_train_epochs is None: - num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override - if output_dir is None and save_strategy == 'steps' and save_steps == 500: - output_dir = 'unsloth_training_checkpoints' - save_strategy = 'no' - try: - from unsloth_zoo.dataset_num_proc import get_dataset_num_proc as _unsloth_get_dataset_num_proc - except Exception: - try: - from unsloth.dataset_num_proc import get_dataset_num_proc as _unsloth_get_dataset_num_proc - except Exception: - _unsloth_get_dataset_num_proc = None - if _unsloth_get_dataset_num_proc is not None: - dataset_num_proc = _unsloth_get_dataset_num_proc(dataset_num_proc, serial_as_none = True) - if os.environ.get('UNSLOTH_ENABLE_FLEX_ATTENTION', '0') == '1': - from unsloth_zoo.flex_attention import HAS_FLEX_ATTENTION - if HAS_FLEX_ATTENTION and pad_to_multiple_of is None: - from unsloth_zoo.flex_attention import FLEX_ATTENTION_BLOCK_SIZE - pad_to_multiple_of = FLEX_ATTENTION_BLOCK_SIZE - - if eval_steps is not None and eval_strategy != 'steps': - print(f'Unsloth: `eval_steps = {eval_steps}` is ignored because `eval_strategy` is {getattr(eval_strategy, "value", eval_strategy)!r}. Set `eval_strategy = "steps"` to evaluate every `eval_steps` steps.') - - - # One dict so the filter sees the mirrored parameters AND `**kwargs`: - # filtering kwargs alone would double-bind any argument TRL renamed, - # since the new name is itself a mirrored parameter. - _unsloth_config_arguments = dict( - output_dir = output_dir, - per_device_train_batch_size = per_device_train_batch_size, - num_train_epochs = num_train_epochs, - max_steps = max_steps, - learning_rate = learning_rate, - lr_scheduler_type = lr_scheduler_type, - lr_scheduler_kwargs = lr_scheduler_kwargs, - warmup_steps = warmup_steps, - optim = optim, - optim_args = optim_args, - weight_decay = weight_decay, - adam_beta1 = adam_beta1, - adam_beta2 = adam_beta2, - adam_epsilon = adam_epsilon, - optim_target_modules = optim_target_modules, - gradient_accumulation_steps = gradient_accumulation_steps, - average_tokens_across_devices = average_tokens_across_devices, - max_grad_norm = max_grad_norm, - label_smoothing_factor = label_smoothing_factor, - bf16 = bf16, - fp16 = fp16, - bf16_full_eval = bf16_full_eval, - fp16_full_eval = fp16_full_eval, - tf32 = tf32, - gradient_checkpointing = gradient_checkpointing, - gradient_checkpointing_kwargs = gradient_checkpointing_kwargs, - torch_compile = torch_compile, - torch_compile_backend = torch_compile_backend, - torch_compile_mode = torch_compile_mode, - use_liger_kernel = use_liger_kernel, - liger_kernel_config = liger_kernel_config, - use_cache = use_cache, - neftune_noise_alpha = neftune_noise_alpha, - torch_empty_cache_steps = torch_empty_cache_steps, - auto_find_batch_size = auto_find_batch_size, - logging_strategy = logging_strategy, - logging_steps = logging_steps, - logging_first_step = logging_first_step, - log_on_each_node = log_on_each_node, - logging_nan_inf_filter = logging_nan_inf_filter, - include_num_input_tokens_seen = include_num_input_tokens_seen, - log_level = log_level, - log_level_replica = log_level_replica, - disable_tqdm = disable_tqdm, - report_to = report_to, - run_name = run_name, - project = project, - trackio_space_id = trackio_space_id, - eval_strategy = eval_strategy, - eval_steps = eval_steps, - eval_delay = eval_delay, - per_device_eval_batch_size = per_device_eval_batch_size, - prediction_loss_only = prediction_loss_only, - eval_on_start = eval_on_start, - eval_do_concat_batches = eval_do_concat_batches, - eval_use_gather_object = eval_use_gather_object, - eval_accumulation_steps = eval_accumulation_steps, - batch_eval_metrics = batch_eval_metrics, - save_only_model = save_only_model, - save_strategy = save_strategy, - save_steps = save_steps, - save_on_each_node = save_on_each_node, - save_total_limit = save_total_limit, - enable_jit_checkpoint = enable_jit_checkpoint, - push_to_hub = push_to_hub, - hub_token = hub_token, - hub_private_repo = hub_private_repo, - hub_model_id = hub_model_id, - hub_strategy = hub_strategy, - hub_always_push = hub_always_push, - hub_revision = hub_revision, - load_best_model_at_end = load_best_model_at_end, - metric_for_best_model = metric_for_best_model, - greater_is_better = greater_is_better, - ignore_data_skip = ignore_data_skip, - restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint, - full_determinism = full_determinism, - seed = seed, - data_seed = data_seed, - use_cpu = use_cpu, - accelerator_config = accelerator_config, - parallelism_config = parallelism_config, - dataloader_drop_last = dataloader_drop_last, - dataloader_num_workers = dataloader_num_workers, - dataloader_pin_memory = dataloader_pin_memory, - dataloader_persistent_workers = dataloader_persistent_workers, - dataloader_prefetch_factor = dataloader_prefetch_factor, - remove_unused_columns = remove_unused_columns, - label_names = label_names, - train_sampling_strategy = train_sampling_strategy, - length_column_name = length_column_name, - ddp_find_unused_parameters = ddp_find_unused_parameters, - ddp_bucket_cap_mb = ddp_bucket_cap_mb, - ddp_broadcast_buffers = ddp_broadcast_buffers, - ddp_backend = ddp_backend, - ddp_timeout = ddp_timeout, - fsdp = fsdp, - fsdp_config = fsdp_config, - deepspeed = deepspeed, - debug = debug, - skip_memory_metrics = skip_memory_metrics, - do_train = do_train, - do_eval = do_eval, - do_predict = do_predict, - resume_from_checkpoint = resume_from_checkpoint, - warmup_ratio = warmup_ratio, - logging_dir = logging_dir, - local_rank = local_rank, - model_init_kwargs = model_init_kwargs, - trust_remote_code = trust_remote_code, - router_aux_loss_coef = router_aux_loss_coef, - disable_dropout = disable_dropout, - dataset_num_proc = dataset_num_proc, - max_length = max_length, - truncation_mode = truncation_mode, - padding_free = padding_free, - pad_to_multiple_of = pad_to_multiple_of, - precompute_ref_log_probs = precompute_ref_log_probs, - precompute_ref_batch_size = precompute_ref_batch_size, - loss_weights = loss_weights, - ld_alpha = ld_alpha, - f_divergence_type = f_divergence_type, - f_alpha_divergence_coef = f_alpha_divergence_coef, - label_smoothing = label_smoothing, - beta = beta, - use_weighting = use_weighting, - discopop_tau = discopop_tau, - activation_offloading = activation_offloading, - sync_ref_model = sync_ref_model, - ref_model_mixup_alpha = ref_model_mixup_alpha, - ref_model_sync_steps = ref_model_sync_steps, - pad_token = pad_token,**kwargs) - super().__init__(**_unsloth_filter_config_init_kwargs(DPOConfig, _unsloth_config_arguments, mirrored_from = __class__)) - self.vllm_sampling_params = vllm_sampling_params - self.unsloth_num_chunks = unsloth_num_chunks - if unsloth_grpo_mini_batch is not None: - if self.generation_batch_size >= unsloth_grpo_mini_batch: - self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch - else: - raise ValueError( - f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, " - f"which is self.per_device_train_batch_size * gradient_accumulation_steps." - ) - self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier - self.max_seq_length = max_seq_length - self.context_parallel_size = context_parallel_size - self.context_parallel_backend = context_parallel_backend - # Unsloth: keep the reentrant checkpoint path - if getattr(self, 'gradient_checkpointing', False): - _gc_kwargs = getattr(self, 'gradient_checkpointing_kwargs', None) or {} - if _gc_kwargs.get('context_fn') is None and not _gc_kwargs.get('debug', False): - _gc_kwargs['use_reentrant'] = True - self.gradient_checkpointing_kwargs = _gc_kwargs - from unsloth.optimizers.flash_adamw import is_flash_adamw_marked, mark_flash_adamw, unmark_flash_adamw - if _unsloth_flash_adamw or is_flash_adamw_marked(self): - mark_flash_adamw(self) - else: - unmark_flash_adamw(self) - -pass - -class _UnslothDPOTrainer(_BaseTrainer): - """""" - - loss_is_scaled_for_ga = False - - _tag_names = ["trl", "dpo"] - _name = "DPO" - _paper = { - "title": "Direct Preference Optimization: Your Language Model is Secretly a Reward Model", - "id": "2305.18290", - # docstyle-ignore - "citation": textwrap.dedent("""\ - @inproceedings{rafailov2023direct, - title = {{Direct Preference Optimization: Your Language Model is Secretly a Reward Model}}, - author = {Rafael Rafailov and Archit Sharma and Eric Mitchell and Christopher D. Manning and Stefano Ermon and Chelsea Finn}, - year = 2023, - booktitle = {Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 - 16, 2023}, - url = {http://papers.nips.cc/paper_files/paper/2023/hash/a85b405ed65c6477a4fe8302b5e06ce7-Abstract-Conference.html}, - editor = {Alice Oh and Tristan Naumann and Amir Globerson and Kate Saenko and Moritz Hardt and Sergey Levine}, - }"""), - } - - def __init__( - self, - model: "str | PreTrainedModel | PeftModel", - ref_model: PreTrainedModel | None = None, - args: DPOConfig | None = None, - data_collator: DataCollator | None = None, - train_dataset: Dataset | IterableDataset | None = None, - eval_dataset: Dataset - | IterableDataset - | DatasetDict - | IterableDatasetDict - | dict[str, Dataset | IterableDataset] - | None = None, - processing_class: PreTrainedTokenizerBase | ProcessorMixin | None = None, - compute_metrics: Callable[[EvalPrediction], dict] | None = None, - callbacks: list[TrainerCallback] | None = None, - optimizers: tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None] = (None, None), - quantization_config: "BitsAndBytesConfig | None" = None, - peft_config: "PeftConfig | None" = None, - ): - # Args - if args is None: - model_name = model if isinstance(model, str) else get_config_model_id(model.config) - model_name = model_name.split("/")[-1] - args = DPOConfig(f"{model_name}-DPO") - - if train_dataset is None: - raise ValueError("`train_dataset` is required") - elif isinstance(train_dataset, IterableDataset): - # IterableDataset requires dispatch_batches=False because Accelerate's dispatch mode may try to concatenate - # batches from multiple processes, leading to mismatch errors. - if args.accelerator_config.dispatch_batches is True: - logger.warning( - "You are using an `IterableDataset` for training with `dispatch_batches=True`. `dispatch_batches` " - "is forced to `False` when using an `IterableDataset`. To remove this warning, unset " - "`dispatch_batches` in `DPOConfig` or set it to `False`." - ) - args.accelerator_config.dispatch_batches = False - elif not isinstance(train_dataset, Dataset): - raise TypeError( - f"`train_dataset` must be a `Dataset` or `IterableDataset`, got `{type(train_dataset).__name__}`." - ) - - # Model - # PEFT initializes the adapter weights randomly, so set_seed must be done before creating the model to ensure - # reproducibility. - set_seed(args.seed) - if isinstance(model, str): - model_init_kwargs = dict(args.model_init_kwargs or {}) # copy to avoid mutating model_init_kwargs - if quantization_config is not None: - if "quantization_config" in model_init_kwargs: - raise ValueError( - "You set `quantization_config` both as a trainer argument and in `args.model_init_kwargs`. " - "Please set it in only one place, preferably as a trainer argument." - ) - model_init_kwargs["quantization_config"] = quantization_config - # Distributed training requires device_map=None ["auto" fails] - if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: - model_init_kwargs["device_map"] = None - model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) - model_revision = model_init_kwargs.get("revision") - model = create_model_from_path(model, **model_init_kwargs) - else: - model_revision = None - if args.model_init_kwargs is not None: - logger.warning( - "You passed `model_init_kwargs` to the `DPOConfig`, but your model is already instantiated. " - "The `model_init_kwargs` will be ignored." - ) - if quantization_config is not None: - raise ValueError( - "You passed `quantization_config` to the trainer, but your model is already instantiated. " - "Quantization can only be applied when the model is loaded from a model identifier (`str`). " - "Either pass the model as a model identifier, or omit `quantization_config`." - ) - # Non-quantized models do not have the `is_loaded_in_{8,4}bit` attributes, whereas quantized models do - _is_quantized_model = getattr(model, "is_loaded_in_4bit", False) or getattr(model, "is_loaded_in_8bit", False) - if ref_model is model: - raise ValueError( - "`model` and `ref_model` cannot be the same object. In most cases you should omit `ref_model` and " - "we'll initialize it to a copy of `model` for you." - ) - - # Processing class - if processing_class is None: - processing_class = AutoProcessor.from_pretrained( - get_config_model_id(model.config), revision=model_revision, trust_remote_code=args.trust_remote_code - ) - - # Handle pad token for processors or tokenizers - if isinstance(processing_class, ProcessorMixin): - self._tokenizer = processing_class.tokenizer - self._is_vlm = True - elif isinstance(processing_class, PreTrainedTokenizerBase): - self._tokenizer = processing_class - self._is_vlm = False - else: - raise TypeError("The `processing_class` must be either a `PreTrainedTokenizerBase` or a `ProcessorMixin`") - - if self._tokenizer.pad_token is None: - self._tokenizer.pad_token = self._tokenizer.eos_token - - # The model must agree with the tokenizer on the pad token from construction, so mirror it onto the model - # configs. - model.config.get_text_config().pad_token_id = self._tokenizer.pad_token_id - model.generation_config.pad_token_id = self._tokenizer.pad_token_id - - # PEFT - if False: - if not is_peft_available(): - raise ImportError( - "You passed `peft_config` but the `peft` library is not installed. " - "Install it with `pip install trl[peft]`." - ) - if not isinstance(peft_config, PeftConfig): - raise TypeError( - f"`peft_config` must be a `peft.PeftConfig` instance (e.g. `peft.LoraConfig`), " - f"got {type(peft_config).__name__}." - ) - if is_peft_model(model): - raise ValueError( - "You passed a `PeftModel` instance together with a `peft_config` to the trainer. Please first merge " - "and unload the existing adapter, save the resulting base model, and then pass that base model along " - "with the new `peft_config` to the trainer." - ) - # Create PEFT model - # ZeRO-3 + PEFT for non-quantized models: - # - PEFT's default autocast_adapter_dtype=True upcasts LoRA adapter params to fp32 even when the base model is bf16. - # - ZeRO-3's _allgather_params_coalesced allocates output buffers using the dtype of the first persistent parameter, - # so mixed-dtype persistent_parameters [bf16 base + fp32 LoRA] cause a TypeError on the first optimizer step. - # - Passing autocast_adapter_dtype=False keeps adapter params in the base model dtype [bf16], fixing the mismatch. - # - This is safe: the fp32 upcast is a QLoRA-specific concern [low-bit quantized base models], not needed for - # non-quantized bf16 training. - # - See: - # - TRL issue: https://github.com/huggingface/trl/issues/6089 - # - Upstream issue: https://github.com/deepspeedai/DeepSpeed/issues/8072 - get_peft_model_kwargs = {} - if args.deepspeed_plugin is not None and args.deepspeed_plugin.zero_stage == 3 and not _is_quantized_model: - get_peft_model_kwargs["autocast_adapter_dtype"] = False - model = get_peft_model(model, peft_config, **get_peft_model_kwargs) - - elif is_peft_model(model) and ref_model is None: - # If the model is a PEFT model with a pretrained adapter, we need to create a "ref" adapter that is a copy - # of the "default" adapter, so that we can use it as the reference model during DPO training. Before PEFT - # 0.20.0, only one adapter per model was supported when the LoRA config uses `target_parameters` [see - # peft#3340, fixed in peft#3350], so in that case we skip the "ref" adapter and compute the reference log - # probs with adapters disabled, i.e. with the base model. The fix only allows adapters targeting the same - # parameters, which holds here since the "ref" adapter reuses the "default" config. `target_parameters` - # itself was only added in PEFT 0.17.0, so the version check is bounded on both sides. - default_config = model.peft_config["default"] - if ( - isinstance(default_config, LoraConfig) - and Version("0.17.0") <= Version(peft.__version__) < Version("0.20.0") - and default_config.target_parameters - ): - logger.warning( - "PEFT<0.20.0 can't add a frozen reference adapter alongside one that uses `target_parameters` " - "(peft#3340], so the reference log probs are computed from the base model [adapters disabled]. " - "Upgrade to `peft>=0.20.0` to train against a copy of your adapter instead. If you wrapped the " - "model only to apply LoRA, pass a `peft_config` to the trainer instead; if you wrapped it " - "deliberately (pretrained adapter or custom init), note that the base model matches your adapter " - "only when it's freshly zero-initialized. If it is, this warning is safe to ignore." - ) - else: - model.add_adapter("ref", default_config) - for name, param in model.named_parameters(): - if ".default." in name: - ref_name = name.replace(".default.", ".ref.") - ref_param = model.get_parameter(ref_name) - ref_param.data.copy_(param.data) - - # PEFT + DeepSpeed ZeRO-3 requires reentrant checkpointing. For more details, see - # https://github.com/huggingface/trl/issues/2514#issuecomment-2692152703. - # Can be removed once https://github.com/deepspeedai/DeepSpeed/pull/8130 is merged and released. - if ( - is_peft_model(model) - and args.deepspeed_plugin is not None - and args.deepspeed_plugin.zero_stage == 3 - and args.gradient_checkpointing - ): - args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} - use_reentrant = args.gradient_checkpointing_kwargs.get("use_reentrant") - if use_reentrant is False: - logger.warning( - "You are using PEFT with DeepSpeed ZeRO-3 and gradient checkpointing with `use_reentrant=False`. " - "`use_reentrant` is forced to `True` in this configuration to ensure correct training. To remove " - "this warning, unset `use_reentrant` in `gradient_checkpointing_kwargs` or set it to `True`." - ) - args.gradient_checkpointing_kwargs["use_reentrant"] = True - - # When using gradient checkpointing with PEFT, we need to enable input gradients. transformers.Trainer normally - # handles this, but a bug currently prevents it; see https://github.com/huggingface/transformers/issues/42489 - if is_peft_model(model) and args.gradient_checkpointing: - model.enable_input_require_grads() - - # When using QLoRA, the PEFT adapter weights are converted to bf16 to follow the recommendations from the - # original paper [see https://huggingface.co/papers/2305.14314, paragraph 3]. Normally, this can be done by - # passing `autocast_adapter_dtype=False` to `get_peft_model`, but this option is not yet supported for - # quantized models. See: https://github.com/huggingface/peft/issues/2889 - # The DoRA magnitude vector is excluded: unlike LoRA A/B, its optimizer updates can be smaller than bf16 can - # represent, silently freezing it, see #7268. - if False: - for name, param in model.named_parameters(): - if param.requires_grad and "lora_magnitude_vector" not in name: - param.data = param.data.to(torch.bfloat16) - - # Data collator - self.padding_free = args.padding_free - if self.padding_free: - logger.warning( - "`padding_free=True` is temporarily unavailable after a refactor and is currently disabled. Falling " - "back to standard padding (`padding_free=False`). This feature is planned to return in a future " - "update; for now, please set `padding_free=False` explicitly." - ) - self.padding_free = False - dataset_sample = next(iter(train_dataset)) - self._is_vision_dataset = "image" in dataset_sample or "images" in dataset_sample - if self._is_vision_dataset and not self._is_vlm: - raise ValueError( - "The dataset appears to be vision-related (contains 'image' or 'images' keys), but the provided " - "model does not seem to be a vision-language model. Please check your model and dataset." - ) - if self._is_vision_dataset and args.max_length is not None and args.truncation_mode == "keep_end": - raise ValueError( - "truncation_mode='keep_end' is not supported for vision-language models. Image tokens reside " - "inside the prompt portion of the sequence; depending on the example, keep_end may silently " - "drop them, causing pixel_values to be forwarded to the model with no corresponding visual " - "tokens in input_ids. Use truncation_mode='keep_start' (the default) or set max_length=None." - ) - if self._is_vision_dataset and args.precompute_ref_log_probs: - raise ValueError( - "`precompute_ref_log_probs=True` is not supported for vision datasets. For vision-language " - "models, all data processing is performed on the fly rather than upfront, and running a full " - "forward pass of the reference model over the entire dataset is not supported for large " - "multimodal models. Set `precompute_ref_log_probs=False`." - ) - if data_collator is None: - # Get the pad token: if not provided, use the one from the processing class or the eos token - # if the processing class does not have a pad token. - pad_token = args.pad_token or self._tokenizer.pad_token or self._tokenizer.eos_token - if pad_token not in self._tokenizer.get_vocab(): - raise ValueError( - f"The specified `pad_token` ('{pad_token}') is not found in the vocabulary of the given " - f"`processing_class` ({processing_class.__class__.__name__}). Ensure that the `pad_token` exists " - "in the vocabulary before using it as a padding token." - ) - self._tokenizer.pad_token = pad_token - # The model must agree with the tokenizer on the pad token from construction, so mirror it onto the model - # configs. - model.config.get_text_config().pad_token_id = self._tokenizer.pad_token_id - model.generation_config.pad_token_id = self._tokenizer.pad_token_id - if self._is_vision_dataset: - data_collator = DataCollatorForVisionPreference( - processor=processing_class, - max_length=args.max_length, - pad_to_multiple_of=args.pad_to_multiple_of, - ) - else: - data_collator = DataCollatorForPreference( - pad_token_id=self._tokenizer.pad_token_id, - max_length=args.max_length, - truncation_mode=args.truncation_mode, - pad_to_multiple_of=args.pad_to_multiple_of, - ) - - # Training arguments - self.beta = args.beta - self.precompute_ref_logps = args.precompute_ref_log_probs - self.loss_types = args.loss_type # args.loss_type is already a list - self.loss_weights = args.loss_weights or [1.0] * len(self.loss_types) - self.ld_alpha = args.ld_alpha - self.f_divergence_type = args.f_divergence_type - self.f_alpha_divergence_coef = args.f_alpha_divergence_coef - # The f-DPO reparameterization leaves a β·log Z[x] term that only cancels in the chosen-rejected reward - # difference, so the other losses have no valid f-divergence generalization. - f_divergence_loss_types = {"sigmoid", "sigmoid_norm", "hinge", "ipo", "exo_pair", "robust", "discopop", "sft"} - if self.f_divergence_type != "reverse_kl" and not set(self.loss_types) <= f_divergence_loss_types: - raise ValueError( - f"`f_divergence_type='{self.f_divergence_type}'` is only supported for the following loss types: " - f"{sorted(f_divergence_loss_types)}. You provided {self.loss_types}. Use the default " - "`f_divergence_type='reverse_kl'` with these losses." - ) - self.label_smoothing = args.label_smoothing - self.use_weighting = args.use_weighting - if self.use_weighting and any(loss_type in {"aot", "aot_unpaired"} for loss_type in self.loss_types): - raise NotImplementedError( - "WPO-style weighting is not implemented for 'aot' or 'aot_unpaired' because those losses sort " - "samples, which would misalign per-pair weights." - ) - if "robust" in self.loss_types and not (0.0 <= self.label_smoothing < 0.5): - logger.warning( - "The `label_smoothing` parameter should lie in [0.0, 0.5) for the 'robust' loss. You provided " - f"{self.label_smoothing}." - ) - if "exo_pair" in self.loss_types and self.label_smoothing == 0.0: - raise ValueError( - "Label smoothing must be greater than 0.0 when using 'exo_pair' loss. The EXO paper recommends a " - "value of 1e-3." - ) - - if is_peft_model(model) and isinstance(model.get_output_embeddings(), BaseTunerLayer): - # The log-probabilities are computed by multiplying the hidden states by `lm_head.weight` directly, so an - # adapter on the LM head would be ignored and never trained. - raise ValueError( - "Applying a PEFT adapter to `lm_head` is not supported: the log-probabilities are computed from " - "`lm_head.weight` directly, so the adapter would never be trained. Remove `'lm_head'` from your " - "`target_modules`." - ) - - # Dataset - # Skip dataset preparation if it's a VLM, where preprocessing [e.g., image-to-pixel conversion] is too costly - # and done on the fly instead. - skip_prepare_dataset = self._is_vision_dataset - if not skip_prepare_dataset: - train_dataset = self._prepare_dataset(train_dataset, processing_class, args, "train") - if eval_dataset is not None: - if isinstance(eval_dataset, dict): - eval_dataset = { - key: self._prepare_dataset(dataset, processing_class, args, key) - for key, dataset in eval_dataset.items() - } - else: - eval_dataset = self._prepare_dataset(eval_dataset, processing_class, args, "eval") - - # Transformers explicitly set use_reentrant=True in the past to silence a PyTorch warning, but the default was - # never updated once PyTorch switched to recommending use_reentrant=False. Until that change lands upstream - # [see https://github.com/huggingface/transformers/pull/43203] and is released [most likely in 5.0.0], we - # default to the recommended non-reentrant behavior here, while preserving any user-provided value. - if args.gradient_checkpointing and Version(transformers.__version__) < Version("5.0.0"): - args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} - args.gradient_checkpointing_kwargs.setdefault("use_reentrant", False) - - super().__init__( - model=model, - args=args, - data_collator=data_collator, - train_dataset=train_dataset, - eval_dataset=eval_dataset, - processing_class=processing_class, - compute_metrics=compute_metrics, - callbacks=callbacks, - optimizers=optimizers, - ) - - # With several GPUs visible and no distributed launcher, `Trainer` wraps the model in `nn.DataParallel`, whose - # replicas would all run the forward bound to the original model. Checked after `Trainer.__init__`, which sets - # `n_gpu` to 1 for a model split across devices with `device_map`. - if args.n_gpu > 1: - raise ValueError( - f"{type(self).__name__} does not support `nn.DataParallel`, which `Trainer` uses when several GPUs are " - "visible to a single process. Launch the script with `accelerate launch` or `torchrun`, or make a " - "single GPU visible with `CUDA_VISIBLE_DEVICES`." - ) - - # Initialize activation offloading context - if self.args.activation_offloading: - self.maybe_activation_offload_context = get_act_offloading_ctx_manager(model=self.model) - else: - self.maybe_activation_offload_context = contextlib.nullcontext() - - # MoE load-balancing auxiliary loss. `output_router_logits` in the config means the model returns its router - # logits; `router_aux_loss_coef` is the architecture's own coefficient, 0.0 when the config declares none. - text_config = model.config.get_text_config() - coef = args.router_aux_loss_coef - self.router_aux_loss_coef = getattr(text_config, "router_aux_loss_coef", 0.0) if coef is None else coef - if coef and not hasattr(text_config, "output_router_logits"): - raise ValueError( - f"`router_aux_loss_coef` is set to {coef} but {type(model).__name__} is not a Mixture-of-Experts model " - f"that returns its router logits, so there is no auxiliary loss to weight." - ) - self.aux_loss_enabled = hasattr(text_config, "output_router_logits") and self.router_aux_loss_coef != 0.0 - - # Reference model - if ref_model is None: - if is_peft_model(self.model) or (args.precompute_ref_log_probs and not self.is_fsdp_enabled): - # If PEFT is used, the reference model is not needed since the adapter can be disabled to revert to the - # initial model. If precompute_ref_log_probs is True, the reference model does not need to be kept in - # memory during training. Under FSDP, the policy isn't sharded until `train`, so the precompute still - # loads a reference model, freed once it is done. - self.ref_model = None - else: - ref_model_init_kwargs = dict(args.model_init_kwargs or {}) # copy to avoid mutating model_init_kwargs - if quantization_config is not None: - ref_model_init_kwargs["quantization_config"] = quantization_config - # Distributed training requires device_map=None ["auto" fails] - if self.args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: - ref_model_init_kwargs["device_map"] = None - ref_model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) - ref_model_path = get_config_model_id(self.model.config) - self.ref_model = create_model_from_path(ref_model_path, **ref_model_init_kwargs) - else: - self.ref_model = ref_model - - # Disable dropout in the models - if args.disable_dropout: - disable_dropout_in_model(model) - if self.ref_model is not None: - disable_dropout_in_model(self.ref_model) - - # Liger's fused linear cross-entropy replaces `model.forward` when training starts, which would drop the fused - # LM head, so only its layer kernels are applied - if args.use_liger_kernel: - warnings.warn( - "`use_liger_kernel=True` is deprecated and will be removed in v2.0.0. Use the Hub kernels instead, " - 'with `model_init_kwargs={"use_kernels": True}`.', - FutureWarning, - stacklevel=2, - ) - liger_kernel_config = args.liger_kernel_config or {} - if liger_kernel_config.get("fused_linear_cross_entropy"): - raise ValueError( - '`liger_kernel_config={"fused_linear_cross_entropy": True}` is not supported: it replaces the ' - "model's forward, which this trainer patches with a fused LM head." - ) - args.liger_kernel_config = {**liger_kernel_config, "fused_linear_cross_entropy": False} - - # Compute the per-token log-probabilities in chunks, without materializing the full logits - # `mean_logits` and `is_top1` feed the `logits/*` and `mean_token_accuracy` metrics, `log_sum_sq_probs` the WPO - # weights - outputs = ("log_probs", "entropy", "mean_logits", "is_top1") - if self.use_weighting: - outputs += ("log_sum_sq_probs",) - add_fused_lm_head(self.model.get_base_model() if is_peft_model(self.model) else self.model, outputs=outputs) - if self.ref_model is not None: - add_fused_lm_head(self.ref_model) - - # Initialize the metrics - self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} - self._total_train_tokens = 0 - - # Tensor parallel ranks are given the same batch, so a token count gathered across all processes repeats - # every token once per rank in the group. Context and sequence parallelism shard the batch before the loss is - # computed, so they need no such correction. `ParallelismConfig.tp_size` requires accelerate 1.10.0. - if Version(accelerate.__version__) >= Version("1.10.0") and self.accelerator.parallelism_config is not None: - self._tp_size = self.accelerator.parallelism_config.tp_size - else: - self._tp_size = 1 - - # Add tags to the model - self.model.add_model_tags(self._tag_names) - - # Hash before the reference model is sharded, since DTensor parameters can't be hashed - self._precompute_model_hash = ( - hash_module(self.ref_model or self.model) if args.precompute_ref_log_probs else None - ) - - if self.ref_model is not None: - if self.is_deepspeed_enabled: - self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator) - elif self.is_fsdp_enabled: - self.ref_model = prepare_fsdp(self.ref_model, self.accelerator) - else: - self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True) - - if args.sync_ref_model: - if is_peft_model(self.model): - raise NotImplementedError( - "You passed `sync_ref_model=True` while using a PEFT model, which is currently not supported. " - "With PEFT, DPOTrainer does not keep a separate reference model in memory; instead, it recovers " - "reference behavior by temporarily disabling the adapter. As a result, there is no standalone " - "`ref_model` instance to synchronize. Use `sync_ref_model=False`, or opt for full fine-tuning if " - "you need a synced reference model. If you need `sync_ref_model` to work with PEFT, please open a " - "feature request at https://github.com/huggingface/trl/issues." - ) - if args.precompute_ref_log_probs: - raise ValueError( - "You cannot use `sync_ref_model=True` together with `precompute_ref_log_probs=True`. " - "`precompute_ref_log_probs=True` assumes a fixed reference model, but with `sync_ref_model=True` " - "the reference model is periodically updated during training, making any precomputed reference " - "log-probs stale. Set `precompute_ref_log_probs=False` or disable `sync_ref_model`." - ) - self.add_callback(SyncRefModelCallback(ref_model=self.ref_model, accelerator=self.accelerator)) - - # Reference forwards during precompute reuse a single DeepSpeed inference engine [see `_precompute_ref_logps`]. - self._precompute_engine = None - if args.precompute_ref_log_probs: - self.train_dataset = self._precompute_ref_logps( - self.train_dataset, - "train", - self.args.precompute_ref_batch_size or self.args.per_device_train_batch_size, - ) - if self.eval_dataset is not None: - if isinstance(self.eval_dataset, dict): - self.eval_dataset = { - name: self._precompute_ref_logps( - dataset, name, self.args.precompute_ref_batch_size or self.args.per_device_eval_batch_size - ) - for name, dataset in self.eval_dataset.items() - } - else: - self.eval_dataset = self._precompute_ref_logps( - self.eval_dataset, - "eval", - self.args.precompute_ref_batch_size or self.args.per_device_eval_batch_size, - ) - if self.is_fsdp_enabled and ref_model is None: - self.ref_model = None - - def _prepare_dataset( - self, - dataset: Dataset | IterableDataset, - processing_class: PreTrainedTokenizerBase | ProcessorMixin, - args: DPOConfig, - dataset_name: str, - ) -> Dataset | IterableDataset: - # Build the kwargs for the `map` function - map_kwargs = {} - if isinstance(dataset, Dataset): # IterableDataset does not support num_proc - map_kwargs["num_proc"] = args.dataset_num_proc - - with global_then_local_main_first(): - # Extract the prompt if needed - first_example = next(iter(dataset)) - if "prompt" not in first_example: - if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` - map_kwargs["desc"] = f"Extracting prompt from {dataset_name} dataset" - dataset = dataset.map(extract_prompt, **map_kwargs) - - # Add EOS token if needed: non-conversational only - first_example = next(iter(dataset)) - if not is_conversational(first_example): - if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` - map_kwargs["desc"] = f"Adding EOS to {dataset_name} dataset" - - def add_eos(example, eos_token): - if not example["chosen"].endswith(eos_token): - example["chosen"] = example["chosen"] + eos_token - if not example["rejected"].endswith(eos_token): - example["rejected"] = example["rejected"] + eos_token - return example - - dataset = dataset.map(add_eos, fn_kwargs={"eos_token": self._tokenizer.eos_token}, **map_kwargs) - - # Tokenize the dataset - if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` - map_kwargs["desc"] = f"Tokenizing {dataset_name} dataset" - - def tokenize_fn(example, processing_class): - tools = example.get("tools") - tools = json.loads(tools) if isinstance(tools, str) else tools - apply_chat_template_kwargs = {"tools": tools, **example.get("chat_template_kwargs", {})} - output = {} - if is_conversational(example): - prompt_ids = _tokenize( - processing_class, - example["prompt"], - add_generation_prompt=True, - **apply_chat_template_kwargs, - )["input_ids"] - prompt_chosen_ids = _tokenize( - processing_class, - example["prompt"] + example["chosen"], - **apply_chat_template_kwargs, - )["input_ids"] - prompt_rejected_ids = _tokenize( - processing_class, - example["prompt"] + example["rejected"], - **apply_chat_template_kwargs, - )["input_ids"] - else: - prompt_ids = _tokenize(processing_class, example["prompt"])["input_ids"] - prompt_chosen_ids = _tokenize(processing_class, example["prompt"] + example["chosen"])["input_ids"] - prompt_rejected_ids = _tokenize(processing_class, example["prompt"] + example["rejected"])[ - "input_ids" - ] - - # The completion starts where the tokenized prompt and prompt+completion diverge, which is not always - # after the prompt (see `common_prefix_length`) - prompt_len = min( - common_prefix_length(prompt_ids, prompt_chosen_ids), - common_prefix_length(prompt_ids, prompt_rejected_ids), - ) - if prompt_len < len(prompt_ids): - logger.warning_once( - "The tokenized prompt is not a prefix of the tokenized prompt+completion, usually because the " - "chat template renders the prompt alone differently or because tokens merge across the " - "boundary. The completion starts where they diverge, and this end of the prompt is left out " - "of the training sequence: " - f"{processing_class.decode(prompt_ids[prompt_len:])!r}" - ". The model is trained on a context that differs from the one it sees at inference. To avoid " - "it, use a chat template that renders the prompt the same way in both cases, or end the " - "prompt on a token boundary. This warning is shown once, but it likely applies to every " - "example in the dataset." - ) - - output["prompt_ids"] = prompt_ids[:prompt_len] - output["chosen_ids"] = prompt_chosen_ids[prompt_len:] - output["rejected_ids"] = prompt_rejected_ids[prompt_len:] - return output - - dataset = dataset.map( - tokenize_fn, - fn_kwargs={"processing_class": processing_class}, - **map_kwargs, - ) - - # Drop examples whose prompt alone fills `max_length`: with `keep_start` truncation the collator would - # remove every completion token, leaving no learning signal. `keep_end` keeps the completion end, so - # nothing is dropped there. - if args.max_length is not None and args.truncation_mode == "keep_start": - if isinstance(dataset, Dataset): # `IterableDataset.filter` does not support `desc` - map_kwargs["desc"] = f"Dropping fully truncated examples from {dataset_name} dataset" - dataset = dataset.filter(lambda example: len(example["prompt_ids"]) < args.max_length, **map_kwargs) - - return dataset - - def _set_signature_columns_if_needed(self): - # If `self.args.remove_unused_columns` is True, non-signature columns are removed. - # By default, this method sets `self._signature_columns` to the model's expected inputs (usually, "input_ids" - # and "attention_mask"). - if self._signature_columns is None: - if self._is_vision_dataset: - self._signature_columns = [ - "prompt", - "chosen", - "rejected", - "image", - "images", - "tools", - "chat_template_kwargs", - ] - else: - self._signature_columns = [ - "prompt_ids", - "chosen_ids", - "rejected_ids", - "ref_chosen_logps", - "ref_rejected_logps", - ] - - def _precompute_ref_logps(self, dataset: Dataset, name: str, batch_size: int) -> Dataset: - if isinstance(dataset, IterableDataset): - raise ValueError( - "`precompute_ref_log_probs=True` is not supported with IterableDataset. Please use a map-style " - "Dataset or set `precompute_ref_log_probs=False`." - ) - model_hash = self._precompute_model_hash - # Both inputs are rank-dependent under distributed training (ZeRO-3 shards the model), so broadcast rank 0's - # value so all ranks share one cache file. - fingerprint = [Hasher.hash((dataset._fingerprint, model_hash))] - broadcast_object_list(fingerprint, from_process=0) - fingerprint = fingerprint[0] - cache_file = dataset._get_cache_file_path(fingerprint) - - if os.path.exists(cache_file): - return concatenate_datasets([dataset, Dataset.from_file(cache_file)], axis=1) - - dataloader = DataLoader( - dataset, - batch_size=batch_size, - collate_fn=self.data_collator, - num_workers=self.args.dataloader_num_workers, - pin_memory=self.args.dataloader_pin_memory, - shuffle=False, - ) - data_loader = self.accelerator.prepare(dataloader) - - # This runs before the parent class prepares the model in `train`, so with DeepSpeed the parameters are still - # on CPU (ZeRO-1/2) and sharded (ZeRO-3). Wrap the model in an inference engine to place and gather them. Build - # it once and reuse it across precompute passes (train, eval, and later `evaluate` calls) - if self.ref_model is None and self.is_deepspeed_enabled: - if self._precompute_engine is None: - self._precompute_engine = prepare_deepspeed(self.model, self.accelerator) - model = self._precompute_engine - else: - model = self.ref_model or self.model - - ref_chosen_logps = [] - ref_rejected_logps = [] - for padded_batch in tqdm( - iterable=data_loader, - desc=f"Computing reference log probs for {name} dataset", - disable=bool(os.environ.get("TQDM_DISABLE", "")), - ): - ref_chosen_logp, ref_rejected_logp = self.compute_ref_log_probs(model, padded_batch) - ref_chosen_logp, ref_rejected_logp = self.accelerator.gather_for_metrics( - (ref_chosen_logp, ref_rejected_logp) - ) - ref_chosen_logps.append(ref_chosen_logp.cpu()) - ref_rejected_logps.append(ref_rejected_logp.cpu()) - - ref_chosen_logps = torch.cat(ref_chosen_logps) - ref_rejected_logps = torch.cat(ref_rejected_logps) - - if self.accelerator.is_main_process: - - def add_ref_logps(batch, indices): - return { - "ref_chosen_logps": ref_chosen_logps[indices], - "ref_rejected_logps": ref_rejected_logps[indices], - } - - dataset.map( - add_ref_logps, - with_indices=True, - batched=True, - remove_columns=dataset.column_names, - new_fingerprint=fingerprint, - cache_file_name=cache_file, - desc=f"Caching reference log probs for {name} dataset", - ) - self.accelerator.wait_for_everyone() - - return concatenate_datasets([dataset, Dataset.from_file(cache_file)], axis=1) - - def compute_ref_log_probs(self, model, inputs): - """Computes reference log probabilities for a single padded batch.""" - device = self.accelerator.device - _non_model_keys = {"completion_mask", "ref_chosen_logps", "ref_rejected_logps"} - model_kwargs = {k: v for k, v in inputs.items() if k not in _non_model_keys} - model_kwargs["use_cache"] = False - labels = inputs["input_ids"].masked_fill(inputs["completion_mask"] == 0, -100) - - adapter_context = contextlib.nullcontext() - if self.ref_model is None and is_peft_model(self.model): - unwrapped_model = self.accelerator.unwrap_model(self.model) - adapter_context = use_adapter( - unwrapped_model, adapter_name="ref" if "ref" in unwrapped_model.peft_config else None - ) - - with ( - torch.no_grad(), - disable_gradient_checkpointing(self.model, self.args.gradient_checkpointing_kwargs), - adapter_context, - self.accelerator.autocast(), - ): - ref_per_token_logps = model(**model_kwargs, labels=labels, fused_lm_head=True).log_probs - - shift_completion_mask = inputs["completion_mask"][..., 1:] - # Prompt-learning PEFT prepends virtual tokens to the outputs; keep the positions of the real tokens - ref_per_token_logps = ref_per_token_logps[:, -shift_completion_mask.size(1) :] - if self.ld_alpha is None: - ref_logps = ref_per_token_logps.sum(dim=1) - else: - comp_pos = shift_completion_mask.cumsum(dim=1) - comp_lens = shift_completion_mask.sum(dim=1).long() - chosen_lens, rejected_lens = comp_lens.chunk(2, dim=0) - shared_lens = torch.minimum(chosen_lens, rejected_lens) - shared_lens = torch.cat([shared_lens, shared_lens], dim=0).to(device) - shared_mask = (comp_pos > 0) & (comp_pos <= shared_lens.unsqueeze(1)) - tail_mask = comp_pos > shared_lens.unsqueeze(1) - shared_logps = (ref_per_token_logps * shared_mask).sum(dim=1) - tail_logps = (ref_per_token_logps * tail_mask).sum(dim=1) - ref_logps = shared_logps + self.ld_alpha * tail_logps - - return ref_logps.chunk(2, dim=0) - - def _compute_loss(self, model, inputs, return_outputs): - mode = "train" if self.model.training else "eval" - device = self.accelerator.device - - _non_model_keys = {"completion_mask", "ref_chosen_logps", "ref_rejected_logps"} - model_kwargs = {k: v for k, v in inputs.items() if k not in _non_model_keys} - model_kwargs["use_cache"] = False - # MoE models: request router logits so the model returns `outputs.aux_loss`. VLM wrappers honor this only - # as a forward kwarg (not from the model config), so it must be passed here. - if self.aux_loss_enabled: - model_kwargs["output_router_logits"] = True - completion_mask = inputs["completion_mask"] - shift_completion_mask = completion_mask[..., 1:] - labels = inputs["input_ids"].masked_fill(completion_mask == 0, -100) - outputs = model(**model_kwargs, labels=labels, fused_lm_head=True) - # Prompt-learning PEFT prepends virtual tokens to the outputs; keep the positions of the real tokens - seq_len = shift_completion_mask.size(1) - per_token_logps = outputs.log_probs[:, -seq_len:] - if self.ld_alpha is None: - logps = per_token_logps.sum(dim=1) # sum over sequence length - else: - comp_pos = shift_completion_mask.cumsum(dim=1) - comp_lens = shift_completion_mask.sum(dim=1).long() - chosen_lens, rejected_lens = comp_lens.chunk(2, dim=0) - shared_lens = torch.minimum(chosen_lens, rejected_lens) - shared_lens = torch.cat([shared_lens, shared_lens], dim=0).to(device) - shared_mask = (comp_pos > 0) & (comp_pos <= shared_lens.unsqueeze(1)) # shared: 1 <= pos <= shared_len - tail_mask = comp_pos > shared_lens.unsqueeze(1) # tail: pos > shared_len - shared_logps = (per_token_logps * shared_mask).sum(dim=1) - tail_logps = (per_token_logps * tail_mask).sum(dim=1) - logps = shared_logps + self.ld_alpha * tail_logps - chosen_logps, rejected_logps = logps.chunk(2, dim=0) # batch is [chosen, rejected] - - if self.precompute_ref_logps: - ref_chosen_logps, ref_rejected_logps = inputs["ref_chosen_logps"], inputs["ref_rejected_logps"] - else: - # The reference forward only needs logits for log-probs. Drop `output_router_logits` so the frozen - # reference model does not materialize router logits and compute a discarded MoE aux loss. - ref_model_kwargs = {k: v for k, v in model_kwargs.items() if k != "output_router_logits"} - # When gradient checkpointing is enabled with use_reentrant=True (non default), calling the model inside a - # torch.no_grad() block triggers a harmless PyTorch warning ("None of the inputs have requires_grad=True"). - # Temporarily disable checkpointing to avoid this warning during inference. - with torch.no_grad(), disable_gradient_checkpointing(self.model, self.args.gradient_checkpointing_kwargs): - if is_peft_model(self.model) and self.ref_model is None: - # When training a PEFT adapter, how we obtain the reference depends on the setup: - # - New adapter: disabling adapters yields the base model. - # - Re-training an existing adapter: an initial copy is loaded under the name "ref". - ref_model_unwrapped = self.accelerator.unwrap_model(self.model) - with use_adapter( - ref_model_unwrapped, adapter_name="ref" if "ref" in ref_model_unwrapped.peft_config else None - ): - ref_outputs = self.model(**ref_model_kwargs, labels=labels, fused_lm_head=True) - else: - ref_outputs = self.ref_model(**ref_model_kwargs, labels=labels, fused_lm_head=True) - ref_per_token_logps = ref_outputs.log_probs[:, -seq_len:] - if self.ld_alpha is None: - ref_logps = ref_per_token_logps.sum(dim=1) # sum over sequence length - else: - # reuse comp_pos/shared_mask/tail_mask computed above (they depend only on completion_mask) - ref_shared_logps = (ref_per_token_logps * shared_mask).sum(dim=1) - ref_tail_logps = (ref_per_token_logps * tail_mask).sum(dim=1) - ref_logps = ref_shared_logps + self.ld_alpha * ref_tail_logps - ref_chosen_logps, ref_rejected_logps = ref_logps.chunk(2, dim=0) # batch is [chosen, rejected] - - # Get the log ratios for the chosen and rejected responses - chosen_logratios = chosen_logps - ref_chosen_logps - rejected_logratios = rejected_logps - ref_rejected_logps - - if self.f_divergence_type == "reverse_kl": # standard DPO - chosen_scores = chosen_logratios - rejected_scores = rejected_logratios - elif self.f_divergence_type == "forward_kl": - # f'(t) = 1 - 1/t - chosen_scores = 1 - torch.exp(-chosen_logratios) - rejected_scores = 1 - torch.exp(-rejected_logratios) - elif self.f_divergence_type == "js_divergence": - # f'(t) = log(2t/(t+1)) = log 2 + logsigmoid(log t) - chosen_scores = math.log(2) + F.logsigmoid(chosen_logratios) - rejected_scores = math.log(2) + F.logsigmoid(rejected_logratios) - elif self.f_divergence_type == "alpha_divergence": - # alpha-divergence: f'(t) = (t^(α-1) - 1)/(α-1) - if abs(self.f_alpha_divergence_coef - 1.0) < 1e-6: # limit case f'(t) -> log(t), fall back to reverse_kl - chosen_scores = chosen_logratios - rejected_scores = rejected_logratios - else: - coef = 1.0 / (self.f_alpha_divergence_coef - 1.0) - t_chosen = (self.f_alpha_divergence_coef - 1.0) * chosen_logratios - t_rejected = (self.f_alpha_divergence_coef - 1.0) * rejected_logratios - dtype = t_chosen.dtype - # Clamp max so exp(.) stays representable after casting back - clamp_max = {torch.float16: 11.0, torch.bfloat16: 80.0, torch.float32: 80.0}[dtype] - t_chosen_float = torch.clamp(t_chosen.float(), max=clamp_max) - t_rejected_float = torch.clamp(t_rejected.float(), max=clamp_max) - chosen_scores = (torch.exp(t_chosen_float) - 1.0).to(dtype) * coef - rejected_scores = (torch.exp(t_rejected_float) - 1.0).to(dtype) * coef - else: - raise ValueError(f"Unknown f_divergence_type: {self.f_divergence_type}") - - delta_score = chosen_scores - rejected_scores - - loss = 0.0 - for loss_type, loss_weight in zip(self.loss_types, self.loss_weights, strict=True): - if loss_type == "sigmoid": - per_sequence_loss = -F.logsigmoid(self.beta * delta_score) - - elif loss_type == "hinge": - per_sequence_loss = torch.relu(1 - self.beta * delta_score) - - elif loss_type == "ipo": - # IPO uses sequence-level log-prob differences; in code these are token-summed over the completion, - # which makes the squared loss scale with completion length. We therefore normalize by the number of - # completion tokens (average per token) to make β/loss comparable across variable lengths. This length - # normalization is not explicitly discussed in the IPO paper; we confirmed this choice with the IPO - # authors, and the results reported in the paper correspond to this normalized form. - chosen_mask, rejected_mask = completion_mask.chunk(2, dim=0) - chosen_avg_score = chosen_scores / chosen_mask.sum(dim=1).clamp(min=1.0) - rejected_avg_score = rejected_scores / rejected_mask.sum(dim=1).clamp(min=1.0) - ipo_delta = chosen_avg_score - rejected_avg_score - # (Eq. 17) of the paper where beta is the regularization parameter for the IPO loss, denoted by τ. - per_sequence_loss = (ipo_delta - 1 / (2 * self.beta)) ** 2 - - elif loss_type == "exo_pair": - # Implements EXO-pref from the paper https://huggingface.co/papers/2402.00856 (Eq. 16) - # Minimize KL(p_fθ || p_rh) for K=2; p_fθ = softmax(βπ * (log πθ − log π_ref)) over {chosen, rejected} - # p_rh = [(1−ε), ε]; expanded KL gives the weighted logsigmoid form below - epsilon = torch.tensor(self.label_smoothing, device=device) - qw = torch.sigmoid(self.beta * delta_score) - log_qw = F.logsigmoid(self.beta * delta_score) - log_pw = torch.log1p(-epsilon) - ql = torch.sigmoid(-self.beta * delta_score) - log_ql = F.logsigmoid(-self.beta * delta_score) - log_pl = torch.log(epsilon) - per_sequence_loss = qw * (log_qw - log_pw) + ql * (log_ql - log_pl) - - elif loss_type == "nca_pair": - chosen_rewards = self.beta * chosen_logratios - rejected_rewards = self.beta * rejected_logratios - per_sequence_loss = ( - -F.logsigmoid(chosen_rewards) - - 0.5 * F.logsigmoid(-chosen_rewards) - - 0.5 * F.logsigmoid(-rejected_rewards) - ) - - elif loss_type == "robust": - clean_loss_term = -(1 - self.label_smoothing) * F.logsigmoid(self.beta * delta_score) - flipped_loss_term = -self.label_smoothing * F.logsigmoid(-self.beta * delta_score) - per_sequence_loss = (clean_loss_term - flipped_loss_term) / (1 - 2 * self.label_smoothing) - - elif loss_type == "bco_pair": - chosen_rewards = self.beta * chosen_logratios - rejected_rewards = self.beta * rejected_logratios - per_sequence_loss = -F.logsigmoid(chosen_rewards) - F.logsigmoid(-rejected_rewards) - - elif loss_type == "sppo_hard": - # In the paper (https://huggingface.co/papers/2405.00675), SPPO employs a soft probability approach, - # estimated using the PairRM score. The probability calculation is conducted outside of the trainer - # class. The version described here is the hard probability version, where P in Equation (4.7) of - # Algorithm 1 is set to 1 for the winner and 0 for the loser. - winner_margin_error = (chosen_logratios - 0.5 / self.beta) ** 2 - loser_margin_error = (rejected_logratios + 0.5 / self.beta) ** 2 - per_sequence_loss = winner_margin_error + loser_margin_error - - elif loss_type == "aot": - logratios = chosen_logps - rejected_logps - ref_logratios = ref_chosen_logps - ref_rejected_logps - logratios_sorted, _ = torch.sort(logratios, dim=0) - ref_logratios_sorted, _ = torch.sort(ref_logratios, dim=0) - delta = logratios_sorted - ref_logratios_sorted - per_sequence_loss = ( - -F.logsigmoid(self.beta * delta) * (1 - self.label_smoothing) - - F.logsigmoid(-self.beta * delta) * self.label_smoothing - ) - - elif loss_type == "aot_unpaired": - chosen_logratios_sorted, _ = torch.sort(chosen_logratios, dim=0) - rejected_logratios_sorted, _ = torch.sort(rejected_logratios, dim=0) - delta = chosen_logratios_sorted - rejected_logratios_sorted - per_sequence_loss = ( - -F.logsigmoid(self.beta * delta) * (1 - self.label_smoothing) - - F.logsigmoid(-self.beta * delta) * self.label_smoothing - ) - - elif loss_type == "apo_zero": - # Eqn (7) of the APO paper (https://huggingface.co/papers/2408.06266) - # Use this loss when you believe the chosen outputs are better than your model's default output - # Increase chosen likelihood and decrease rejected likelihood - losses_chosen = 1 - torch.sigmoid(self.beta * chosen_logratios) - losses_rejected = torch.sigmoid(self.beta * rejected_logratios) - per_sequence_loss = losses_chosen + losses_rejected - - elif loss_type == "apo_down": - # Eqn (8) of the APO paper (https://huggingface.co/papers/2408.06266) - # Use this loss when you believe the chosen outputs are worse than your model's default output. - # Decrease chosen likelihood and decrease rejected likelihood more - losses_chosen = torch.sigmoid(self.beta * chosen_logratios) - losses_rejected = 1 - torch.sigmoid(self.beta * (chosen_logratios - rejected_logratios)) - per_sequence_loss = losses_chosen + losses_rejected - - elif loss_type == "discopop": - # Eqn (5) of the DiscoPOP paper (https://huggingface.co/papers/2406.08414) - logits = delta_score * self.beta - # Modulate the mixing coefficient based on the log ratio magnitudes - log_ratio_modulation = torch.sigmoid(logits / self.args.discopop_tau) - logistic_component = -F.logsigmoid(logits) - exp_component = torch.exp(-logits) - # Blend between logistic and exponential component based on log ratio modulation - per_sequence_loss = ( - logistic_component * (1 - log_ratio_modulation) + exp_component * log_ratio_modulation - ) - - elif loss_type == "sft": - chosen_mask, _ = shift_completion_mask.chunk(2, dim=0) - chosen_per_token_logps, _ = per_token_logps.chunk(2, dim=0) - batch_loss = -chosen_per_token_logps[chosen_mask.bool()].mean() - batch_size = chosen_per_token_logps.size(0) - # Implementation convenience: expand the scalar SFT loss to a per-sequence tensor so it matches the - # shape of other losses; only the mean is used, so this is a no-op numerically. - per_sequence_loss = batch_loss.expand(batch_size) - - elif loss_type == "sigmoid_norm": - chosen_mask, rejected_mask = completion_mask.chunk(2, dim=0) - chosen_avg_score = chosen_scores / chosen_mask.sum(dim=1).clamp(min=1.0) - rejected_avg_score = rejected_scores / rejected_mask.sum(dim=1).clamp(min=1.0) - delta = chosen_avg_score - rejected_avg_score - per_sequence_loss = -F.logsigmoid(self.beta * delta) - - else: - raise ValueError( - f"Unknown loss type: {loss_type}. Should be one of ['sigmoid', 'hinge', 'ipo', 'exo_pair', " - "'nca_pair', 'robust', 'bco_pair', 'sppo_hard', 'aot', 'aot_unpaired', 'apo_zero', 'apo_down', " - "'discopop', 'sft', 'sigmoid_norm']" - ) - - if self.use_weighting: - # Eq (2) of the WPO paper: https://huggingface.co/papers/2406.11827 - completion_lengths = shift_completion_mask.sum(dim=1).clamp_min(1) - with torch.no_grad(): - aligned_logps = (per_token_logps - outputs.log_sum_sq_probs[:, -seq_len:]) * shift_completion_mask - mean_logps = aligned_logps.sum(dim=1) / completion_lengths - weights = torch.exp(mean_logps) - chosen_weights, rejected_weights = weights.chunk(2, dim=0) - per_sequence_loss *= chosen_weights * rejected_weights - - loss += per_sequence_loss.mean() * loss_weight - - if self.aux_loss_enabled: - aux_loss = outputs.aux_loss - loss = loss + self.router_aux_loss_coef * aux_loss - self._metrics[mode]["aux_loss"].append(self.accelerator.gather_for_metrics(aux_loss).mean().item()) - - # Log the metrics - # Entropy - per_token_entropy = outputs.entropy[:, -seq_len:].detach() - mask = shift_completion_mask - entropy_sum = (per_token_entropy * mask).sum() - total_tokens = mask.sum() - - # Gather counts across ranks and weight-average - entropy_sum = self.accelerator.gather_for_metrics(entropy_sum).sum() - total_tokens = self.accelerator.gather_for_metrics(total_tokens).sum() - entropy = (entropy_sum / total_tokens).item() if total_tokens > 0 else 0.0 - self._metrics[mode]["entropy"].append(entropy) - - # Number of tokens - if mode == "train": - num_tokens_in_batch = self.accelerator.gather_for_metrics(inputs["attention_mask"].sum()).sum().item() - self._total_train_tokens += num_tokens_in_batch // self._tp_size - self._metrics[mode]["num_tokens"] = [self._total_train_tokens] - - # Average logits for chosen and rejected completions - chosen_mean_logits, rejected_mean_logits = outputs.mean_logits[:, -seq_len:].chunk(2, dim=0) - chosen_mask, rejected_mask = shift_completion_mask.chunk(2, dim=0) - total_chosen_logits = chosen_mean_logits[chosen_mask.bool()].sum() - total_chosen_tokens = chosen_mask.sum() - total_rejected_logits = rejected_mean_logits[rejected_mask.bool()].sum() - total_rejected_tokens = rejected_mask.sum() - total_chosen_logits = self.accelerator.gather_for_metrics(total_chosen_logits).sum().item() - total_chosen_tokens = self.accelerator.gather_for_metrics(total_chosen_tokens).sum().item() - total_rejected_logits = self.accelerator.gather_for_metrics(total_rejected_logits).sum().item() - total_rejected_tokens = self.accelerator.gather_for_metrics(total_rejected_tokens).sum().item() - avg_chosen_logits = total_chosen_logits / total_chosen_tokens if total_chosen_tokens > 0 else 0.0 - avg_rejected_logits = total_rejected_logits / total_rejected_tokens if total_rejected_tokens > 0 else 0.0 - self._metrics[mode]["logits/chosen"].append(avg_chosen_logits) - self._metrics[mode]["logits/rejected"].append(avg_rejected_logits) - - # Token accuracy for the chosen completions - chosen_is_top1, _ = outputs.is_top1[:, -seq_len:].chunk(2, dim=0) - correct_predictions = chosen_is_top1 & chosen_mask.bool() - total_tokens = chosen_mask.sum() - correct_tokens = correct_predictions.sum() - correct_tokens = self.accelerator.gather_for_metrics(correct_tokens) - total_tokens = self.accelerator.gather_for_metrics(total_tokens) - total_sum = total_tokens.sum() - accuracy = (correct_tokens.sum() / total_sum).item() if total_sum > 0 else 0.0 - self._metrics[mode]["mean_token_accuracy"].append(accuracy) - - # Rewards for chosen and rejected completions - chosen_rewards = self.beta * chosen_logratios.detach() - rejected_rewards = self.beta * rejected_logratios.detach() - agg_chosen_rewards = self.accelerator.gather(chosen_rewards) - agg_rejected_rewards = self.accelerator.gather(rejected_rewards) - self._metrics[mode]["rewards/chosen"].append(agg_chosen_rewards.mean().item()) - self._metrics[mode]["rewards/rejected"].append(agg_rejected_rewards.mean().item()) - - # Reward accuracy - reward_accuracies = (chosen_rewards > rejected_rewards).float() - agg_reward_accuracies = self.accelerator.gather(reward_accuracies) - self._metrics[mode]["rewards/accuracies"].append(agg_reward_accuracies.mean().item()) - - # Reward margins - margins = chosen_rewards - rejected_rewards - agg_margins = self.accelerator.gather(margins) - self._metrics[mode]["rewards/margins"].append(agg_margins.mean().item()) - - # Average log probabilities for chosen and rejected completions - self._metrics[mode]["logps/chosen"].append(self.accelerator.gather(chosen_logps).mean().item()) - self._metrics[mode]["logps/rejected"].append(self.accelerator.gather(rejected_logps).mean().item()) - - if return_outputs: - warnings.warn( - "`return_outputs=True`, and the logits it passes to `compute_metrics`, are deprecated and will be removed " - "in v2.0.0. The fused LM head does not build the logits, so they are recomputed with an extra forward pass.", - FutureWarning, - stacklevel=2, - ) - with torch.no_grad(): - outputs = model(**model_kwargs) - return (loss, outputs) if return_outputs else loss - - def evaluate( - self, - eval_dataset: Dataset - | IterableDataset - | DatasetDict - | IterableDatasetDict - | dict[str, Dataset | IterableDataset] - | None = None, - ignore_keys: list[str] | None = None, - metric_key_prefix: str = "eval", - ) -> dict[str, float]: - # When a dataset is passed directly to `evaluate` (e.g. a held-out test set), preprocess it the same way - # `__init__` does, so that `evaluate` accepts the same dataset types as the trainer. `_prepare_dataset` is - # idempotent: it skips datasets that are already tokenized. A `str` selects a dataset that was already prepared - # at init time, so it's left untouched. - if not self._is_vision_dataset and eval_dataset is not None and not isinstance(eval_dataset, str): - # Full fine-tuning with no `ref_model` uses `self.model` as the reference, which is only valid before - # training. After a step (`global_step > 0`) it's the trained policy, so we can't precompute a correct - # reference here. (PEFT is exempt: the reference is recovered by disabling the adapter.) Checked before - # tokenizing so we fail fast. - if ( - self.precompute_ref_logps - and self.ref_model is None - and not is_peft_model(self.model) - and self.state.global_step > 0 - ): - raise ValueError( - "Cannot compute reference log-probs for a dataset passed to `evaluate()` after training has " - "started, because `precompute_ref_log_probs=True` and no `ref_model` was provided (full " - "fine-tuning). In this setup the reference model is not kept in memory, so it is only available " - "before training. Provide this dataset as `eval_dataset` at initialization, pass an explicit " - "`ref_model`, or set `precompute_ref_log_probs=False`." - ) - if isinstance(eval_dataset, dict): - eval_dataset = { - key: self._prepare_dataset(dataset, self.processing_class, self.args, key) - for key, dataset in eval_dataset.items() - } - else: - eval_dataset = self._prepare_dataset(eval_dataset, self.processing_class, self.args, "eval") - # With `precompute_ref_log_probs`, `_compute_loss` reads the reference log-probs from the batch, so they - # must be precomputed here as well, mirroring `__init__`. - if self.precompute_ref_logps: - batch_size = self.args.precompute_ref_batch_size or self.args.per_device_eval_batch_size - if isinstance(eval_dataset, dict): - eval_dataset = { - name: self._precompute_ref_logps(dataset, name, batch_size) - for name, dataset in eval_dataset.items() - } - else: - eval_dataset = self._precompute_ref_logps(eval_dataset, "eval", batch_size) - # Call `super().evaluate()` once per split ourselves instead of handing the whole dict to it: `Trainer.evaluate` - # would otherwise recurse into `self.evaluate` per split, re-entering this override on an already-prepared dataset. - if isinstance(eval_dataset, dict): - metrics = {} - for name, dataset in eval_dataset.items(): - metrics.update( - super().evaluate( - eval_dataset=dataset, - ignore_keys=ignore_keys, - metric_key_prefix=f"{metric_key_prefix}_{name}", - ) - ) - return metrics - return super().evaluate( - eval_dataset=eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix - ) - - def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): - try: - return self._compute_loss(model, inputs, return_outputs) - except ValueError as e: - if "Image features and image tokens do not match" in str(e) and self.args.max_length is not None: - raise ValueError( - f"The current `max_length` ({self.args.max_length}) is too short and causes image placeholder " - f"tokens in `input_ids` to be truncated, while the corresponding image features remain intact. " - f"Please increase `max_length` or set it to `None` to disable truncation." - ) from e - raise - - # Override training step to add activation offloading context. - def training_step(self, *args, **kwargs): - with self.maybe_activation_offload_context: - return super().training_step(*args, **kwargs) - - def log(self, logs: dict[str, float], start_time: float | None = None) -> None: - mode = "train" if self.model.training else "eval" - metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()} # average the metrics - # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` - # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. - if mode == "eval": - metrics = {f"eval_{key}": val for key, val in metrics.items()} - logs.update(metrics) - super().log(logs, start_time) - self._metrics[mode].clear() - - # During eval, Trainer calls prediction_step. If no labels are present in the inputs, it only runs forward and - # returns logits. We override prediction_step to force compute_loss, because this trainer doesn't involve labels. - def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys: list[str] | None = None): - inputs = self._prepare_inputs(inputs) - with torch.no_grad(), self.compute_loss_context_manager(): - if prediction_loss_only: - loss = self.compute_loss(model, inputs, return_outputs=False) - logits, labels = None, None - else: - loss, outputs = self.compute_loss(model, inputs, return_outputs=True) - logits, labels = outputs.logits, inputs["input_ids"] - return loss, logits, labels - - # Ensure the model card is saved along with the checkpoint - def _save_checkpoint(self, model, trial): - if self.args.hub_model_id is None: - model_name = Path(self.args.output_dir).name - else: - model_name = self.args.hub_model_id.split("/")[-1] - self.create_model_card(model_name=model_name) - super()._save_checkpoint(model, trial) -class UnslothDPOTrainer(_UnslothDPOTrainer): - """ - - Trainer for Direct Preference Optimization (DPO) method. This algorithm was initially proposed in the paper [Direct - Preference Optimization: Your Language Model is Secretly a Reward Model](https://huggingface.co/papers/2305.18290). - This class is a wrapper around the [`~transformers.Trainer`] class and inherits all of its attributes and methods. - - Example: - - ```python - >>> from trl import DPOTrainer - >>> from datasets import load_dataset - - >>> dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train") - - >>> trainer = DPOTrainer( - ... model="Qwen/Qwen2.5-0.5B-Instruct", - ... train_dataset=dataset, - ... ) - >>> trainer.train() - ``` - - Args: - model (`str` or [`~transformers.PreTrainedModel`] or [`~peft.PeftModel`]): - Model to be trained. Can be either: - - - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a - path to a *directory* containing model weights saved using - [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded - using `.from_pretrained` (where `` is derived from the model - config) with the keyword arguments in `args.model_init_kwargs`. If `dtype` is not specified in - `args.model_init_kwargs`, it defaults to `float32`. This differs from - [`~transformers.PreTrainedModel.from_pretrained`], where (since Transformers v5) the dtype is inferred - from the model config. - - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. - - A [`~peft.PeftModel`] object. Only causal language models are supported. - ref_model ([`~transformers.PreTrainedModel`], *optional*): - Reference model used to compute the reference log probabilities. - - - If provided, this model is used directly as the reference policy. - - If `None`, the trainer will automatically use the initial policy corresponding to `model`, i.e. the model - state before DPO training starts. - args ([`DPOConfig`], *optional*): - Configuration for this trainer. If `None`, a default configuration is used. - data_collator ([`~transformers.DataCollator`], *optional*): - Function to use to form a batch from a list of elements of the processed `train_dataset` or `eval_dataset`. - Will default to [`~trainer.dpo_trainer.DataCollatorForPreference`] if the model is a language model and - [`~trainer.dpo_trainer.DataCollatorForVisionPreference`] if the model is a vision-language model. Custom - collators must truncate sequences before padding; the trainer does not apply post-collation truncation. - train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): - Dataset to use for training. This trainer supports both [language modeling](#language-modeling) type and - [prompt-completion](#prompt-completion) type. The format of the samples can be either: - - - [Standard](dataset_formats#standard): Each sample contains plain text. - - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role - and content). - - When `train_dataset` is an [`~datasets.IterableDataset`] (e.g. a streaming dataset), `max_steps` must be - set in the training arguments, since its length cannot be inferred and the total number of training steps - is required to bound the training loop and configure the learning rate scheduler. - eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`], [`~datasets.DatasetDict`], [`~datasets.IterableDatasetDict`] or `dict[str, Dataset | IterableDataset]`): - Dataset to use for evaluation. It must meet the same requirements as `train_dataset`. - processing_class ([`~transformers.PreTrainedTokenizerBase`] or [`~transformers.ProcessorMixin`], *optional*): - Processing class used to process the data. The padding side must be set to "left". If `None`, the - processing class is loaded from the model's name with [`~transformers.AutoProcessor.from_pretrained`]. A - padding token, `tokenizer.pad_token`, must be set. If the processing class has not set a padding token, - `tokenizer.eos_token` will be used as the default. - compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): - The function that will be used to compute metrics at evaluation. Must take a - [`~transformers.EvalPrediction`] and return a dictionary string to metric values. When passing - [`SFTConfig`] with `batch_eval_metrics` set to `True`, your `compute_metrics` function must take a boolean - `compute_result` argument. This will be triggered after the last eval batch to signal that the function - needs to calculate and return the global summary statistics rather than accumulating the batch-level - statistics. The logits it receives are recomputed with an extra forward pass. This is deprecated and will - be removed in v2.0.0. - callbacks (list of [`~transformers.TrainerCallback`], *optional*): - List of callbacks to customize the training loop. Will add those to the list of default callbacks detailed - in [here](https://huggingface.co/docs/transformers/main_classes/callback). - - If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`] - method. - optimizers (`tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None]`, *optional*, defaults to `(None, None)`): - A tuple containing the optimizer and the scheduler to use. Will default to an instance of `AdamW` on your - model and a scheduler given by [`~transformers.get_linear_schedule_with_warmup`] controlled by `args`. - quantization_config ([`~transformers.BitsAndBytesConfig`], *optional*): - Quantization configuration used when loading the model from a model identifier. Combine with `peft_config` - for QLoRA training. Ignored if the model is already instantiated. - peft_config ([`~peft.PeftConfig`], *optional*): - PEFT configuration used to wrap the model. If `None`, the model is not wrapped. - - """ - def __init__( - self, - model, - ref_model = None, - args = None, - data_collator = None, - train_dataset = None, - eval_dataset = None, - processing_class = None, - compute_metrics = None, - callbacks = None, - quantization_config = None, - peft_config = None, - **kwargs - ): - if args is None: args = UnslothDPOConfig() - _unsloth_precision_requested = tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) - use_bf16 = getattr(args, 'bf16', False) - if type(use_bf16) is not bool: use_bf16 = False - use_fp16 = getattr(args, 'fp16', False) - if type(use_fp16) is not bool: use_fp16 = False - force_float32 = False - try: - from unsloth_zoo.device_type import device_is_bf16_supported as _bf16_supported - except Exception: - _bf16_supported = torch.cuda.is_bf16_supported - try: - from unsloth_zoo.device_type import DEVICE_TYPE as _unsloth_device_type - except Exception: - _unsloth_device_type = None - full_finetuning = getattr(model, '_unsloth_full_finetuning', None) - if full_finetuning is None: full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1' - model_forced_float32 = getattr(model, '_unsloth_forced_float32', None) - if model_forced_float32 is None: model_forced_float32 = os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1' - if model_forced_float32 and not (full_finetuning and _bf16_supported()): - print('Unsloth: Switching to float32 training since model cannot work with float16') - force_float32 = True - mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') - dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None) - if dtype is None: dtype = model.get_input_embeddings().weight.dtype - from unsloth_zoo.utils import _get_dtype - dtype = _get_dtype(dtype) - float16 = dtype == torch.float16 - bfloat16 = dtype == torch.bfloat16 - float32 = dtype == torch.float32 - user_float32 = bool(getattr(model, '_unsloth_user_float32', False)) - if full_finetuning: - if bfloat16 and use_fp16: use_fp16 = False - if float16 and use_bf16: use_bf16 = False - if not force_float32 and bfloat16 and use_fp16 and _unsloth_device_type == 'xpu' and not _bf16_supported(): - print('Unsloth: The model was loaded in bfloat16 on purpose, so training uses bfloat16 mixed precision instead of float16.') - use_fp16 = False - use_bf16 = True - args.fp16 = False - args.bf16 = True - if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`') - if not force_float32 and (bfloat16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`') - if force_float32: - # Forced float32 training - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32' and float32 and user_float32 and not _bf16_supported(): - print('Unsloth: Model is in float32 and this GPU has no bfloat16 support, so training stays in float32. Pass fp16 = True to force float16 mixed precision instead.') - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32': - # Mixed precision training. bf16 only if the GPU supports it; V100/T4 use fp16. - use_bf16_amp = (not float16) and (_bf16_supported() or (bfloat16 and _unsloth_device_type == 'xpu')) - args.fp16 = not use_bf16_amp - args.bf16 = use_bf16_amp - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16_amp else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16_amp else 'fp16' - # args.mixed_precision is a new argument which needs to be set now - elif mixed_precision_dtype == 'bfloat16': - # Both False since bfloat16 full finetuning doesn't do any autocasting. - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif use_bf16 or use_fp16: - # transformers <5 exported this itself from fp16/bf16; 5.x dropped the write, so an - # explicit flag left it unset and GRPO readers defaulted to 'fp16', wrapping a - # bfloat16 model in a float16 autocast. See unslothai/unsloth#4891. - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16 else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16 else 'fp16' - - if getattr(args, 'eval_dataset', None) is not None and getattr(args, 'eval_strategy', 'no') == 'no': - args.eval_strategy = 'steps' - if getattr(args, 'eval_steps', None) is None: args.eval_steps = 0.1 - ga_steps = getattr(args, 'gradient_accumulation_steps', None) - if ga_steps is not None and ga_steps > 1: - from transformers import __version__ as transformers_version - if Version(transformers_version) <= Version('4.45.2'): - print('**** Unsloth: Please use our fixed gradient_accumulation_steps by updating transformers, TRL and Unsloth!\n' - '`pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth transformers trl unsloth_zoo`') - if getattr(args, 'eval_strategy', 'no') != 'no': - eval_bsz = getattr(args, 'per_device_eval_batch_size', 8) - eval_generations = getattr(args, 'num_generations_eval', None) or getattr(args, 'num_generations', None) or 1 - if eval_bsz == 8 and args.per_device_train_batch_size < eval_bsz and (args.per_device_train_batch_size * args.world_size) % eval_generations == 0: args.per_device_eval_batch_size = args.per_device_train_batch_size - if getattr(args, 'eval_accumulation_steps', None) is None and ga_steps is not None: args.eval_accumulation_steps = ga_steps - fp16_full_eval = getattr(args, 'fp16_full_eval', False) - if type(fp16_full_eval) is not bool: fp16_full_eval = False - bf16_full_eval = getattr(args, 'bf16_full_eval', False) - if type(bf16_full_eval) is not bool: bf16_full_eval = False - if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True - if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False - if force_float32: - args.bf16_full_eval = False - args.fp16_full_eval = False - elif os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16': - args.bf16_full_eval = True - args.fp16_full_eval = False - elif not bf16_full_eval and not fp16_full_eval: - args.bf16_full_eval = args.bf16 - args.fp16_full_eval = args.fp16 - if tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) != _unsloth_precision_requested: - from unsloth.models.loader_utils import sync_deepspeed_mixed_precision - sync_deepspeed_mixed_precision(args) - if model is not None: - from unsloth.models.loader_utils import apply_deepspeed_zero3_model_setup - apply_deepspeed_zero3_model_setup(model) - if model is not None: - _warnings_issued = getattr(model, 'warnings_issued', None) - if _warnings_issued is None: - model.warnings_issued = {} - elif not isinstance(_warnings_issued, dict): - try: - model.warnings_issued = dict(_warnings_issued) - except Exception: - model.warnings_issued = {} - if 'max_seq_length' not in locals() and not hasattr(args, 'max_seq_length'): - pass - else: - model_max_seq_length = getattr(model, 'max_seq_length', None) - args_max_seq_length = getattr(args, 'max_seq_length', None) - if args_max_seq_length is None and model_max_seq_length is not None: - max_seq_length = model.max_seq_length - if hasattr(args, 'max_seq_length'): args.max_seq_length = max_seq_length - elif args_max_seq_length is not None and model_max_seq_length is not None: - if args_max_seq_length > model_max_seq_length: - print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but ' - 'the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.') - args.max_seq_length = model_max_seq_length - _unsloth_model_msl = getattr(model, 'max_seq_length', None) - if isinstance(_unsloth_model_msl, int) and _unsloth_model_msl > 0 and hasattr(args, 'max_length'): - if args.max_length is None or args.max_length > _unsloth_model_msl: - print('Unsloth: `max_length = ' + str(args.max_length) + '` exceeds the model max_seq_length of ' + str(_unsloth_model_msl) + ', so it is reduced to ' + str(_unsloth_model_msl) + '.') - args.max_length = _unsloth_model_msl - _unsloth_mpl = getattr(args, 'max_prompt_length', 0) - if (_unsloth_mpl is None and args.max_length <= 128) or (_unsloth_mpl is not None and _unsloth_mpl >= args.max_length): - args.max_prompt_length = args.max_length // 2 - if model is not None and hasattr(model, 'for_training'): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right' - if 'processing_class' in locals(): - if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right' - if hasattr(processing_class, 'tokenizer') and hasattr(processing_class.tokenizer, 'padding_side'): processing_class.tokenizer.padding_side = 'right' - __tokenizer = processing_class if 'processing_class' in locals() else tokenizer - from unsloth_zoo.vision_utils import UnslothVisionDataCollator - if not isinstance(data_collator, UnslothVisionDataCollator): - if isinstance(data_collator, DataCollatorForSeq2Seq) and 'labels' not in _unsloth_dataset_column_names(train_dataset): - data_collator = TransformersDataCollatorForLanguageModeling( - __tokenizer, - mlm = False, - mlm_probability = 0.0, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling) and 'labels' in _unsloth_dataset_column_names(train_dataset): - data_collator = DataCollatorForSeq2Seq( - __tokenizer, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - else: - if hasattr(args, 'remove_unused_columns'): args.remove_unused_columns = False - if hasattr(args, 'dataset_text_field'): args.dataset_text_field = '' - if hasattr(args, 'dataset_kwargs'): args.dataset_kwargs = {'skip_prepare_dataset': True} - if not isinstance(data_collator, UnslothVisionDataCollator): - if not hasattr(__tokenizer, 'pad') and hasattr(__tokenizer, 'tokenizer'): - if isinstance(data_collator, DataCollatorForSeq2Seq): - data_collator = DataCollatorForSeq2Seq( - __tokenizer.tokenizer, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling): - data_collator = TransformersDataCollatorForLanguageModeling( - __tokenizer.tokenizer, - mlm = False, - mlm_probability = 0.0, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - other_metrics = [] - - from unsloth_zoo.logging_utils import PatchRLStatistics - PatchRLStatistics('dpo_trainer', other_metrics) - if hasattr(train_dataset, 'column_names'): - column_names = set(train_dataset.column_names) - check = ['chosen', 'rejected', 'prompt', 'chosen_input_ids', 'chosen_attention_mask', - 'chosen_labels', 'rejected_input_ids', 'rejected_attention_mask', 'rejected_labels', - 'prompt_input_ids', 'prompt_attention_mask'] - if all(x in column_names for x in check): - train_dataset = train_dataset.remove_columns(['chosen', 'rejected', 'prompt']) - del check, column_names - if hasattr(train_dataset, 'column_names'): - column_names = set(train_dataset.column_names) - is_dpo_dataset = ({'chosen', 'rejected'}.issubset(column_names) or - {'prompt_input_ids', 'chosen_input_ids', 'rejected_input_ids'}.issubset(column_names)) - if is_dpo_dataset and isinstance(data_collator, TransformersDataCollatorForLanguageModeling): - data_collator = None - del is_dpo_dataset, column_names - from trl.trainer.dpo_trainer import DataCollatorForPreference - if not hasattr(DataCollatorForPreference, '_unsloth_vision_keys_patch'): - _old_dpo_collator_torch_call = DataCollatorForPreference.torch_call - - def _unsloth_dpo_torch_call(self, examples): - output = _old_dpo_collator_torch_call(self, examples) - import torch as _unsloth_torch - try: - from trl.trainer.utils import pad as _unsloth_trl_pad - except Exception: - _unsloth_trl_pad = None - for _k in ('pixel_position_ids', 'image_position_ids', 'mm_token_type_ids'): - if not all(_k in example for example in examples): - continue - _is_position_key = _k.endswith('position_ids') - _padding_value = -1 if _is_position_key else 0 - _padding_side = 'right' if _is_position_key else 'left' - _values = [_unsloth_torch.as_tensor(example[_k]) for example in examples] - try: - if _unsloth_trl_pad is not None: - output[_k] = _unsloth_trl_pad(_values, padding_value=_padding_value, padding_side=_padding_side) - else: - from torch.nn.utils.rnn import pad_sequence as _unsloth_pad_sequence - output[_k] = _unsloth_pad_sequence(_values, batch_first=True, padding_value=_padding_value) - except Exception: - from torch.nn.utils.rnn import pad_sequence as _unsloth_pad_sequence - output[_k] = _unsloth_pad_sequence(_values, batch_first=True, padding_value=_padding_value) - return output - - DataCollatorForPreference.torch_call = _unsloth_dpo_torch_call - DataCollatorForPreference._unsloth_vision_keys_patch = True - - # [TODO] Fix up DataParallel multiplying batch sizes - # [TODO] DDP works, but DP seems to not work? [TODO] - if getattr(args, "parallel_mode", None) == ParallelMode.NOT_DISTRIBUTED and args.n_gpu > 1: - if getattr(args, "_n_gpu", 1) != 1: - args._n_gpu = 1 - if "model" in locals(): - try: - from unsloth.models._gc_config_default import honour_gradient_checkpointing_off - honour_gradient_checkpointing_off(args, model) - except Exception: - pass - if "model" in locals() and hasattr(model, "for_training"): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if "model" in locals(): - try: - from unsloth.models._ddp_defaults import unsloth_default_ddp_find_unused_parameters - unsloth_default_ddp_find_unused_parameters(args, model) - except Exception: - pass - super().__init__( - model = model, - ref_model = ref_model, - args = args, - data_collator = data_collator, - train_dataset = train_dataset, - eval_dataset = eval_dataset, - processing_class = processing_class, - compute_metrics = compute_metrics, - callbacks = callbacks, - quantization_config = quantization_config, - peft_config = peft_config,**kwargs) - # A tensor parallel model (tp_native.tp_trainer_setup) stays on its training path: for_inference would - # route the decoder layers onto decode paths outside the tensor parallel regions. - if "model" in locals() and hasattr(model, "for_inference") and not getattr(self, "_unsloth_tensor_parallel", False): - model.for_inference() - if hasattr(self, 'neftune_hook_handle'): - self.neftune_hook_handle.remove() - if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle - if getattr(args, 'neftune_noise_alpha', None) is not None: - model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha - pass - if hasattr(self, 'accelerator'): - scaler = self.accelerator.scaler - current_model = model - while hasattr(current_model, 'model'): - current_model.accelerator_scaler = scaler - current_model = current_model.model - current_model.accelerator_scaler = scaler - pass - if hasattr(self, 'train'): - self.train = MethodType(prepare_for_training_mode(self.__class__.train), self) - pass - if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'): - _vllm_tok = self.llm.get_tokenizer() - _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None) - if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None: - _vllm_tok.chat_template = _pc.chat_template - pass - if getattr(self, 'aux_loss_enabled', False) and hasattr(getattr(self, 'model', None), 'config'): - _text_config = self.model.config - if hasattr(_text_config, 'get_text_config'): _text_config = _text_config.get_text_config() - _n_experts = [getattr(_text_config, _k) for _k in ('num_local_experts', 'num_experts', 'n_routed_experts', 'moe_num_experts') if isinstance(getattr(_text_config, _k, None), int)] - if _n_experts and all(_n <= 1 for _n in _n_experts): - self.aux_loss_enabled = False - _text_config.output_router_logits = False - pass - -pass - - -if hasattr(logger, "addFilter"): - import logging - class HideLoggingMessage(logging.Filter): - def __init__(self, text): self.text = text - def filter(self, x): return not (self.text in x.getMessage()) - pass - logger.addFilter(HideLoggingMessage("`use_cache=True`")) - diff --git a/unsloth_compiled_cache/UnslothDistillationTrainer.py b/unsloth_compiled_cache/UnslothDistillationTrainer.py deleted file mode 100644 index 2be9f4df8c059543b23f799a81e9c8f93b90d4d8..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/UnslothDistillationTrainer.py +++ /dev/null @@ -1,3026 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from unsloth_zoo.temporary_patches.common import torch_compile -from unsloth_zoo.temporary_patches.common import _maybe_compile -import functools -import transformers -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from trl.trainer.distillation_trainer import (Any, AutoProcessor, BaseTunerLayer, BitsAndBytesConfig, Callable, DataLoader, Dataset, DistillationConfig, DistillationTrainer, DistributedBackend, GenerationConfig, IterableDataset, Path, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, PromptLearningConfig, RepeatSampler, Sampler, TrainerCallback, VLLMGeneration, Version, _BaseTrainer, _ForwardRedirection, _SUPPORTS_RESPONSE_TEMPLATE, _chunked_divergence_loss, add_response_schema, apply_chat_template, copy, create_model_from_path, defaultdict, deque, disable_dropout_in_model, gather_object, get_config_model_id, get_peft_model, get_training_chat_template, identity, inspect, is_chat_template_prefix_preserving, is_conversational, is_jmespath_available, is_peft_available, is_peft_model, is_rich_available, is_vllm_available, logger, math, np, os, pad, parse_response, prepare_deepspeed, prepare_multimodal_messages, print_prompt_completions_sample, profiling_context, profiling_decorator, repeat_iterable_dataset, set_seed, shuffle_sequence_dict, split_pixel_values_by_grid, split_tensor_dict, supports_tool_calling, sys, textwrap, time, torch, transformers, unsplit_pixel_values_by_grid, unwrap_model_for_generation, warnings, AutoProcessor, BaseTunerLayer, BitsAndBytesConfig, Callable, Dataset, DistillationConfig, DistillationTrainer, DistributedBackend, GenerationConfig, IterableDataset, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, PromptLearningConfig, RepeatSampler, Sampler, TrainerCallback, VLLMGeneration, Version, add_response_schema, copy, create_model_from_path, defaultdict, deque, disable_dropout_in_model, get_config_model_id, get_peft_model, get_training_chat_template, identity, inspect, is_chat_template_prefix_preserving, is_jmespath_available, is_peft_available, is_peft_model, is_vllm_available, logger, np, os, pad, prepare_deepspeed, repeat_iterable_dataset, set_seed, supports_tool_calling, sys, time, torch, transformers, warnings, Any, np, profiling_decorator, shuffle_sequence_dict, split_pixel_values_by_grid, split_tensor_dict, torch, unsplit_pixel_values_by_grid, PeftModel, PreTrainedModel, is_peft_available, logger, os, torch) - - -import os -import math -import logging -from typing import * -from dataclasses import dataclass, field -from packaging.version import Version -import torch -import numpy as np -from contextlib import nullcontext -from torch.nn import functional as F -import inspect -from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling -from transformers.training_args import ParallelMode -from unsloth_zoo.device_type import DEVICE_TYPE, DEVICE_TYPE_TORCH, device_synchronize - -# Wrap trainer with padding to right and enable training mode -import functools -from types import MethodType -try: - from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers -except: - def reset_unsloth_gradient_checkpointing_buffers(): pass -# Canonical reset lives in unsloth.models._utils so the SFT auto-packing wrapper and the plain -# Trainer loop can import the same helper; fall back to a no-op only if it can't be imported. -try: - from unsloth.models._utils import _unsloth_reset_stray_compile_cache -except Exception: - def _unsloth_reset_stray_compile_cache(self): pass -try: - from unsloth.models._utils import _unsloth_dataset_column_names -except Exception: - def _unsloth_dataset_column_names(dataset): return dataset.column_names -# Drops/renames config arguments the installed TRL no longer accepts, so a -# script pinned to an older TRL keeps working after an upgrade. Falls back to -# the historical raw passthrough so this can never break trainer construction. -try: - from unsloth.models.rl_config_compat import filter_config_init_kwargs as _unsloth_filter_config_init_kwargs - # A cache file generated here can be imported by an older Unsloth whose filter - # predates `mirrored_from`, so drop the argument rather than raise TypeError. - if "mirrored_from" not in inspect.signature(_unsloth_filter_config_init_kwargs).parameters: - _unsloth_filter_config_init_kwargs_old = _unsloth_filter_config_init_kwargs - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): - return _unsloth_filter_config_init_kwargs_old(config_class, kwargs) -except Exception: - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): return kwargs -def prepare_for_training_mode(f): - @functools.wraps(f) - def wrapper(self, *args, **kwargs): - # Drop any torch.compile graph cache poisoned by a stray pre-train forward. - try: - _unsloth_reset_stray_compile_cache(self) - except Exception: - pass - # Finish the previous W&B run if this is a subsequent train() call. - # We do this at the START of train() (not the end) so that - # evaluate() / log() still work after train() completes. - # HF's WandbCallback.setup() will call wandb.init() for the new run. - # See: https://github.com/unslothai/unsloth/issues/3954 - if getattr(self, '_unsloth_training_completed', False): - try: - import wandb - if wandb.run is not None: - wandb.finish() - # Reset HF's WandbCallback so it calls wandb.init() for the new run - for cb in self.callback_handler.callbacks: - if type(cb).__name__ == 'WandbCallback': - cb._initialized = False - break - except: - pass - # Enable training mode - _was_training = None - # Restore the GC mode the model was configured with at setup; fall back to - # the training args only when it wasn't recorded (issue #4735). Use hasattr, - # not a None sentinel, so a deliberately-recorded None is restored verbatim. - _model = getattr(self, 'model', None) - if hasattr(_model, '_unsloth_gradient_checkpointing'): - use_gc = _model._unsloth_gradient_checkpointing - else: - use_gc = getattr(self.args, 'gradient_checkpointing', True) - if hasattr(self, 'model') and hasattr(self.model, "training"): - _was_training = self.model.training - if hasattr(self, 'model') and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - output = f(self, *args, **kwargs) - # Restore previous mode when possible - if hasattr(self, 'model') and hasattr(self.model, "for_inference"): - if _was_training is False: - self.model.for_inference() - elif _was_training is True and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - # Reset gradient checkpointing buffers to free memory while staying ready for next run - try: - reset_unsloth_gradient_checkpointing_buffers() - except: - pass - # Mark that training completed so the next train() call can - # finish this W&B run before starting a new one - self._unsloth_training_completed = True - return output - return wrapper -pass - -torch_compile_options = { - "epilogue_fusion" : True, - "max_autotune" : False, - "shape_padding" : True, - "trace.enabled" : False, - "triton.cudagraphs" : False, -} - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_hidden_states_selective_log_softmax( - hidden_states: torch.Tensor, - lm_head: torch.Tensor, - index: torch.Tensor, - chunks: int = 4, - logit_scale_multiply: float = 0.0, - logit_scale_divide: float = 0.0, - logit_softcapping: float = 0.0, - temperature: float = 1.0, - # Rows per chunk cap. Read HERE, in the default, not in the body: the body - # is traced with fullgraph = True, and `os.environ` is an unsupported op - # there on torch 2.4 -- Dynamo raises - # torch._dynamo.exc.Unsupported: const method call bytes.decode - # from os._Environ.__getitem__, which the eager fallback does not catch - # (it only catches recompile-limit and disabled-hook breaks), so the very - # first call would die even with the variable unset. A default is evaluated - # once when the def runs, which is import time, outside any traced region. - # It is also a plain int argument, so Dynamo guards on it instead of - # constant-folding an unguarded read (2.7+ never notice a later change). - # A non-numeric value is ignored rather than raised on, so a typo cannot - # break the import. 0 keeps the previous chunk boundaries exactly. - max_rows_per_chunk: int = ( - int(os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip()) - if os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip().isdigit() - else 0 - ), - lm_head_bias: Optional[torch.Tensor] = None, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, - # Read in the default for the same Dynamo reason as max_rows_per_chunk. Pairs with vLLM's - # head_dtype=float32 that load_vllm requests under the same flag. - fp32_lm_head: bool = ( - os.environ.get("UNSLOTH_VLLM_FP32_LM_HEAD", "0").strip().lower() in ("1", "true", "on", "yes") - ), - # Older Unsloth scores old / ref log-probs here and the policy with the fused kernel: one scorer keeps ratio at 1. - fused_linear: bool = ( - os.environ.get("UNSLOTH_FUSED_LINEAR", "1") != "0" - and os.environ.get("UNSLOTH_GRPO_CHUNKED_LOGPROBS", "1") != "0" - ), -) -> torch.Tensor: - # All Unsloth Zoo code licensed under AGPL3 - from unsloth_zoo.fused_losses import chunked_logprob as _chunked_logprob - if (fused_linear and _chunked_logprob._FUSED_LINEAR and not fp32_lm_head - and support_top_k == 0 and support_top_p == 1.0 and support_min_p == 0.0 - and support_excluded_ids is None): - return _chunked_logprob.unsloth_chunked_logprobs( - hidden_states, lm_head, index, - temperature = temperature, - logit_scale_multiply = logit_scale_multiply, - logit_scale_divide = logit_scale_divide, - logit_softcapping = logit_softcapping, - lm_head_bias = lm_head_bias, - max_rows_per_chunk = max_rows_per_chunk if max_rows_per_chunk > 0 else None, - ).reshape((hidden_states.shape[0], hidden_states.shape[1])) - # Reshape on this tensor's own last dim: a no-op, so a wrong-width caller - # cannot have its row count silently rewritten and instead fails at the - # matmul below, which prints both operands. Do not swap in a bare - # torch._check: it reports only "Expected cond to be True", naming neither - # operand, and Dynamo rejects a message-carrying one. Callers dispatch on - # the width first -- see `compute_logprobs_chunk`, the packed path and - # `_pg_grad_forward`. - flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) - flat_index = index.reshape(-1) - - # Each chunk materialises rows x vocab logits and then a float32 copy of - # them, all on the device holding the output head. With a large vocabulary - # and a fixed chunk count that grows with the batch, so the peak scales with - # the batch rather than staying bounded. max_rows_per_chunk caps the rows - # per chunk instead, which is pure loop splitting: more, smaller chunks, - # same concatenated result. 0 (the default) keeps the previous chunk - # boundaries exactly. - if max_rows_per_chunk > 0: - n_rows = flat_hidden_states.shape[0] - chunks = max(chunks, -(-n_rows // max_rows_per_chunk)) - chunks = min(chunks, max(n_rows, 1)) - - # A one-row chunk's head gradient is a K=1 matmul Inductor mis-lowers under dynamic shapes (illegal - # memory access or NaN head gradients, torch 2.13), so keep chunks at 2+ rows when the head trains. - # Exactly `chunks` row ranges, so the unrolled loop length never depends on T: torch.chunk - # returns ceil(T / ceil(T / chunks)) pieces and recompiled per new count; tensor_split guards - # on T % chunks. At most rows // 2 chunks when the head trains keeps every piece at 2+ rows. - lm_head_grad = torch.is_grad_enabled() and lm_head.requires_grad - if lm_head_grad and flat_hidden_states.shape[0] < 2 * chunks: - chunks = max(flat_hidden_states.shape[0] // 2, 1) - - n_rows = flat_hidden_states.shape[0] - chunks = max(1, min(chunks, n_rows)) - bounds = [(i * n_rows) // chunks for i in range(chunks + 1)] - chunked_hidden_states = [flat_hidden_states[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - - # Cast once, outside the loop, so every chunk shares one fp32 head copy (also for autograd). - if fp32_lm_head: - lm_head = lm_head.to(torch.float32) - if lm_head_bias is not None: - lm_head_bias = lm_head_bias.to(torch.float32) - - all_per_token_logps = [] - - for chunk_hidden_states, chunk_index in zip(chunked_hidden_states, chunked_index): - # When the model is dispatched over several devices, the output head can - # sit on a different one from the hidden states, because accelerate - # places the tail of the model on the last device it fills. Co-locate on - # the head's device before the matmul, otherwise this raises - # Unhandled FakeTensor Device Propagation for aten.mm.default, - # found two different devices cuda:0, cuda:1 - # On a single device every .to() here is a no-op and the result is - # bit-identical to before. - chunk_hidden_states = chunk_hidden_states.to(device = lm_head.device, dtype = lm_head.dtype) - chunk_index = chunk_index.to(lm_head.device) - if fp32_lm_head: - # The GRPO loss runs this under bf16 / fp16 autocast, which would downcast the matmul. - with torch.autocast(device_type = chunk_hidden_states.device.type, enabled = False): - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - elif lm_head_bias is None: - chunk_logits = chunk_hidden_states @ lm_head.t() - else: - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - - if logit_scale_multiply != 0.0: - chunk_logits = chunk_logits * logit_scale_multiply - if logit_scale_divide != 0.0: - chunk_logits = chunk_logits / logit_scale_divide - if logit_softcapping != 0.0: - chunk_logits = logit_softcapping * torch.tanh(chunk_logits / logit_softcapping) - - chunk_logits = chunk_logits.to(torch.float32) - - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - per_token_logps = _truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - ) - else: - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - # Return to the caller's device so the concatenation below and every - # downstream consumer see the device they started on. - all_per_token_logps.append(per_token_logps.to(hidden_states.device)) - - all_per_token_logps = torch.concat(all_per_token_logps) - - all_per_token_logps = all_per_token_logps.reshape((hidden_states.shape[0], hidden_states.shape[1])) - return all_per_token_logps - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_selective_log_softmax( - logits, - index, - temperature: float = 1.0, - chunks: int = 4, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, -): - # Fixed `chunks` row ranges; see chunked_hidden_states_selective_log_softmax. - flat_logits = logits.reshape(-1, logits.shape[-1]) - flat_index = index.reshape(-1) - # Fewer rows than chunks: one row per range, as torch.chunk did, never empty ranges. - chunks = max(1, min(chunks, flat_logits.shape[0])) - bounds = [(i * flat_logits.shape[0]) // chunks for i in range(chunks + 1)] - chunked_logits = [flat_logits[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - all_per_token_logps = [] - # Per-chunk selective_log_softmax. - for chunk_logits, chunk_index in zip(chunked_logits, chunked_index): - chunk_logits = chunk_logits.to(torch.float32) - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - all_per_token_logps.append(_truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - )) - continue - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - all_per_token_logps.append(per_token_logps) - pass - all_per_token_logps = torch.concat(all_per_token_logps) - all_per_token_logps = all_per_token_logps.reshape((logits.shape[0], logits.shape[1])) - return all_per_token_logps - -def calculate_pad_tokens_in_prompt( - input_ids: torch.Tensor, - logits_to_keep: int, - pad_token_id: int -) -> torch.Tensor: - """Count left-padded tokens per sequence, e.g. [pad, pad, pad, cat] -> 3.""" - if logits_to_keep >= input_ids.shape[1]: - raise ValueError("logits_to_keep must be smaller than the sequence length.") - - prompt_section = input_ids[:, :-logits_to_keep] - - padding_mask = (prompt_section == pad_token_id) - - pad_token_counts = padding_mask.sum(dim=1) - - return pad_token_counts - -def create_completion_attention_mask( - completion_input_ids: torch.Tensor, - left_pad_tokens_per_prompt: torch.Tensor, - max_left_pad: int, - pad_token_id: int -) -> torch.Tensor: - """Build a completion mask that zeros leading prompt and trailing pad tokens. - - For [p,p,p,c,c,c,pad,pad,pad] (p=sliced prompt, c=completion, pad=padding) - this returns [0,0,0,1,1,1,0,0,0]. - """ - batch_size, completion_len = completion_input_ids.shape - device = completion_input_ids.device - - num_tokens_to_mask = max_left_pad - left_pad_tokens_per_prompt - - indices = torch.arange(completion_len, device=device).unsqueeze(0) - shift_mask = indices >= num_tokens_to_mask.unsqueeze(1) - - non_padding_mask = (completion_input_ids != pad_token_id) - - final_mask = shift_mask & non_padding_mask - - return final_mask - -def left_pack_padding(tensor: torch.Tensor, pad_id: int) -> torch.Tensor: - """Move all padding tokens in each sequence to the right.""" - mask = (tensor != pad_id) - # stable=True since the binary mask is unordered. - sorted_indices = torch.argsort(mask, dim=1, descending=True, stable=True) - packed_tensor = torch.gather(tensor, 1, sorted_indices) - return packed_tensor - -def align_logprobs_with_mask( - logprob_tensor: torch.Tensor, - attention_mask: torch.Tensor, - pad_value: float = 0.0 -) -> torch.Tensor: - """Align a log probability tensor with a given attention mask.""" - - device = logprob_tensor.device - batch_size, logprob_seq_len = logprob_tensor.shape - mask_seq_len = attention_mask.shape[1] - - padded_logprobs = torch.full( - attention_mask.shape, - fill_value=pad_value, - dtype=logprob_tensor.dtype, - device=device - ) - - left_pad_counts = torch.argmax(attention_mask, dim=1) - - cols = torch.arange(logprob_seq_len, device=device) - dest_indices = left_pad_counts.unsqueeze(1) + cols - - # Destination row indices, shape [batch_size, logprob_seq_len]. - row_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand_as(dest_indices) - - # Keep only in-bounds destinations, then scatter via advanced indexing. - valid_mask = dest_indices < mask_seq_len - valid_rows = row_indices[valid_mask] - valid_cols = dest_indices[valid_mask] - valid_vals = logprob_tensor[valid_mask] - padded_logprobs[valid_rows, valid_cols] = valid_vals - - return padded_logprobs - -def align_completion_tool_mask( - tool_mask: torch.Tensor, - completion_mask: torch.Tensor, -) -> torch.Tensor: - """Align a raw completion-length tool/env mask with Unsloth's repacked loss mask.""" - if tool_mask is None: - return completion_mask - if tool_mask.shape[0] != completion_mask.shape[0]: - raise ValueError("tool_mask batch size must match completion_mask batch size.") - - tool_mask = tool_mask.to(device=completion_mask.device) - if tool_mask.shape == completion_mask.shape: - aligned_tool_mask = tool_mask - else: - aligned_tool_mask = align_logprobs_with_mask( - tool_mask, - completion_mask, - pad_value=0, - ) - return completion_mask * aligned_tool_mask.to(dtype=completion_mask.dtype) - -def autotune_batch_and_chunks( - total_input_rows, - seq_len, - hidden_size, - vocab_size, - dtype_bytes=16, - multiplier=None -): - if multiplier is None: - final_m = max(4, seq_len // 4096) - else: - final_m = multiplier - - if torch.cuda.is_available(): - free_bytes, _ = torch.cuda.mem_get_info() - limit_gb = (free_bytes / (1024**3))*.80 - elif hasattr(torch, "xpu") and torch.xpu.is_available(): - # XPU: estimate free memory as total - reserved. - total_mem = torch.xpu.get_device_properties(0).total_memory - reserved_mem = torch.xpu.memory_reserved() - free_bytes = total_mem - reserved_mem - limit_gb = (free_bytes / (1024**3)) * 0.80 - else: - # Fallback: assume 8GB available. - limit_gb = 8.0 - - bytes_to_gb = 1024**3 - - b_vals = torch.arange(total_input_rows, 0, -1, device='cpu', dtype=torch.float32) - - hidden_gb = (b_vals * seq_len * hidden_size * dtype_bytes) / bytes_to_gb - - base_logits = ((b_vals/total_input_rows) * b_vals * seq_len * vocab_size * dtype_bytes) / bytes_to_gb - logits_gb = base_logits / final_m - - total_mem_gb = hidden_gb + logits_gb - - valid_mask = total_mem_gb <= limit_gb - valid_indices = torch.nonzero(valid_mask, as_tuple=False) - - if valid_indices.shape[0] == 0: - #This means your GPU will OOM - # Capped at the row count: unsloth's no-grad pass divides rows by this without max(1, ...). - return max(1, min(4, total_input_rows)), final_m - - best_idx = valid_indices[0].item() - final_b = int(b_vals[best_idx].item()) - - return final_b, final_m - -def sanitize_logprob(logprob): - """Local port of trl.scripts.vllm_serve.sanitize_logprob. - Filters NaN logprobs from vLLM outputs.""" - value = logprob.logprob - if math.isnan(value): - logging.getLogger(__name__).warning( - f"Generated NaN logprob, token logprob '{logprob}' will be ignored" - ) - return None - return value -@dataclass -class UnslothDistillationConfig(DistillationConfig): - """ - - Configuration class for the [`DistillationTrainer`]. - - Extends [`~transformers.TrainingArguments`] with parameters specific to knowledge distillation. All necessary - fields are declared here. - - Using [`~transformers.HfArgumentParser`] we can turn this class into - [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the - command line. - - Parameters: - > Parameters that control the model and the teacher model - - model_init_kwargs (`str` or `dict[str, Any]`, *optional*): - Keyword arguments for `AutoModelForCausalLM.from_pretrained`, used when the `model` argument of the trainer - is provided as a string. The `revision` value is also used when loading the processing class. - trust_remote_code (`bool`, *optional*, defaults to `False`): - Whether to allow loading models and tokenizers that ship custom Python code from the Hub. Forwarded to - [`~transformers.AutoModelForCausalLM.from_pretrained`] and [`~transformers.AutoTokenizer.from_pretrained`], - for both the student and teacher. - teacher_model_name_or_path (`str`, *optional*): - Model name or path for the teacher model. Used when the teacher is loaded locally. - teacher_model_revision (`str`, *optional*): - Model revision of the teacher model (e.g., branch name, tag, or commit hash). - teacher_model_init_kwargs (`str` or `dict[str, Any]`, *optional*): - Keyword arguments passed to `AutoModelForCausalLM.from_pretrained` when instantiating the teacher model - from a string. - disable_dropout (`bool`, *optional*, defaults to `False`): - Whether to disable dropout in the student model during training. - - > Parameters that control the data preprocessing - - remove_unused_columns (`bool`, *optional*, defaults to `False`): - Whether to only keep the column `"prompt"` in the dataset. The trainer consumes the raw prompt column and - generates completions on-policy, so it defaults to `False`. - max_completion_length (`int` or `None`, *optional*, defaults to `512`): - Maximum number of tokens to generate per completion during on-policy generation. - ds3_gather_for_generation (`bool`, *optional*, defaults to `True`): - This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for generation, - improving generation speed. However, disabling this option allows training models that exceed the VRAM - capacity of a single GPU, albeit at the cost of slower generation. Disabling this option is not compatible - with vLLM generation. - shuffle_dataset (`bool`, *optional*, defaults to `True`): - Whether to shuffle the training dataset. - pad_to_multiple_of (`int`, *optional*): - If set, the prompts ids and completions ids will be padded to a multiple of this value. - - > Parameters that control generation - - temperature (`float`, *optional*, defaults to `1.0`): - Temperature for sampling during generation and for computing the distillation loss. Higher values produce - softer probability distributions. - top_p (`float`, *optional*, defaults to `1.0`): - Top-p (nucleus) sampling parameter for on-policy generation. - top_k (`int`, *optional*, defaults to `0`): - Top-k sampling parameter for on-policy generation. `0` disables top-k filtering. - min_p (`float`, *optional*): - Minimum token probability, which will be scaled by the probability of the most likely token. It must be a - value between `0.0` and `1.0`. Typical values are in the `0.01-0.2` range. - generation_kwargs (`dict[str, Any]`, *optional*): - Additional keyword arguments to pass to [`~transformers.GenerationConfig`] (if using transformers) or - `SamplingParams` (if using vLLM) when sampling completions. This can be used to further customize the - generation behavior, such as setting `suppress_tokens`, `num_beams`, etc. If it contains keys that conflict - with the other generation parameters (like `min_p`, `top_p`, etc.), they will override them. - chat_template_kwargs (`dict[str, Any]`, *optional*): - Additional keyword arguments to pass to the `apply_chat_template` function when generating completions. - repetition_penalty (`float`, *optional*, defaults to `1.0`): - Float that penalizes new tokens based on whether they appear in the prompt and the generated text so far. - Values > `1.0` encourage the model to use new tokens, while values < `1.0` encourage the model to repeat - tokens. - cache_implementation (`str`, *optional*): - Implementation of the cache method for faster generation when `use_vllm` is set to `False`. - - > Parameters that control generation acceleration powered by vLLM - - use_vllm (`bool`, *optional*, defaults to `False`): - Whether to use vLLM for generating on-policy completions from the student model. - vllm_mode (`str`, *optional*, defaults to `"colocate"`): - Mode for student vLLM integration. Either `"server"` or `"colocate"`. - vllm_model_impl (`str`, *optional*, defaults to `"vllm"`): - Model implementation backend for vLLM. Use `"vllm"` or `"transformers"`. - vllm_enable_sleep_mode (`bool`, *optional*, defaults to `False`): - Enable vLLM sleep mode to offload student weights during the optimizer step. - vllm_structured_outputs_regex (`str`, *optional*): - Regex pattern for vLLM structured outputs. - - > Parameters that control the vLLM server (only used when `vllm_mode` is `"server"`) - - vllm_server_base_url (`str`, *optional*): - Base URL for the student vLLM server. If provided, `vllm_server_host` and `vllm_server_port` are ignored. - vllm_server_host (`str`, *optional*, defaults to `"0.0.0.0"`): - Host of the student vLLM server. - vllm_server_port (`int`, *optional*, defaults to `8000`): - Port of the student vLLM server. - vllm_server_timeout (`float`, *optional*, defaults to `240.0`): - Timeout for connecting to the student vLLM server. - vllm_group_port (`int`, *optional*, defaults to `51216`): - Port for the vLLM weight-update group (NCCL communicator). - - > Parameters that control colocated vLLM execution (only used when `vllm_mode` is `"colocate"`) - - vllm_gpu_memory_utilization (`float`, *optional*, defaults to `0.3`): - GPU memory utilization for the colocated student vLLM engine. - vllm_max_model_length (`int`, *optional*): - Maximum model sequence length for the colocated vLLM engine. - vllm_tensor_parallel_size (`int`, *optional*, defaults to `1`): - Tensor parallel size for the colocated student vLLM engine. - - > Parameters that control the training - - beta (`float`, *optional*, defaults to `1.0`): - Interpolation coefficient for the Generalized Jensen-Shannon Divergence loss. When `0.0`, the loss is the - forward KL divergence. When `1.0`, the loss is the reverse KL divergence. When `0.5`, it is the standard - JSD. Unlike GRPO's `beta` (a KL-penalty coefficient against a reference model), here it selects the - divergence itself; there is no reference-model KL penalty. - max_tool_calling_iterations (`int`, *optional*): - Maximum number of tool-calling turns when training an agent. If `None`, there is no limit and generation - stops when the model generates a response turn with no tool calls or when the total response length reaches - `max_model_length`. - - > Parameters that control the logging - - log_completions (`bool`, *optional*, defaults to `False`): - Whether to log a sample of (prompt, completion) pairs every `logging_steps` steps. If `rich` is installed, - it prints the sample. If `wandb` and/or `trackio` logging is enabled, it logs it to `wandb` and/or - `trackio`. - num_completions_to_print (`int`, *optional*): - Number of completions to print with `rich`. If `None`, all completions are logged. - log_unique_prompts (`bool`, *optional*, defaults to `False`): - Whether to log unique prompts. If `True`, only unique prompts are logged. If `False`, all prompts are - logged. - - > [!NOTE] - > These parameters have default values different from [`~transformers.TrainingArguments`]: - > - `logging_steps`: Defaults to `10` instead of `500`. - > - `gradient_checkpointing`: Defaults to `True` instead of `False`. - > - `bf16`: Defaults to `True` if `fp16` is not set, instead of `False`. - > - `learning_rate`: Defaults to `1e-6` instead of `5e-5`. - - """ - _unsloth_accepts_adamw_flash = True - vllm_sampling_params: Optional[Any] = field( - default = None, - metadata = {'help': 'vLLM SamplingParams'}, - ) - unsloth_num_chunks : Optional[int] = field( - default = -1, - metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'}, - ) - unsloth_logit_chunk_multiplier : Optional[int] = field( - default = None, - metadata = {'help': 'Multiplier for chunked logit computations.'}, - ) - unsloth_grpo_mini_batch : Optional[int] = field( - default = None, - metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'}, - ) - - def __init__( - self, - output_dir = None, - per_device_train_batch_size = 4, - num_train_epochs = 3.0, - max_steps = -1, - learning_rate = 5e-05, - lr_scheduler_type = 'linear', - lr_scheduler_kwargs = None, - warmup_steps = 0.1, - optim = 'adamw_8bit', - optim_args = None, - weight_decay = 0.001, - adam_beta1 = 0.9, - adam_beta2 = 0.999, - adam_epsilon = 1e-08, - optim_target_modules = None, - gradient_accumulation_steps = 2, - average_tokens_across_devices = True, - max_grad_norm = 1.0, - label_smoothing_factor = 0.0, - bf16 = False, - fp16 = False, - bf16_full_eval = False, - fp16_full_eval = False, - tf32 = None, - gradient_checkpointing = True, - gradient_checkpointing_kwargs = None, - torch_compile = False, - torch_compile_backend = None, - torch_compile_mode = None, - use_liger_kernel = False, - liger_kernel_config = None, - use_cache = False, - neftune_noise_alpha = None, - torch_empty_cache_steps = 250, - auto_find_batch_size = False, - logging_strategy = 'steps', - logging_steps = 1, - logging_first_step = False, - log_on_each_node = True, - logging_nan_inf_filter = False, - include_num_input_tokens_seen = False, - log_level = 'passive', - log_level_replica = 'warning', - disable_tqdm = None, - report_to = 'none', - run_name = None, - project = 'huggingface', - trackio_space_id = 'trackio', - eval_strategy = 'no', - eval_steps = None, - eval_delay = 0, - per_device_eval_batch_size = 4, - prediction_loss_only = False, - eval_on_start = False, - eval_do_concat_batches = True, - eval_use_gather_object = False, - eval_accumulation_steps = 2, - batch_eval_metrics = False, - save_only_model = False, - save_strategy = 'steps', - save_steps = 500, - save_on_each_node = False, - save_total_limit = None, - enable_jit_checkpoint = False, - push_to_hub = False, - hub_token = None, - hub_private_repo = None, - hub_model_id = None, - hub_strategy = 'every_save', - hub_always_push = False, - hub_revision = None, - load_best_model_at_end = False, - metric_for_best_model = None, - greater_is_better = None, - ignore_data_skip = False, - restore_callback_states_from_checkpoint = False, - full_determinism = False, - seed = 3407, - data_seed = None, - use_cpu = False, - accelerator_config = None, - parallelism_config = None, - dataloader_drop_last = False, - dataloader_num_workers = 0, - dataloader_pin_memory = True, - dataloader_persistent_workers = False, - dataloader_prefetch_factor = None, - remove_unused_columns = False, - label_names = None, - train_sampling_strategy = 'random', - length_column_name = 'length', - ddp_find_unused_parameters = None, - ddp_bucket_cap_mb = None, - ddp_broadcast_buffers = None, - ddp_backend = None, - ddp_timeout = 1800, - fsdp = None, - fsdp_config = None, - deepspeed = None, - debug = '', - skip_memory_metrics = True, - do_train = False, - do_eval = False, - do_predict = False, - resume_from_checkpoint = None, - warmup_ratio = None, - logging_dir = None, - local_rank = -1, - model_init_kwargs = None, - trust_remote_code = False, - teacher_model_name_or_path = None, - teacher_model_revision = None, - teacher_model_init_kwargs = None, - disable_dropout = False, - max_completion_length = 512, - ds3_gather_for_generation = True, - shuffle_dataset = True, - pad_to_multiple_of = None, - temperature = 1.0, - top_p = 1.0, - top_k = 0, - min_p = None, - generation_kwargs = {}, - chat_template_kwargs = None, - repetition_penalty = 1.0, - cache_implementation = None, - use_vllm = False, - vllm_mode = 'colocate', - vllm_model_impl = 'vllm', - vllm_enable_sleep_mode = False, - vllm_structured_outputs_regex = None, - vllm_server_base_url = None, - vllm_server_host = '0.0.0.0', - vllm_server_port = 8000, - vllm_server_timeout = 240.0, - vllm_group_port = 51216, - vllm_gpu_memory_utilization = 0.3, - vllm_max_model_length = None, - vllm_tensor_parallel_size = 1, - beta = 1.0, - max_tool_calling_iterations = None, - log_completions = False, - num_completions_to_print = None, - log_unique_prompts = False, - vllm_sampling_params = None, - unsloth_num_chunks = -1, - unsloth_logit_chunk_multiplier = None, - unsloth_grpo_mini_batch = None, - - **kwargs, - ): - if learning_rate < 1e-7: print(f'Unsloth: Your learning rate of `{learning_rate}` is too small and less than 1e-7! Consider increasing it, otherwise gradient updates will be close to 0!') - if learning_rate > 1: print(f'Unsloth: Your learning rate of `{learning_rate}` is way too larger > 1! Consider decreasing it to 1e-1, otherwise gradient updates will explode!') - _unsloth_fsdp2 = os.environ.get('FSDP_VERSION', '').strip() == '2' - _unsloth_fsdp_cfg = fsdp_config - if fsdp and isinstance(_unsloth_fsdp_cfg, str) and os.path.isfile(_unsloth_fsdp_cfg): - try: - import json as _unsloth_json - with open(_unsloth_fsdp_cfg, encoding = 'utf-8') as _unsloth_f: - _unsloth_fsdp_cfg = _unsloth_json.load(_unsloth_f) - except Exception: - _unsloth_fsdp_cfg = None - if fsdp and _unsloth_fsdp_cfg is None: - _unsloth_fsdp_cfg = {} - if fsdp and isinstance(_unsloth_fsdp_cfg, dict): - _unsloth_fsdp2 = _unsloth_fsdp2 or str(_unsloth_fsdp_cfg.get('fsdp_version', _unsloth_fsdp_cfg.get('version', 1))) == '2' - _unsloth_flash_adamw = str(getattr(optim, 'value', optim)).strip().lower() == 'adamw_flash' - if _unsloth_flash_adamw: optim = 'adamw_torch' - if _unsloth_fsdp2 and str(getattr(optim, 'value', optim)) == 'adamw_8bit': - print("Unsloth: FSDP2 shards parameters as DTensors, which bitsandbytes' 8-bit optimizers cannot step. Switching optim = 'adamw_8bit' to 'adamw_torch_fused'.") - optim = 'adamw_torch_fused' - if num_train_epochs is None: - num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override - if output_dir is None and save_strategy == 'steps' and save_steps == 500: - output_dir = 'unsloth_training_checkpoints' - save_strategy = 'no' - if os.environ.get('UNSLOTH_ENABLE_FLEX_ATTENTION', '0') == '1': - from unsloth_zoo.flex_attention import HAS_FLEX_ATTENTION - if HAS_FLEX_ATTENTION and pad_to_multiple_of is None: - from unsloth_zoo.flex_attention import FLEX_ATTENTION_BLOCK_SIZE - pad_to_multiple_of = FLEX_ATTENTION_BLOCK_SIZE - - if eval_steps is not None and eval_strategy != 'steps': - print(f'Unsloth: `eval_steps = {eval_steps}` is ignored because `eval_strategy` is {getattr(eval_strategy, "value", eval_strategy)!r}. Set `eval_strategy = "steps"` to evaluate every `eval_steps` steps.') - - if temperature <= 0: - raise ValueError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.') - elif temperature >= 10: - raise ValueError('Unsloth: Please set a positive non-zero temperature less than 10, since sampling will be quite erratic.') - - - # One dict so the filter sees the mirrored parameters AND `**kwargs`: - # filtering kwargs alone would double-bind any argument TRL renamed, - # since the new name is itself a mirrored parameter. - _unsloth_config_arguments = dict( - output_dir = output_dir, - per_device_train_batch_size = per_device_train_batch_size, - num_train_epochs = num_train_epochs, - max_steps = max_steps, - learning_rate = learning_rate, - lr_scheduler_type = lr_scheduler_type, - lr_scheduler_kwargs = lr_scheduler_kwargs, - warmup_steps = warmup_steps, - optim = optim, - optim_args = optim_args, - weight_decay = weight_decay, - adam_beta1 = adam_beta1, - adam_beta2 = adam_beta2, - adam_epsilon = adam_epsilon, - optim_target_modules = optim_target_modules, - gradient_accumulation_steps = gradient_accumulation_steps, - average_tokens_across_devices = average_tokens_across_devices, - max_grad_norm = max_grad_norm, - label_smoothing_factor = label_smoothing_factor, - bf16 = bf16, - fp16 = fp16, - bf16_full_eval = bf16_full_eval, - fp16_full_eval = fp16_full_eval, - tf32 = tf32, - gradient_checkpointing = gradient_checkpointing, - gradient_checkpointing_kwargs = gradient_checkpointing_kwargs, - torch_compile = torch_compile, - torch_compile_backend = torch_compile_backend, - torch_compile_mode = torch_compile_mode, - use_liger_kernel = use_liger_kernel, - liger_kernel_config = liger_kernel_config, - use_cache = use_cache, - neftune_noise_alpha = neftune_noise_alpha, - torch_empty_cache_steps = torch_empty_cache_steps, - auto_find_batch_size = auto_find_batch_size, - logging_strategy = logging_strategy, - logging_steps = logging_steps, - logging_first_step = logging_first_step, - log_on_each_node = log_on_each_node, - logging_nan_inf_filter = logging_nan_inf_filter, - include_num_input_tokens_seen = include_num_input_tokens_seen, - log_level = log_level, - log_level_replica = log_level_replica, - disable_tqdm = disable_tqdm, - report_to = report_to, - run_name = run_name, - project = project, - trackio_space_id = trackio_space_id, - eval_strategy = eval_strategy, - eval_steps = eval_steps, - eval_delay = eval_delay, - per_device_eval_batch_size = per_device_eval_batch_size, - prediction_loss_only = prediction_loss_only, - eval_on_start = eval_on_start, - eval_do_concat_batches = eval_do_concat_batches, - eval_use_gather_object = eval_use_gather_object, - eval_accumulation_steps = eval_accumulation_steps, - batch_eval_metrics = batch_eval_metrics, - save_only_model = save_only_model, - save_strategy = save_strategy, - save_steps = save_steps, - save_on_each_node = save_on_each_node, - save_total_limit = save_total_limit, - enable_jit_checkpoint = enable_jit_checkpoint, - push_to_hub = push_to_hub, - hub_token = hub_token, - hub_private_repo = hub_private_repo, - hub_model_id = hub_model_id, - hub_strategy = hub_strategy, - hub_always_push = hub_always_push, - hub_revision = hub_revision, - load_best_model_at_end = load_best_model_at_end, - metric_for_best_model = metric_for_best_model, - greater_is_better = greater_is_better, - ignore_data_skip = ignore_data_skip, - restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint, - full_determinism = full_determinism, - seed = seed, - data_seed = data_seed, - use_cpu = use_cpu, - accelerator_config = accelerator_config, - parallelism_config = parallelism_config, - dataloader_drop_last = dataloader_drop_last, - dataloader_num_workers = dataloader_num_workers, - dataloader_pin_memory = dataloader_pin_memory, - dataloader_persistent_workers = dataloader_persistent_workers, - dataloader_prefetch_factor = dataloader_prefetch_factor, - remove_unused_columns = remove_unused_columns, - label_names = label_names, - train_sampling_strategy = train_sampling_strategy, - length_column_name = length_column_name, - ddp_find_unused_parameters = ddp_find_unused_parameters, - ddp_bucket_cap_mb = ddp_bucket_cap_mb, - ddp_broadcast_buffers = ddp_broadcast_buffers, - ddp_backend = ddp_backend, - ddp_timeout = ddp_timeout, - fsdp = fsdp, - fsdp_config = fsdp_config, - deepspeed = deepspeed, - debug = debug, - skip_memory_metrics = skip_memory_metrics, - do_train = do_train, - do_eval = do_eval, - do_predict = do_predict, - resume_from_checkpoint = resume_from_checkpoint, - warmup_ratio = warmup_ratio, - logging_dir = logging_dir, - local_rank = local_rank, - model_init_kwargs = model_init_kwargs, - trust_remote_code = trust_remote_code, - teacher_model_name_or_path = teacher_model_name_or_path, - teacher_model_revision = teacher_model_revision, - teacher_model_init_kwargs = teacher_model_init_kwargs, - disable_dropout = disable_dropout, - max_completion_length = max_completion_length, - ds3_gather_for_generation = ds3_gather_for_generation, - shuffle_dataset = shuffle_dataset, - pad_to_multiple_of = pad_to_multiple_of, - temperature = temperature, - top_p = top_p, - top_k = top_k, - min_p = min_p, - generation_kwargs = generation_kwargs, - chat_template_kwargs = chat_template_kwargs, - repetition_penalty = repetition_penalty, - cache_implementation = cache_implementation, - use_vllm = use_vllm, - vllm_mode = vllm_mode, - vllm_model_impl = vllm_model_impl, - vllm_enable_sleep_mode = vllm_enable_sleep_mode, - vllm_structured_outputs_regex = vllm_structured_outputs_regex, - vllm_server_base_url = vllm_server_base_url, - vllm_server_host = vllm_server_host, - vllm_server_port = vllm_server_port, - vllm_server_timeout = vllm_server_timeout, - vllm_group_port = vllm_group_port, - vllm_gpu_memory_utilization = vllm_gpu_memory_utilization, - vllm_max_model_length = vllm_max_model_length, - vllm_tensor_parallel_size = vllm_tensor_parallel_size, - beta = beta, - max_tool_calling_iterations = max_tool_calling_iterations, - log_completions = log_completions, - num_completions_to_print = num_completions_to_print, - log_unique_prompts = log_unique_prompts,**kwargs) - super().__init__(**_unsloth_filter_config_init_kwargs(DistillationConfig, _unsloth_config_arguments, mirrored_from = __class__)) - self.vllm_sampling_params = vllm_sampling_params - self.unsloth_num_chunks = unsloth_num_chunks - if unsloth_grpo_mini_batch is not None: - if self.generation_batch_size >= unsloth_grpo_mini_batch: - self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch - else: - raise ValueError( - f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, " - f"which is self.per_device_train_batch_size * gradient_accumulation_steps." - ) - self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier - - # Unsloth: keep the reentrant checkpoint path - if getattr(self, 'gradient_checkpointing', False): - _gc_kwargs = getattr(self, 'gradient_checkpointing_kwargs', None) or {} - if _gc_kwargs.get('context_fn') is None and not _gc_kwargs.get('debug', False): - _gc_kwargs['use_reentrant'] = True - self.gradient_checkpointing_kwargs = _gc_kwargs - from unsloth.optimizers.flash_adamw import is_flash_adamw_marked, mark_flash_adamw, unmark_flash_adamw - if _unsloth_flash_adamw or is_flash_adamw_marked(self): - mark_flash_adamw(self) - else: - unmark_flash_adamw(self) - -pass - -class _UnslothDistillationTrainer(_BaseTrainer): - """""" - - loss_is_scaled_for_ga = True - - _tag_names = ["trl", "distillation"] - _name = "Distillation" - _paper = { - "title": "On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes", - "id": "2306.13649", - # docstyle-ignore - "citation": textwrap.dedent("""\ - @inproceedings{agarwal2024on-policy, - title = {{On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes}}, - author = {Rishabh Agarwal and Nino Vieillard and Yongchao Zhou and Piotr Stanczyk and Sabela Ramos Garea and Matthieu Geist and Olivier Bachem}, - year = 2024, - booktitle = {The Twelfth International Conference on Learning Representations, {ICLR} 2024, Vienna, Austria, May 7-11, 2024}, - publisher = {OpenReview.net}, - url = {https://openreview.net/forum?id=3zKtaqxLhW}, - }"""), - } - - def __init__( - self, - model: "str | PreTrainedModel | PeftModel", - teacher_model: str | PreTrainedModel | None = None, - args: DistillationConfig | None = None, - train_dataset: Dataset | None = None, - eval_dataset: Dataset | dict[str, Dataset] | None = None, - processing_class: PreTrainedTokenizerBase | ProcessorMixin | None = None, - callbacks: list[TrainerCallback] | None = None, - optimizers: tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None] = (None, None), - quantization_config: "BitsAndBytesConfig | None" = None, - peft_config: "PeftConfig | None" = None, - tools: list[Callable] | None = None, - ): - - if hasattr(model, 'vllm_engine') and hasattr(args, 'use_vllm'): - _unsloth_fsdp_cfg = getattr(args, 'fsdp_config', None) - if getattr(args, 'fsdp', None) and isinstance(_unsloth_fsdp_cfg, dict) and str(_unsloth_fsdp_cfg.get('version', _unsloth_fsdp_cfg.get('fsdp_version', 1))) == '2': - raise NotImplementedError('Unsloth: `fast_inference = True` (vLLM) is not supported with FSDP2.\nLoad the model with `fast_inference = False` to train with FSDP2, or launch without FSDP2 to use vLLM.') - if (getattr(args, 'use_vllm', False) == False): - args.use_vllm = True - if getattr(args, 'top_k', -1) is None or getattr(args, 'top_k', -1) == 0: - args.top_k = -1 - if args is None: - model_name = model if isinstance(model, str) else get_config_model_id(model.config) - model_name = model_name.split("/")[-1] - args = DistillationConfig(f"{model_name}-Distillation") - - # Student model loading - # PEFT initializes the adapter weights randomly, so set_seed must be done before creating the model to ensure - # reproducibility. - set_seed(args.seed) - # `_VALID_DICT_FIELDS` already parses any JSON-string form of these in `DistillationConfig.__post_init__`, so - # they are dicts [or None] here; copy so the setdefaults below don't mutate the config. - model_init_kwargs = dict(args.model_init_kwargs or {}) - teacher_model_init_kwargs = dict(args.teacher_model_init_kwargs or {}) - if isinstance(model, str): - model_name_or_path = model - if quantization_config is not None: - if "quantization_config" in model_init_kwargs: - raise ValueError( - "You set `quantization_config` both as a trainer argument and in `args.model_init_kwargs`. " - "Please set it in only one place, preferably as a trainer argument." - ) - model_init_kwargs["quantization_config"] = quantization_config - # Distributed training requires device_map=None ["auto" fails] - if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: - model_init_kwargs["device_map"] = None - model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) - model_revision = model_init_kwargs.get("revision") - model = create_model_from_path(model, **model_init_kwargs) - else: - model_name_or_path = get_config_model_id(model.config) - model_revision = None - if args.model_init_kwargs is not None: - logger.warning( - "You passed `model_init_kwargs` to the `DistillationConfig`, but your model is already " - "instantiated. The `model_init_kwargs` will be ignored." - ) - if quantization_config is not None: - logger.warning( - "You passed `quantization_config` to the trainer, but your model is already instantiated. The " - "`quantization_config` will be ignored." - ) - # Non-quantized models do not have the `is_loaded_in_{8,4}bit` attributes, whereas quantized models do - _is_quantized_model = getattr(model, "is_loaded_in_4bit", False) or getattr(model, "is_loaded_in_8bit", False) - - # Some models [SmolVLM/Idefics3] don't support `logits_to_keep` argument and error out if we pass it - # Inspect the forward method before we wrap the model with PEFT - self.model_kwarg_keys = ( - inspect.signature(model.forward).parameters.keys() - if not hasattr(model, "get_base_model") - else inspect.signature(model.get_base_model().forward).parameters.keys() - ) - - # Processing class - if processing_class is None: - processing_class = AutoProcessor.from_pretrained( - model_name_or_path, - revision=model_revision, - truncation_side="left", - padding_side="left", - trust_remote_code=args.trust_remote_code, - ) - - # Handle pad token for processors or tokenizers - if isinstance(processing_class, ProcessorMixin): - self._tokenizer = processing_class.tokenizer - self._is_vlm = True - elif isinstance(processing_class, PreTrainedTokenizerBase): - self._tokenizer = processing_class - self._is_vlm = False - else: - raise TypeError("The `processing_class` must be either a `PreTrainedTokenizerBase` or a `ProcessorMixin`") - - if self._tokenizer.pad_token is None: - self._tokenizer.pad_token = self._tokenizer.eos_token - - # The model must agree with the tokenizer on the pad token from construction, so mirror it onto the model - # configs. - model.config.get_text_config().pad_token_id = self._tokenizer.pad_token_id - model.generation_config.pad_token_id = self._tokenizer.pad_token_id - - # Every eos id the model declares [e.g. an end-of-turn token] ends a completion, not only the tokenizer's - eos_token_ids = model.generation_config.eos_token_id - if eos_token_ids is None: - eos_token_ids = [] - elif isinstance(eos_token_ids, int): - eos_token_ids = [eos_token_ids] - if self._tokenizer.eos_token_id not in eos_token_ids: - eos_token_ids = [self._tokenizer.eos_token_id, *eos_token_ids] - self.eos_token_ids = eos_token_ids - - # Resolve vision placeholder token IDs once. Used by the forward pass to rebuild mm_token_type_ids - # when tool responses inject images into the completion [see _generate forward_kwargs block]. - self._image_pad_token_id = None - self._video_pad_token_id = None - if self._is_vlm: - for candidate in ("<|image_pad|>", "<|image|>"): - tid = self._tokenizer.convert_tokens_to_ids(candidate) - if tid != self._tokenizer.unk_token_id: - self._image_pad_token_id = tid - break - tid = self._tokenizer.convert_tokens_to_ids("<|video_pad|>") - if tid != self._tokenizer.unk_token_id: - self._video_pad_token_id = tid - - # Tools - if tools: - if not Version(transformers.__version__) >= Version("5.0.0"): - raise ImportError( - "Using tools with DistillationTrainer requires transformers version 5.0.0 or higher. Please " - "upgrade transformers with `pip install --upgrade transformers` to use this feature." - ) - # jmespath is only needed by the legacy `response_schema` parser, which is all transformers < 5.13 ships. - # The new-style `response_template` parser doesn't use it, so don't require it on newer versions. - if not _SUPPORTS_RESPONSE_TEMPLATE and not is_jmespath_available(): - raise ImportError( - "Using tools with DistillationTrainer on transformers below 5.13.0 requires the jmespath library " - "for response parsing. Please install it with `pip install jmespath`, or upgrade transformers to " - "5.13.0 or higher, which doesn't need it." - ) - if not supports_tool_calling(processing_class): - raise ValueError( - "The provided chat template does not support tool calling. The template must be able to render a " - "full tool-calling conversation (user -> assistant with tool_calls -> tool)." - ) - # Async tools are deferred to a future PR; reject them now rather than fail cryptically in the tool loop. - if any(inspect.iscoroutinefunction(tool) for tool in tools): - raise ValueError("Async tools are not yet supported by `DistillationTrainer`. Pass synchronous tools.") - self.tools = tools or [] - # Tools are constant across examples, so the tool dict is built once here. - self._tool_dict = {tool.__name__: tool for tool in self.tools} - - # `add_response_schema` sets the response template [transformers >= 5.13] or legacy schema for known chat - # templates, so tool calls can be parsed. Skip if one is already set; warn if it's a migratable legacy schema. - if self.tools: - has_template = getattr(self._tokenizer, "response_template", None) is not None - has_schema = getattr(self._tokenizer, "response_schema", None) is not None - if not has_template and not has_schema: - processing_class = add_response_schema(processing_class) - elif has_schema and not has_template and _SUPPORTS_RESPONSE_TEMPLATE: - warnings.warn( - "The tokenizer has a legacy `response_schema` set but no `response_template`. The installed " - "transformers supports the new-style `response_template`; consider migrating, as `response_schema` " - "support will eventually be removed. See the Transformers response-parsing docs.", - FutureWarning, - ) - # In multi-turn training, the chat template *must* be prefix-preserving. If the tokenizer's original template - # isn't, we replace it at initialization with a training-safe, prefix-preserving template. - if self.tools and not is_chat_template_prefix_preserving(processing_class): - self.chat_template = get_training_chat_template(processing_class) - else: - self.chat_template = None - - # PEFT - if False: - if not is_peft_available(): - raise ImportError( - "You passed `peft_config` but the `peft` library is not installed. " - "Install it with `pip install trl[peft]`." - ) - if not isinstance(peft_config, PeftConfig): - raise TypeError( - f"`peft_config` must be a `peft.PeftConfig` instance (e.g. `peft.LoraConfig`), " - f"got {type(peft_config).__name__}." - ) - if is_peft_model(model): - raise ValueError( - "You passed a `PeftModel` instance together with a `peft_config` to the trainer. Please first merge " - "and unload the existing adapter, save the resulting base model, and then pass that base model along " - "with the new `peft_config` to the trainer." - ) - # Create PEFT model - # ZeRO-3 + PEFT for non-quantized models: - # - PEFT's default autocast_adapter_dtype=True upcasts LoRA adapter params to fp32 even when the base model is bf16. - # - ZeRO-3's _allgather_params_coalesced allocates output buffers using the dtype of the first persistent parameter, - # so mixed-dtype persistent_parameters [bf16 base + fp32 LoRA] cause a TypeError on the first optimizer step. - # - Passing autocast_adapter_dtype=False keeps adapter params in the base model dtype [bf16], fixing the mismatch. - # - This is safe: the fp32 upcast is a QLoRA-specific concern [low-bit quantized base models], not needed for - # non-quantized bf16 training. - # - See: - # - TRL issue: https://github.com/huggingface/trl/issues/6089 - # - Upstream issue: https://github.com/deepspeedai/DeepSpeed/issues/8072 - get_peft_model_kwargs = {} - if args.deepspeed_plugin is not None and args.deepspeed_plugin.zero_stage == 3 and not _is_quantized_model: - get_peft_model_kwargs["autocast_adapter_dtype"] = False - model = get_peft_model(model, peft_config, **get_peft_model_kwargs) - - # PEFT + DeepSpeed ZeRO-3 requires reentrant checkpointing. For more details, see - # https://github.com/huggingface/trl/issues/2514#issuecomment-2692152703. - # Can be removed once https://github.com/deepspeedai/DeepSpeed/pull/8130 is merged and released. - if ( - is_peft_model(model) - and args.deepspeed_plugin is not None - and args.deepspeed_plugin.zero_stage == 3 - and args.gradient_checkpointing - ): - args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} - use_reentrant = args.gradient_checkpointing_kwargs.get("use_reentrant") - if use_reentrant is False: - logger.warning( - "You are using PEFT with DeepSpeed ZeRO-3 and gradient checkpointing with `use_reentrant=False`. " - "`use_reentrant` is forced to `True` in this configuration to ensure correct training. To remove " - "this warning, unset `use_reentrant` in `gradient_checkpointing_kwargs` or set it to `True`." - ) - args.gradient_checkpointing_kwargs["use_reentrant"] = True - - # When using gradient checkpointing with PEFT, we need to enable input gradients. transformers.Trainer normally - # handles this, but a bug currently prevents it; see https://github.com/huggingface/transformers/issues/42489 - if is_peft_model(model) and args.gradient_checkpointing: - model.enable_input_require_grads() - - # When using QLoRA, the PEFT adapter weights are converted to bf16 to follow the recommendations from the - # original paper [see https://huggingface.co/papers/2305.14314, paragraph 3]. Normally, this can be done by - # passing `autocast_adapter_dtype=False` to `get_peft_model`, but this option is not yet supported for - # quantized models. See: https://github.com/huggingface/peft/issues/2889 - # The DoRA magnitude vector is excluded: unlike LoRA A/B, its optimizer updates can be smaller than bf16 can - # represent, silently freezing it, see #7268. - if False: - for name, param in model.named_parameters(): - if param.requires_grad and "lora_magnitude_vector" not in name: - param.data = param.data.to(torch.bfloat16) - - # The chunked JSD loss reads `lm_head.weight` directly and runs the backbone via - # `_get_last_hidden_state`, bypassing `PeftModel.forward[]` — so PEFT setups that live outside the backbone - # weights are silently ignored. Fail loudly rather than train on a silently-wrong objective. - if is_peft_model(model): - # When the LM head is targeted by a PEFT adapter [`"lm_head"` in `target_modules`], `lm_head.weight` is the - # frozen base weight and the trainable adapter lives in separate submodules the loss never sees, so the head - # adapter would receive no gradient. - output_embeddings = model.get_output_embeddings() - if isinstance(output_embeddings, BaseTunerLayer): - raise ValueError( - "Applying a PEFT adapter to `lm_head` is not supported. The distillation loss reads " - "`lm_head.weight` directly, so the adapter on the head is ignored and never trained. Remove " - "`'lm_head'` from your `target_modules`." - ) - # Prompt-learning methods [PromptTuning, PrefixTuning, P-Tuning] inject virtual tokens via - # `PeftModel.forward[]`, which the loss bypasses by calling the backbone directly, so virtual tokens are - # never prepended and the loss is computed on the wrong sequence. - if any(isinstance(cfg, PromptLearningConfig) for cfg in model.peft_config.values()): - raise ValueError( - "Prompt-learning PEFT methods (PromptTuning, PrefixTuning, P-Tuning) are not supported. The " - "distillation loss bypasses `PeftModel.forward()` by calling the backbone directly, so virtual " - "tokens are never prepended and the loss is computed on the wrong sequence. Use a weight-based " - "adapter such as LoRA instead." - ) - - # The chunked JSD loss calls the student backbone directly, bypassing the DDP/FSDP wrapper's forward; route it - # through `_forward_redirection` so `prepare_for_backward[]` still fires. - self._forward_redirection = _ForwardRedirection() - - # Teacher model setup - # `teacher_model` may be None: subclasses [e.g. ServerDistillationTrainer] supply the teacher another way. - if teacher_model is not None: - if isinstance(teacher_model, str): - dtype = teacher_model_init_kwargs.get("dtype") - teacher_model_init_kwargs["dtype"] = dtype if dtype in ["auto", None] else getattr(torch, dtype) - if args.teacher_model_revision is not None: - teacher_model_init_kwargs.setdefault("revision", args.teacher_model_revision) - # Distributed training requires device_map=None ["auto" fails] - if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: - teacher_model_init_kwargs["device_map"] = None - teacher_model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) - teacher_model = create_model_from_path(teacher_model, **teacher_model_init_kwargs) - elif args.teacher_model_init_kwargs is not None: - raise ValueError( - "You passed teacher_model_init_kwargs to the config, but your teacher_model is already " - "instantiated." - ) - - # Iterable datasets can't be indexed, so the RepeatSampler can't be attached to them. Instead, the sampler's - # ordering is reproduced by streaming [see `get_train_dataloader`/`get_eval_dataloader` and - # `repeat_iterable_dataset`]. This requires `dispatch_batches=False`: with the default dispatch path, batches - # are collated on the main process and Accelerate tries to concatenate the string `prompt` column, which fails; - # `dispatch_batches=False` also lets each process shard the stream into contiguous slices, as the sampler does. - # See https://github.com/huggingface/trl/issues/3213 - uses_iterable_dataset = ( - isinstance(train_dataset, IterableDataset) - or isinstance(eval_dataset, IterableDataset) - or ( - isinstance(eval_dataset, dict) and any(isinstance(ds, IterableDataset) for ds in eval_dataset.values()) - ) - ) - if uses_iterable_dataset: - if args.accelerator_config.dispatch_batches: - raise ValueError( - "Iterable datasets require `dispatch_batches=False`, but it is set to `True` in " - "`accelerator_config`. Please set it to `False`." - ) - args.accelerator_config.dispatch_batches = False - # An iterable train set bakes the generation-batch repeats into the stream, so it must be read by a single - # worker: multiple workers would shard and interleave it, breaking the generation-batch ordering that - # `_prepare_inputs` relies on. Map-style train keeps its workers. - if isinstance(train_dataset, IterableDataset) and args.dataloader_num_workers != 0: - logger.warning( - f"Iterable datasets require `dataloader_num_workers=0` to preserve prompt grouping; overriding the " - f"provided value ({args.dataloader_num_workers})." - ) - args.dataloader_num_workers = 0 - - super().__init__( - model=model, - args=args, - data_collator=identity, # No data collation is needed in distillation - train_dataset=train_dataset, - eval_dataset=eval_dataset, - processing_class=processing_class, - callbacks=callbacks, - optimizers=optimizers, - ) - - # With several GPUs visible and no distributed launcher, `Trainer` wraps the model in `nn.DataParallel`, whose - # replicas would all run the forward bound to the original model. Checked after `Trainer.__init__`, which sets - # `n_gpu` to 1 for a model split across devices with `device_map`. - if args.n_gpu > 1: - raise ValueError( - f"{type(self).__name__} does not support `nn.DataParallel`, which `Trainer` uses when several GPUs are " - "visible to a single process. Launch the script with `accelerate launch` or `torchrun`, or make a " - "single GPU visible with `CUDA_VISIBLE_DEVICES`." - ) - - self._dist = DistributedBackend(self.accelerator) - - # Add tags to the model - self.model.add_model_tags(self._tag_names) - - # Prepare teacher model [after super[].__init__ so accelerator is ready] - if teacher_model is not None: - # The divergence compares the full next-token distribution of the student against the teacher's, so both - # must be defined over the same vocabulary. - student_vocab_size = self.model.config.get_text_config().vocab_size - teacher_vocab_size = teacher_model.config.get_text_config().vocab_size - if student_vocab_size != teacher_vocab_size: - raise ValueError( - f"The student model has vocab_size {student_vocab_size} but the teacher model has vocab_size " - f"{teacher_vocab_size}. Distillation compares the teacher's full next-token distribution, which " - f"requires a shared vocabulary. Use a teacher with the same vocab_size, or GOLD for " - f"cross-tokenizer distillation." - ) - if self.is_deepspeed_enabled: - self.teacher_model = prepare_deepspeed(teacher_model, self.accelerator) - else: - self.teacher_model = self.accelerator.prepare_model(teacher_model, evaluation_mode=True) - else: - self.teacher_model = None - - if args.disable_dropout: - disable_dropout_in_model(self.model) - - # Store config values - self.beta = args.beta - self.temperature = args.temperature - self.top_p = args.top_p - self.top_k = args.top_k - self.min_p = args.min_p - self.repetition_penalty = args.repetition_penalty - self.max_completion_length = args.max_completion_length - self.max_tool_calling_iterations = ( - args.max_tool_calling_iterations if args.max_tool_calling_iterations is not None else sys.maxsize - ) - self.chat_template_kwargs = args.chat_template_kwargs or {} - self.pad_to_multiple_of = args.pad_to_multiple_of - self.shuffle_dataset = args.shuffle_dataset - - # Tracks the number of iterations [forward + backward passes], including those within a grad accum cycle - self._step = 0 - # Buffer the batch to reuse generated outputs across multiple updates. For more details, see - # `_get_train_sampler` and `_prepare_inputs`. - self._buffered_inputs = None - - # Ensure each process receives a unique seed so different processes generate different completions when - # generating with transformers. We could skip it if we use vLLM, but it's safer to set it in all cases. - set_seed(args.seed, device_specific=True) - - # Generation config - generation_kwargs = { - "max_new_tokens": self.max_completion_length, - "do_sample": True, - "pad_token_id": self._tokenizer.pad_token_id, - "bos_token_id": self._tokenizer.bos_token_id, - "eos_token_id": self.eos_token_ids, - "temperature": self.temperature, - "top_p": self.top_p, - "top_k": self.top_k, - "min_p": self.min_p, - "repetition_penalty": self.repetition_penalty, - "cache_implementation": args.cache_implementation, - } - if args.generation_kwargs is not None: - generation_kwargs.update(args.generation_kwargs) - self.generation_config = GenerationConfig(**generation_kwargs, disable_compile=True) - # Keep training-specific generation kwargs to overwrite model's original generation config - self.generation_kwargs = generation_kwargs - - # Metrics & Logging - self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} - self._total_train_tokens = 0 - self._current_train_step_time = 0.0 - self.log_completions = args.log_completions - self.log_unique_prompts = args.log_unique_prompts - self.num_completions_to_print = args.num_completions_to_print - # Keep logs sized to the generation batch to record only outputs from the latest model update. - generation_batch_size = ( - args.per_device_train_batch_size * self.accelerator.num_processes * args.gradient_accumulation_steps - ) - self._logs = { - "prompt": deque(maxlen=generation_batch_size), - "completion": deque(maxlen=generation_batch_size), - } - - if self.accelerator.is_main_process and self.log_completions: - os.makedirs(os.path.join(self.args.output_dir, "completions"), exist_ok=True) - - # vLLM for student generation - self.use_vllm = args.use_vllm - self.vllm_mode = args.vllm_mode - if self.use_vllm: - if not is_vllm_available(): - raise ImportError( - "vLLM is not available and `use_vllm` is set to True. Please install vLLM with " - "`pip install trl[vllm]` to use it." - ) - self.vllm_generation = VLLMGeneration( - model=self.model, - accelerator=self.accelerator, - processing_class=self.processing_class, - # vLLM configuration - mode=args.vllm_mode, - structured_outputs_regex=args.vllm_structured_outputs_regex, - # Server mode configuration - server_base_url=args.vllm_server_base_url, - server_host=args.vllm_server_host, - server_port=args.vllm_server_port, - group_port=args.vllm_group_port, - server_timeout=args.vllm_server_timeout, - # Colocate mode configuration - tensor_parallel_size=args.vllm_tensor_parallel_size, - gpu_memory_utilization=args.vllm_gpu_memory_utilization, - max_model_length=args.vllm_max_model_length, - max_num_seqs=args.per_device_train_batch_size - * args.vllm_tensor_parallel_size - * args.gradient_accumulation_steps, - enable_sleep_mode=args.vllm_enable_sleep_mode, - model_impl=args.vllm_model_impl, - trust_remote_code=args.trust_remote_code, - # Generation configuration - repetition_penalty=self.repetition_penalty, - temperature=self.temperature, - top_p=self.top_p, - top_k=self.top_k, - min_p=self.min_p, - max_completion_length=self.max_completion_length, - logprobs=None, # distillation trains on the teacher distribution, not sampled logprobs - generation_kwargs=args.generation_kwargs, - ) - self._last_loaded_step = -1 # tag to avoid useless loading during grad accumulation - - def _set_signature_columns_if_needed(self): - # If `self.args.remove_unused_columns` is True, non-signature columns are removed. - # By default, this method sets `self._signature_columns` to the model's expected inputs (usually, "input_ids" - # and "attention_mask"). In DistillationTrainer, we preprocess data, so using the model's signature columns - # doesn't work. Instead, we set them to the columns expected by the `training_step` method, hence the override. - if self._signature_columns is None: - self._signature_columns = ["prompt", "image", "images"] - - # Instead of returning a standard per-step batch (i.e., `per_device_batch_size), our dataloader loads an - # *generation* batch (i.e., `per_device_batch_size × gradient_accumulation_steps`). This allows us to generate - # completions once every gradient_accumulation_steps step—rather than once per accumulation step—which is - # significantly more efficient. Thus, `_prepare_inputs` is called with this *generation* batch, and it handles the - # splitting internally. - # Maintenance note: this method is a copy-paste of the original `Trainer.get_train_dataloader` with two changes: the - # batch size is multiplied by `gradient_accumulation_steps`, and iterable datasets are wrapped (see - # `repeat_iterable_dataset`). - def get_train_dataloader(self): - dataset = self.train_dataset - if isinstance(dataset, IterableDataset): - # Iterable datasets can't be indexed, so RepeatSampler can't be attached. Reproduce its ordering by - # transforming the stream instead (see `repeat_iterable_dataset`). The full permutation done by - # RepeatSampler becomes a buffered shuffle here. - if self.shuffle_dataset: - dataset = dataset.shuffle(seed=self.args.seed) - # Effective training batch = the generation batch (the deferred `generation_batch_size` slot-in). - generation_batch_size = ( - self.args.per_device_train_batch_size - * self.accelerator.num_processes - * self.args.gradient_accumulation_steps - ) - dataset = repeat_iterable_dataset( - dataset, - mini_repeat_count=1, - batch_size=generation_batch_size, - repeat_count=self.args.gradient_accumulation_steps, - ) - return self._get_dataloader( - dataset=dataset, - description="Training", - batch_size=self._train_batch_size * self.args.gradient_accumulation_steps, # < this is the change - sampler_fn=self._get_train_sampler, - is_training=True, - ) - - def _get_train_sampler(self, dataset: Dataset | None = None) -> Sampler: - # Repeat each generation batch `gradient_accumulation_steps` times so the completions generated once per - # generation batch (see `_prepare_inputs`) are reused across the accumulation window. Distillation is n=1, - # so there is no per-prompt repeat (`mini_repeat_count=1`). - if dataset is None: - dataset = self.train_dataset - # The generation batch is the full effective training batch. GRPO names it `generation_batch_size` (paired with - # the deferred `steps_per_generation`); both are deferred here, so derive it inline — slot-in point. - generation_batch_size = ( - self.args.per_device_train_batch_size - * self.accelerator.num_processes - * self.args.gradient_accumulation_steps - ) - return RepeatSampler( - data_source=dataset, - mini_repeat_count=1, - batch_size=generation_batch_size, - repeat_count=self.args.gradient_accumulation_steps, - shuffle=self.shuffle_dataset, - seed=self.args.seed, - ) - - def _get_eval_sampler(self, eval_dataset) -> Sampler: - # See _get_train_sampler for an explanation of the sampler. - return RepeatSampler( - data_source=eval_dataset, - mini_repeat_count=1, - seed=self.args.seed, - ) - - # This method overrides `Trainer.get_eval_dataloader` to wrap iterable eval datasets, reproducing the - # RepeatSampler ordering that can't be attached to them (see `get_train_dataloader`). Map-style datasets keep the - # default path via `_get_eval_sampler`, which shuffles with `seed`, so the iterable wrap shuffles too (buffered) - # to walk prompts in a matching order. - # Maintenance note: this method is a copy-paste of the original `Trainer.get_eval_dataloader`, with the iterable - # wrapping as the only addition. - def get_eval_dataloader(self, eval_dataset: str | Dataset | IterableDataset | None = None) -> DataLoader: - if eval_dataset is None and self.eval_dataset is None: - raise ValueError("Trainer: evaluation requires an eval_dataset.") - - # If we have persistent workers, don't do a fork bomb especially as eval datasets - # don't change during training - dataloader_key = eval_dataset if isinstance(eval_dataset, str) else "eval" - if ( - hasattr(self, "_eval_dataloaders") - and dataloader_key in self._eval_dataloaders - and self.args.dataloader_persistent_workers - ): - return self._eval_dataloaders[dataloader_key] - - eval_dataset = ( - self.eval_dataset[eval_dataset] - if isinstance(eval_dataset, str) - else eval_dataset - if eval_dataset is not None - else self.eval_dataset - ) - - if isinstance(eval_dataset, IterableDataset): - # Apply the `__init__` iterable config here too - if self.args.accelerator_config.dispatch_batches: - raise ValueError( - "Iterable datasets require `dispatch_batches=False`, but it is set to `True` in " - "`accelerator_config`. Please set it to `False`." - ) - self.accelerator.dataloader_config.dispatch_batches = False - eval_dataset = eval_dataset.shuffle(seed=self.args.seed) - eval_dataset = repeat_iterable_dataset(eval_dataset, mini_repeat_count=1) - # Force a single worker for this loader only, without persisting the change - num_workers = self.args.dataloader_num_workers - self.args.dataloader_num_workers = 0 - - try: - return self._get_dataloader( - dataset=eval_dataset, - description="Evaluation", - batch_size=self.args.eval_batch_size, - sampler_fn=self._get_eval_sampler, - dataloader_key=dataloader_key, - ) - finally: - if isinstance(eval_dataset, IterableDataset): - self.args.dataloader_num_workers = num_workers - - def _tokenize_prompts(self, prompts: list): - """Tokenize prompts and extract images/multimodal fields for generation.""" - if is_conversational({"prompt": prompts[0]}): - # Normalize string content to content blocks for VLM processors that don't handle plain strings. - if self._is_vlm: - prompts = [prepare_multimodal_messages(prompt) for prompt in prompts] - - # Extract images from messages for VLM support - images = [] - has_images = False - for prompt in prompts: - prompt_images = [] - for message in prompt: - if isinstance(message["content"], list): - for part in message["content"]: - if part["type"] == "image": - prompt_images.append(part["image"]) - has_images = True - images.append(prompt_images if prompt_images else None) - images = images if has_images else None - - # Workaround for a bug in transformers 5.3.0 where some processors (e.g. Qwen2.5-VL) crash on - # batched unpadded input (transformers#44514). - # Fixed in transformers 5.4.0 (transformers#44563). - needs_padding_workaround = Version("5.3.0") <= Version(transformers.__version__) < Version("5.4.0") - tokenized = self.processing_class.apply_chat_template( - conversation=prompts, - tools=self.tools or None, # `or None`: Llama bug: it renders tool boilerplate for tools=[] - chat_template=self.chat_template, - add_generation_prompt=True, - tokenize=True, - return_dict=True, - **({"padding": True} if needs_padding_workaround else {}), - **self.chat_template_kwargs, - ) - if needs_padding_workaround: - # Unpad input_ids: remove padding tokens using attention_mask to get per-sequence lists - prompt_ids = [ - [tok for tok, m in zip(ids, mask, strict=True) if m] - for ids, mask in zip(tokenized["input_ids"], tokenized["attention_mask"], strict=True) - ] - else: - prompt_ids = tokenized["input_ids"] - # For VLMs, the processor returns extra multimodal fields (pixel_values, image_grid_thw, etc.) - multimodal_fields = {k: v for k, v in tokenized.items() if k not in ("input_ids", "attention_mask")} - else: - prompt_ids = self.processing_class(text=prompts)["input_ids"] - images = None - multimodal_fields = {} - return prompt_ids, images, multimodal_fields - - def _generate_single_turn(self, prompt_ids, images, multimodal_fields, has_tool_images=False): - device = self.accelerator.device - - # Generate completions using either vLLM or regular generation - if self.use_vllm: - # Sync weights if training step changed - if self.state.global_step != self._last_loaded_step: - with profiling_context(self, "sync_weights"): - self.vllm_generation.sync_weights() - self._last_loaded_step = self.state.global_step - - # Generate using vLLM with raw token IDs. Distillation is n=1 and uses the teacher distribution rather than - # sampled logprobs, so we request one completion per prompt and discard vLLM's logprobs. - _, completion_ids, _, _ = self.vllm_generation.generate( - prompts=prompt_ids, - images=images, - num_generations=1, - profiler=profiling_context(self, "vLLM.generate"), - ) - - else: - # Regular generation path: left-pad token IDs into tensors - prompt_tensors = [torch.tensor(ids) for ids in prompt_ids] - padded_ids = pad(prompt_tensors, padding_value=self._tokenizer.pad_token_id, padding_side="left") - attention_mask = pad([torch.ones_like(t) for t in prompt_tensors], padding_value=0, padding_side="left") - generate_inputs = {"input_ids": padded_ids, "attention_mask": attention_mask} - # For VLMs, include multimodal fields as tensors (pixel_values, image_grid_thw, etc.) - for k, v in multimodal_fields.items(): - if isinstance(v, torch.Tensor): - generate_inputs[k] = v - elif isinstance(v, list) and v and isinstance(v[0], list): - # Per-token field (e.g., token_type_ids): left-pad like input_ids - generate_inputs[k] = pad([torch.tensor(x) for x in v], padding_value=0, padding_side="left") - else: - generate_inputs[k] = torch.tensor(np.array(v)) - - # For VLM tool images: build token type IDs from the padded input IDs. - if self._is_vlm and self.tools and has_tool_images: - mm_ids = torch.zeros_like(padded_ids) - if self._image_pad_token_id is not None: - mm_ids[padded_ids == self._image_pad_token_id] = 1 - if self._video_pad_token_id is not None: - mm_ids[padded_ids == self._video_pad_token_id] = 2 - - # Use the same key the model expects: token_type_ids for models like Gemma, - # mm_token_type_ids for models like Qwen. - if "image_grid_thw" in generate_inputs: - generate_inputs["mm_token_type_ids"] = mm_ids - else: - generate_inputs["token_type_ids"] = mm_ids - - generate_inputs = super()._prepare_inputs(generate_inputs) - - with ( - profiling_context(self, "transformers.generate"), - unwrap_model_for_generation( - self.model_wrapped, - self.accelerator, - gather_deepspeed3_params=self.args.ds3_gather_for_generation, - generation_kwargs=self.generation_kwargs, # Override model.generation_config with generation_kwargs to fix transformers#42762 - ) as unwrapped_model, - torch.no_grad(), - self._dist.summon_full_params(self.model_wrapped, recurse=False), - ): - prompt_completion_ids = unwrapped_model.generate( - **generate_inputs, generation_config=self.generation_config - ) - # Compute prompt length and extract completion ids - prompt_length = generate_inputs["input_ids"].size(1) - completion_ids = prompt_completion_ids[:, prompt_length:] - - # Mask everything after the first EOS token - is_eos = torch.isin(completion_ids, torch.tensor(self.eos_token_ids, device=device)) - eos_idx = torch.full((is_eos.size(0),), is_eos.size(1), dtype=torch.long, device=device) - eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)] - sequence_indices = torch.arange(is_eos.size(1), device=device).expand(is_eos.size(0), -1) - completion_mask = (sequence_indices <= eos_idx.unsqueeze(1)).int() - completion_ids = [ - c[m].tolist() for c, m in zip(completion_ids.cpu(), completion_mask.bool().cpu(), strict=True) - ] - - return completion_ids - - def _get_tool_suffix_ids(self, tool_messages): - """Get token IDs for tool result formatting by using a minimal dummy conversation.""" - # Use the real tool name instead of a dummy: some templates (e.g. GPT-OSS) derive the tool response - # header from the assistant's tool call name. - dummy_tool_calls = [{"type": "function", "function": {"name": tool_messages[0]["name"], "arguments": {}}}] - dummy_messages = [ - {"role": "user", "content": "dummy"}, - { - "role": "assistant", - # "content" is required here because VLM processors crash on tokenize=True without it - # (KeyError in processing_utils.py). See huggingface/transformers#45290. - "content": "", - "tool_calls": dummy_tool_calls, - }, - ] - if self._is_vlm: - dummy_messages = prepare_multimodal_messages(dummy_messages) - tool_messages = prepare_multimodal_messages(tool_messages) - - prefix_ids = self.processing_class.apply_chat_template( - dummy_messages, - add_generation_prompt=False, - tokenize=True, - chat_template=self.chat_template, - return_dict=False, - **self.chat_template_kwargs, - ) - full_ids = self.processing_class.apply_chat_template( - dummy_messages + tool_messages, - add_generation_prompt=True, - tokenize=True, - chat_template=self.chat_template, - return_dict=False, - **self.chat_template_kwargs, - ) - # VLM processors return batched output (list of lists), unbatch for single conversation - if self._is_vlm: - prefix_ids = prefix_ids[0] - full_ids = full_ids[0] - - # Some chat templates (notably Qwen3/Qwen3.5) render "...<|im_end|>\n" after an assistant/tool block. - # When we compute `suffix_ids` by slicing `full_ids`, we must align the slicing boundary to - # EOS (not EOS + newline). Templates that don't use EOS as end-of-turn (e.g. Gemma uses - # ) skip this trimming. - eos_positions = [i for i, tok_id in enumerate(prefix_ids) if tok_id == self._tokenizer.eos_token_id] - if eos_positions: - prefix_ids = prefix_ids[: eos_positions[-1] + 1] - - if full_ids[: len(prefix_ids)] != prefix_ids: - raise ValueError("Unexpected tokenization: the EOS-trimmed prefix IDs are not a prefix of the full IDs.") - return full_ids[len(prefix_ids) :] - - def _tool_call_loop(self, prompts, prompt_ids, completion_ids, completions, images, multimodal_fields): - # Tool execution loop: execute tools, then regenerate completions with tool results appended to the prompt - tool_calls = [completion[0].get("tool_calls") for completion in completions] - idxs_with_tool = [idx for idx, tool_call in enumerate(tool_calls) if tool_call] - tool_calls = [tool_calls[idx] for idx in idxs_with_tool] - tool_mask = [[1] * len(ids) for ids in completion_ids] # 0 for tool result tokens, 1 elsewhere - # Collect images from multimodal tool responses for the forward pass - tool_images = [[] for _ in completion_ids] - tool_call_count = 0 - tool_failure_count = 0 - iteration_num = 0 - - while idxs_with_tool and iteration_num < self.max_tool_calling_iterations: - prompt_completion_tools = [prompts[i] for i in idxs_with_tool] # select only prompts that need tool calls - # Snapshot state so we can rollback tool results that would exceed max_completion_length - completions_len_before = [len(completions[i]) for i in idxs_with_tool] - tool_images_len_before = [len(tool_images[i]) for i in idxs_with_tool] - prompts_len_before = [len(prompts[i]) for i in idxs_with_tool] - - # Call the tools, and build the new prompt for generation - for idx in range(len(idxs_with_tool)): - idx_with_tool = idxs_with_tool[idx] - tool_call_list = tool_calls[idx] - prompt_completion_tool = prompt_completion_tools[idx] - # Append the last assistant message (which triggered tool_calls) to the prompt - prompt_completion_tool.append(completions[idx_with_tool][-1]) - tool_call_results = [] - for tool_call in tool_call_list: - tool_call_count += 1 - if tool_call["type"] == "function": - function = tool_call["function"] - name = function["name"] - try: - if name in self._tool_dict: - tool_call_results.append((name, self._tool_dict[name](**function["arguments"]))) - else: - raise ValueError(f"Tool {name} not found.") - except Exception as e: - tool_failure_count += 1 - result = {"error": str(e)} - tool_call_results.append((name, result)) - else: - tool_failure_count += 1 - name = tool_call.get("name", "unknown") - tool_call_results.append((name, {"error": f"Unsupported tool call type: {tool_call['type']}"})) - - for name, result in tool_call_results: - # Support multimodal tool responses: if the tool returns a list of content blocks - # (e.g., [{"type": "image", "image": ...}, {"type": "text", "text": "..."}]), - # pass them through directly so _tokenize_prompts can extract images for VLMs. - content = result if isinstance(result, list) else str(result) - tool_message = {"role": "tool", "name": name, "content": content} - # Collect images from multimodal tool responses - if isinstance(content, list): - for part in content: - if isinstance(part, dict) and part.get("type") == "image": - tool_images[idx_with_tool].append(part["image"]) - prompt_completion_tool.append(tool_message) - completions[idx_with_tool].append(tool_message) - - # Build token IDs by concatenation: prompt + completion + tool_suffix. - prompt_completion_tool_ids = [] - for idx in range(len(idxs_with_tool)): - idx_with_tool = idxs_with_tool[idx] - # Extract trailing tool messages from completions - tool_messages = [] - for message in reversed(completions[idx_with_tool]): - if message["role"] == "tool": - tool_messages.insert(0, message) - else: - break - suffix_ids = self._get_tool_suffix_ids(tool_messages) - prompt_completion_tool_ids.append( - prompt_ids[idx_with_tool] + completion_ids[idx_with_tool] + suffix_ids - ) - - # Drop tool results whose addition would push the sequence past max_completion_length (the completion - # budget) or past the backend context ceiling (vLLM and transformers will error out on inputs longer than - # the model's max length). The sample exits the loop with its completion as-is, and the tool - # messages/images appended this iteration are rolled back so completions and tool_images stay consistent - # with completion_ids. - if self.use_vllm and self.vllm_mode == "colocate": - max_model_len = self.vllm_generation.llm.llm_engine.model_config.max_model_len - else: - config = self.model.config.text_config if self._is_vlm else self.model.config - max_model_len = config.max_position_embeddings - overlong = [ - len(pct) - len(prompt_ids[i]) > self.max_completion_length or len(pct) >= max_model_len - for i, pct in zip(idxs_with_tool, prompt_completion_tool_ids, strict=True) - ] - for idx in range(len(idxs_with_tool)): - if overlong[idx]: - idx_with_tool = idxs_with_tool[idx] - del completions[idx_with_tool][completions_len_before[idx] :] - del tool_images[idx_with_tool][tool_images_len_before[idx] :] - del prompts[idx_with_tool][prompts_len_before[idx] :] - # Keep only non-overlong items for further processing - idxs_with_tool = [idx for idx, o in zip(idxs_with_tool, overlong, strict=True) if not o] - prompt_completion_tool_ids = [ - pct for pct, o in zip(prompt_completion_tool_ids, overlong, strict=True) if not o - ] - if not idxs_with_tool: - break # all overlong, exit tool loop - - # Filter images and multimodal fields to match the current subset (index into full batch). - # Merge tool response images so the model can see visual feedback during generation. - merged_images = images - if any(imgs for imgs in tool_images): - if merged_images is None: - merged_images = [imgs if imgs else None for imgs in tool_images] - else: - merged_images = [ - (existing or []) + new for existing, new in zip(merged_images, tool_images, strict=True) - ] - loop_images = [merged_images[i] for i in idxs_with_tool] if merged_images else None - if self._is_vlm and self.tools and any(imgs for imgs in tool_images): - flat_loop_images = [img for img_list in loop_images if img_list for img in img_list] - if flat_loop_images: - image_inputs = self.processing_class.image_processor(images=flat_loop_images, return_tensors="pt") - image_inputs = super()._prepare_inputs(image_inputs) - loop_multimodal_fields = dict(image_inputs) - else: - loop_multimodal_fields = {} - elif multimodal_fields: - if "num_images" not in multimodal_fields and images is not None: - multimodal_fields = { - **multimodal_fields, - "num_images": [len(img_list) if img_list else 0 for img_list in images], - } - split_fields = split_pixel_values_by_grid(multimodal_fields) - loop_multimodal_fields = {} - for k, v in split_fields.items(): - selected = [v[i] for i in idxs_with_tool] - # Per-token type-id fields may arrive as tensors (e.g. via the padding workaround); - # convert them to lists so they follow the same left-padding path as list-valued - # per-token fields below, keeping them aligned with the left-padded input_ids. - if k in ("token_type_ids", "mm_token_type_ids") and isinstance(selected[0], torch.Tensor): - selected = [s.tolist() for s in selected] - # Per-token fields (e.g. token_type_ids) need zero-padding to match extended prompt length - if isinstance(selected[0], list): - selected = [ - s + [0] * (len(pct) - len(s)) - for s, pct in zip(selected, prompt_completion_tool_ids, strict=True) - ] - elif isinstance(v, torch.Tensor): - selected = torch.stack(selected) - loop_multimodal_fields[k] = selected - loop_multimodal_fields = unsplit_pixel_values_by_grid(loop_multimodal_fields) - loop_multimodal_fields.pop("num_images", None) - if "pixel_values" in loop_multimodal_fields and loop_multimodal_fields["pixel_values"].numel() == 0: - for key in ["pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw"]: - loop_multimodal_fields.pop(key, None) - else: - loop_multimodal_fields = {} - - # Generate new completions after tool execution (using concatenated IDs, no re-tokenization) - post_tool_ids = self._generate_single_turn( - prompt_completion_tool_ids, - loop_images, - loop_multimodal_fields, - has_tool_images=any(imgs for imgs in tool_images), - ) - - # Truncate so that pct[len(prompt_ids[idx]) :] + post_tool does not exceed max_completion_length. - # The pre-regen check guarantees len(completion_tool_ids) <= max_completion_length, so any - # excess can only come from post_tool_ids. post_tool_ids is model-generated text and never - # contains image tokens, so a plain slice is safe. - for idx in range(len(idxs_with_tool)): - idx_with_tool = idxs_with_tool[idx] - completion_tool_length = len(prompt_completion_tool_ids[idx]) - len(prompt_ids[idx_with_tool]) - excess_length = completion_tool_length + len(post_tool_ids[idx]) - self.max_completion_length - if excess_length > 0: - new_len = len(post_tool_ids[idx]) - excess_length - post_tool_ids[idx] = post_tool_ids[idx][:new_len] - - # Update tool_mask: the tool result should be 0 and the post-tool 1 - for idx in range(len(idxs_with_tool)): - idx_with_tool = idxs_with_tool[idx] - prompt_completion_tool_length = len(prompt_completion_tool_ids[idx]) - prompt_length = len(prompt_ids[idx_with_tool]) - completion_length = len(completion_ids[idx_with_tool]) - post_tool_length = len(post_tool_ids[idx]) - tool_length = prompt_completion_tool_length - prompt_length - completion_length - tool_mask[idx_with_tool] += [0] * tool_length + [1] * post_tool_length - - # Update completion_ids with the new completions (after tool execution) - for idx in range(len(idxs_with_tool)): - idx_with_tool = idxs_with_tool[idx] - prompt_length = len(prompt_ids[idx_with_tool]) - pct = prompt_completion_tool_ids[idx] # = prompt-completion-tool - completion_ids[idx_with_tool] = pct[prompt_length:] + post_tool_ids[idx] - - # Decode post-tool completions - post_tool_completions = [ - parse_response(self._tokenizer, ids, prefix=prompt_completion_tool_ids[idx]) if ids else {} - for idx, ids in enumerate(post_tool_ids) - ] - - # Add post-tool completions to the existing completions - for idx in range(len(idxs_with_tool)): - idx_with_tool = idxs_with_tool[idx] - if post_tool_completions[idx]: # {} if post-tool completions completely truncated - completions[idx_with_tool].append(post_tool_completions[idx]) - - # Check for further tool calls - tool_calls = [completion.get("tool_calls") for completion in post_tool_completions] - idxs_with_tool = [idx for idx, tool_call in zip(idxs_with_tool, tool_calls, strict=True) if tool_call] - tool_calls = [tool_call for tool_call in tool_calls if tool_call] - iteration_num += 1 - - return tool_mask, completions, completion_ids, tool_call_count, tool_failure_count, tool_images - - def _generate(self, prompts: list): - device = self.accelerator.device - mode = "train" if self.model.training else "eval" - - # Copy the prompts to avoid modifying the original list - prompts = copy.deepcopy(prompts) - - prompt_ids, images, multimodal_fields = self._tokenize_prompts(prompts) - completion_ids = self._generate_single_turn(prompt_ids, images, multimodal_fields) - - # Extract tool calls from the completions and (possibly) execute them - tool_images = [] - if self.tools: - # Decode completions. It's important to use `parse_response` when possible, because it handles tool calls. - if is_conversational({"prompt": prompts[0]}): - if Version(transformers.__version__) >= Version("5.0.0") and ( # parse_response added in v5 - getattr(self._tokenizer, "response_template", None) is not None # new-style - or getattr(self._tokenizer, "response_schema", None) is not None # old-style - ): - completions = [ - [parse_response(self._tokenizer, ids, prefix=prompt_ids[i])] - for i, ids in enumerate(completion_ids) - ] - else: - contents = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) - completions = [[{"role": "assistant", "content": content}] for content in contents] - else: - completions = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) - - ( - tool_mask, - completions, - completion_ids, - tool_call_count, - tool_failure_count, - tool_images, - ) = self._tool_call_loop(prompts, prompt_ids, completion_ids, completions, images, multimodal_fields) - # Merge tool response images into the images list for the forward pass - if any(imgs for imgs in tool_images): - if images is None: - images = [imgs if imgs else None for imgs in tool_images] - else: - images = [(existing or []) + new for existing, new in zip(images, tool_images, strict=True)] - else: - tool_mask = None - - # Get completion length per sequence, used for logging - prompt_lengths = torch.tensor([len(ids) for ids in prompt_ids], device=device) - if tool_mask is not None: # count only model-generated tokens (tool_mask=1) - completion_lengths = torch.tensor([sum(mask) for mask in tool_mask], device=device) - else: - completion_lengths = torch.tensor([len(ids) for ids in completion_ids], device=device) - agg_prompt_lengths = self.accelerator.gather(prompt_lengths) - agg_completion_lengths = self.accelerator.gather(completion_lengths) - # Fail clearly if the generation backend returned no completions (avoids a cryptic min() error below). - if agg_completion_lengths.numel() == 0: - raise RuntimeError( - "No completions were generated. This usually means the generation backend failed to return any " - "results; see the generation logs above for the underlying error." - ) - total_prompt_tokens = agg_prompt_lengths.sum() - - # Log the metrics - if mode == "train": - self.state.num_input_tokens_seen += (total_prompt_tokens + agg_completion_lengths.sum()).item() - self._metrics[mode]["num_tokens"] = [self.state.num_input_tokens_seen] - - # Log completion lengths, mean, min, max - self._metrics[mode]["completions/mean_length"].append(agg_completion_lengths.float().mean().item()) - self._metrics[mode]["completions/min_length"].append(agg_completion_lengths.float().min().item()) - self._metrics[mode]["completions/max_length"].append(agg_completion_lengths.float().max().item()) - - # Identify sequences that terminated with EOS and log their lengths - eos_and_pad = [*self.eos_token_ids, self._tokenizer.pad_token_id] - is_truncated = torch.tensor([ids[-1] not in eos_and_pad for ids in completion_ids], device=device) - agg_is_truncated = self.accelerator.gather(is_truncated) - self._metrics[mode]["completions/clipped_ratio"].append(agg_is_truncated.float().mean().item()) - term_completion_lengths = agg_completion_lengths[~agg_is_truncated] - if len(term_completion_lengths) == 0: # edge case where no terminated sequences are found - term_completion_lengths = torch.zeros(1, device=device) - self._metrics[mode]["completions/mean_terminated_length"].append(term_completion_lengths.float().mean().item()) - self._metrics[mode]["completions/min_terminated_length"].append(term_completion_lengths.float().min().item()) - self._metrics[mode]["completions/max_terminated_length"].append(term_completion_lengths.float().max().item()) - - if self.tools: - agg_tool_call_count = self.accelerator.gather(torch.tensor(tool_call_count, device=device)).sum() - tool_call_frequency = (agg_tool_call_count / len(agg_prompt_lengths)).item() - self._metrics[mode]["tools/call_frequency"].append(tool_call_frequency) - agg_tool_failure_count = self.accelerator.gather(torch.tensor(tool_failure_count, device=device)).sum() - failure_frequency = ( - (agg_tool_failure_count / agg_tool_call_count).item() if agg_tool_call_count > 0 else 0.0 - ) - self._metrics[mode]["tools/failure_frequency"].append(failure_frequency) - - return prompt_ids, completion_ids, tool_mask, images, tool_images - - @profiling_decorator - def _get_last_hidden_state( - self, - unwrapped_model, - input_ids, - attention_mask, - logits_to_keep, - pixel_values=None, - image_grid_thw=None, - pixel_attention_mask=None, - spatial_shapes=None, - image_sizes=None, - token_type_ids=None, - mm_token_type_ids=None, - image_position_ids=None, - ): - if is_peft_model(unwrapped_model): - unwrapped_model = unwrapped_model.base_model.model - - # Build model inputs - check if the model supports logits_to_keep (some models and VLMs don't) - model_inputs = {"input_ids": input_ids, "attention_mask": attention_mask} - - # For Qwen models: - if image_grid_thw is not None and pixel_values is not None: - model_inputs["image_grid_thw"] = image_grid_thw - # For Gemma, SmolVLM2, LLaVa-Next etc.: - if pixel_values is not None: - model_inputs["pixel_values"] = pixel_values - # For SmolVLM2 - if pixel_attention_mask is not None: - model_inputs["pixel_attention_mask"] = pixel_attention_mask - # For LFM2-VL - if spatial_shapes is not None: - model_inputs["spatial_shapes"] = spatial_shapes - # For LLaVa-Next - if image_sizes is not None: - model_inputs["image_sizes"] = image_sizes - if token_type_ids is not None: - model_inputs["token_type_ids"] = token_type_ids - if mm_token_type_ids is not None: - model_inputs["mm_token_type_ids"] = mm_token_type_ids - if image_position_ids is not None: - model_inputs["image_position_ids"] = image_position_ids - - # Only add logits_to_keep if the model supports it - if "logits_to_keep" in self.model_kwarg_keys: - # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded - model_inputs["logits_to_keep"] = logits_to_keep + 1 - - model_inputs["use_cache"] = False # only used in generation; set False to suppress warnings - - # `base_model` gives the backbone model (skipping `lm_head`) — text decoder for LMs, multimodal wrapper for - # VLMs (so vision-token injection runs before the text decoder). `get_decoder()` won't do: on VLMs it - # returns just the text stack and feeds image-placeholder IDs through it. - # Pre-5.0 transformers VLMs set `base_model_prefix = ""` so `base_model is self` (re-runs `lm_head`). - # Fall back to `.model` there. - if self._is_vlm and Version(transformers.__version__) < Version("5.0.0"): - backbone = unwrapped_model.model - else: - backbone = unwrapped_model.base_model - last_hidden_state = backbone(**model_inputs).last_hidden_state - # Exclude the last value: it corresponds to the next token pred - last_hidden_state = last_hidden_state[:, :-1, :] # (B, L-1, H) - # Only keep the last logits_to_keep. For model that support logits_to_keep, this is a no-op. - last_hidden_state = last_hidden_state[:, -logits_to_keep:, :] # (B, logits_to_keep, H) - return last_hidden_state - - # Name kept aligned with GRPO/RLOO for consistency; distillation has no rewards, so nothing is actually scored. - def _generate_and_score_completions(self, inputs: list[dict[str, torch.Tensor | Any]]) -> dict[str, Any]: - device = self.accelerator.device - - prompts = [x["prompt"] for x in inputs] - - if "images" in inputs[0]: - images = [example.get("images") for example in inputs] - elif "image" in inputs[0]: - images = [[example.get("image")] if example.get("image") is not None else None for example in inputs] - else: - images = None - # Transformers requires at least one image in the batch, otherwise it throws an error - if images is not None and all(img_list == [] for img_list in images): - images = None - - # If the prompts are conversational and the inputs contain images, we need to convert the prompts from - # [{"role": "user", "content": "What color is the sky?"}] to - # [{"role": "user", "content": [{"type": "image", "image": }, {"type": "text", "text": "What color is the sky?"}]}] - if images is not None: - if not is_conversational(inputs[0]): - raise ValueError( - "Multimodal training requires conversational prompts. It looks like the dataset contains " - "non-conversational inputs, likely because a chat template was applied before passing the dataset " - "to the trainer. Please provide the raw conversational prompts and let the trainer apply the chat " - "template internally." - ) - prompts = [ - prepare_multimodal_messages(prompt, images=image_list) - for prompt, image_list in zip(prompts, images, strict=True) - ] - - prompt_ids_list, completion_ids_list, tool_mask_list, images, tool_images = self._generate(prompts) - - # Convert lists of token IDs to padded tensors - prompt_ids = [torch.tensor(ids) for ids in prompt_ids_list] - prompt_mask = [torch.ones_like(ids, dtype=torch.long) for ids in prompt_ids] - prompt_ids = pad( - prompt_ids, - padding_value=self._tokenizer.pad_token_id, - padding_side="left", - pad_to_multiple_of=self.pad_to_multiple_of, - ).to(device=device) - prompt_mask = pad( - prompt_mask, padding_value=0, padding_side="left", pad_to_multiple_of=self.pad_to_multiple_of - ).to(device=device) - completion_ids = [torch.tensor(ids) for ids in completion_ids_list] - completion_mask = [torch.ones_like(ids, dtype=torch.long) for ids in completion_ids] - completion_ids = pad( - completion_ids, - padding_value=self._tokenizer.pad_token_id, - padding_side="right", - pad_to_multiple_of=self.pad_to_multiple_of, - ).to(device=device) - completion_mask = pad( - completion_mask, padding_value=0, padding_side="right", pad_to_multiple_of=self.pad_to_multiple_of - ).to(device=device) - if tool_mask_list is not None: - tool_mask = [torch.tensor(mask) for mask in tool_mask_list] - tool_mask = pad( - tool_mask, padding_value=1, padding_side="right", pad_to_multiple_of=self.pad_to_multiple_of - ).to(device=device) - else: - tool_mask = None - - loss_mask = completion_mask if tool_mask is None else completion_mask * tool_mask - num_items_in_batch = self.accelerator.gather(loss_mask.sum()).sum() - - num_images = [len(img_list) if img_list else 0 for img_list in images] if images is not None else None - - # Get forward_kwargs for models with multimodal inputs. - # When tool images are present (from _tool_call_loop), use image_processor directly and build - # mm_token_type_ids from prompt_completion_ids. Otherwise, use the full processor pipeline - # which returns model-specific keys (image_sizes, pixel_attention_mask, etc.). - if self.tools and any(imgs for imgs in tool_images) and self._is_vlm: - flat_images = [img for img_list in images if img_list for img in img_list] - image_inputs = self.processing_class.image_processor(images=flat_images, return_tensors="pt") - image_inputs = super()._prepare_inputs(image_inputs) - forward_kwargs = dict(image_inputs) - elif images is not None: - prompts_text = [ - apply_chat_template( - {"prompt": prompt}, self.processing_class, tools=self.tools or None, **self.chat_template_kwargs - )["prompt"] - for prompt in prompts - ] - prompt_inputs = self.processing_class(images=images, text=prompts_text, padding=True, return_tensors="pt") - prompt_inputs = super()._prepare_inputs(prompt_inputs) - forward_kwargs = {k: v for k, v in prompt_inputs.items() if k not in ["input_ids", "attention_mask"]} - else: - forward_kwargs = {} - - # Recover LFM2-VL tile counts; the full processor drops row/column metadata. - num_tiles = None - if images is not None and "spatial_shapes" in forward_kwargs: - image_info = self.processing_class.image_processor( - images=images, return_tensors="pt", return_row_col_info=True - ) - tiles_per_image = image_info["image_rows"] * image_info["image_cols"] - if self.processing_class.image_processor.use_thumbnail: - tiles_per_image = tiles_per_image + (tiles_per_image > 1).to(tiles_per_image.dtype) - num_tiles = [group.sum().item() for group in torch.split(tiles_per_image, num_images)] - # Same for InternVL, whose pixel_values is tile-indexed ([total_tiles, channels, height, width]). - elif ( - images is not None - and forward_kwargs["pixel_values"].ndim == 4 - and forward_kwargs["pixel_values"].size(0) != sum(num_images) - ): - num_patches = self.processing_class.image_processor( - images=images, crop_to_patches=True, return_tensors="pt" - )["num_patches"] - num_tiles = [group.sum().item() for group in torch.split(num_patches, num_images)] - - # If token_type_ids are used, extend them with zeros for the completion part - if "token_type_ids" in forward_kwargs: - token_type_ids = forward_kwargs["token_type_ids"] - if self.pad_to_multiple_of is not None: - # Needed only with pad_to_multiple_of: otherwise prompt_ids and token_type_ids must have equal len - padding_size = prompt_ids.size(1) - token_type_ids.size(1) - if padding_size > 0: - token_type_ids = torch.cat( - [token_type_ids.new_zeros((token_type_ids.size(0), padding_size)), token_type_ids], dim=1 - ) - forward_kwargs["token_type_ids"] = torch.cat( - [token_type_ids, token_type_ids.new_zeros(completion_ids.shape)], dim=1 - ) - # If mm_token_type_ids are used, extend them with zeros for the completion part - if "mm_token_type_ids" in forward_kwargs: - mm_token_type_ids = forward_kwargs["mm_token_type_ids"] - if self.pad_to_multiple_of is not None: - # Needed only with pad_to_multiple_of: otherwise prompt_ids and mm_token_type_ids must have equal len - padding_size = prompt_ids.size(1) - mm_token_type_ids.size(1) - if padding_size > 0: - mm_token_type_ids = torch.cat( - [mm_token_type_ids.new_zeros((mm_token_type_ids.size(0), padding_size)), mm_token_type_ids], - dim=1, - ) - forward_kwargs["mm_token_type_ids"] = torch.cat( - [mm_token_type_ids, mm_token_type_ids.new_zeros(completion_ids.shape)], dim=1 - ) - - # For VLM tool images: build token type IDs from the full prompt_completion_ids. - # This must happen AFTER the token_type_ids/mm_token_type_ids extension blocks above, - # because our version already covers the full sequence (images are in the completion, - # not just the prompt). - if self.tools and any(imgs for imgs in tool_images) and self._is_vlm: - prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1) # (B, P+C) - mm_ids = torch.zeros_like(prompt_completion_ids) - if self._image_pad_token_id is not None: - mm_ids[prompt_completion_ids == self._image_pad_token_id] = 1 - if self._video_pad_token_id is not None: - mm_ids[prompt_completion_ids == self._video_pad_token_id] = 2 - - # Use the same key the model expects: token_type_ids for models like Gemma, - # mm_token_type_ids for models like Qwen. - image_grid_thw = forward_kwargs.get("image_grid_thw") - if image_grid_thw is not None: - forward_kwargs["mm_token_type_ids"] = mm_ids - else: - forward_kwargs["token_type_ids"] = mm_ids - - # Truncation safety (Qwen-style models with image_grid_thw only): if - # max_completion_length truncated some image tokens, the number of image pad tokens - # in input_ids won't match pixel_values features. Check per-sample and drop ALL - # images for any sample with a mismatch (safe fallback). - if image_grid_thw is not None and num_images is not None: - merge_length = getattr(self.processing_class.image_processor, "merge_size", 2) ** 2 - img_offset = 0 - has_mismatch = False - for b in range(mm_ids.shape[0]): - sample_tokens = (mm_ids[b] == 1).sum().item() - sample_features = 0 - for i in range(num_images[b]): - grid_idx = img_offset + i - if grid_idx < image_grid_thw.shape[0]: - sample_features += image_grid_thw[grid_idx].prod().item() // merge_length - if sample_tokens != sample_features: - has_mismatch = True - break - img_offset += num_images[b] - - if has_mismatch: - # Drop all images: safer than partial trim which is error-prone - forward_kwargs.pop("pixel_values", None) - forward_kwargs.pop("image_grid_thw", None) - mm_ids.zero_() - forward_kwargs["mm_token_type_ids"] = mm_ids - num_images = None - - # Log prompt and completion texts - if self.log_completions: - prompts_text = self.processing_class.batch_decode(prompt_ids, skip_special_tokens=True) - completions_text = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) - self._logs["prompt"].extend(gather_object(prompts_text)) - self._logs["completion"].extend(gather_object(completions_text)) - - output = { - "prompt_ids": prompt_ids, - "prompt_mask": prompt_mask, - "completion_ids": completion_ids, - "completion_mask": completion_mask, - "num_items_in_batch": num_items_in_batch, - } - if "pixel_values" in forward_kwargs: - output["pixel_values"] = forward_kwargs["pixel_values"] - if "image_grid_thw" in forward_kwargs: - output["image_grid_thw"] = forward_kwargs["image_grid_thw"] - if "pixel_attention_mask" in forward_kwargs: - output["pixel_attention_mask"] = forward_kwargs["pixel_attention_mask"] - if "spatial_shapes" in forward_kwargs: - output["spatial_shapes"] = forward_kwargs["spatial_shapes"] - if "image_sizes" in forward_kwargs: - output["image_sizes"] = forward_kwargs["image_sizes"] - if "token_type_ids" in forward_kwargs: - output["token_type_ids"] = forward_kwargs["token_type_ids"] - if "mm_token_type_ids" in forward_kwargs: - output["mm_token_type_ids"] = forward_kwargs["mm_token_type_ids"] - if "image_position_ids" in forward_kwargs: - output["image_position_ids"] = forward_kwargs["image_position_ids"] - if images is not None: - output["num_images"] = num_images - if num_tiles is not None: - output["num_tiles"] = num_tiles - if tool_mask is not None: - output["tool_mask"] = tool_mask - return output - - @profiling_decorator - def _prepare_inputs(self, generation_batch: dict[str, torch.Tensor | Any]) -> dict[str, torch.Tensor | Any]: - # Prepares inputs for model training/evaluation by managing completion generation and batch handling. - # During training: - # - Receives the local generation batch (Per-GPU batch size × gradient accumulation steps) - # from the modified training dataloader instead of the standard local batch - # - Generates completions once for the entire generation batch and splits it into batches of size - # `per_device_train_batch_size` - # - Buffers these completions and returns the appropriate slice for the current accumulation step - # - Optimizes by regenerating completions only periodically (every gradient_accumulation_steps) - # During evaluation: - # - The input is treated as a standard local batch (no accumulation) - # - Completions are generated for each batch without buffering or reuse - # Returns a single local batch in both cases. - - mode = "train" if self.model.training else "eval" - if mode == "train": - generate_every = self.args.gradient_accumulation_steps - if self._step % generate_every == 0 or self._buffered_inputs is None: - # self._buffered_inputs=None can occur when resuming from a checkpoint - generation_batch = self._generate_and_score_completions(generation_batch) - generation_batch = split_pixel_values_by_grid(generation_batch) - - try: generation_batch = shuffle_sequence_dict(generation_batch) - - except: pass - generation_batches = split_tensor_dict(generation_batch, self.args.gradient_accumulation_steps) - self._buffered_inputs = [unsplit_pixel_values_by_grid(batch) for batch in generation_batches] - inputs = self._buffered_inputs[self._step % self.args.gradient_accumulation_steps] - else: - # In evaluation, there is neither batch grouping for generation, nor multiple iterations, hence - # local generation batch == local eval batch - inputs = self._generate_and_score_completions(generation_batch) - return inputs - - @profiling_decorator - def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): - # transformers computes `num_items_in_batch` from the raw dataloader labels, before on-policy generation - # replaces the completions; use the count over the generated completions instead (computed in - # `_generate_and_score_completions`). Divide by the process count so the per-process loss compensates for DDP - # gradient averaging. - if self.model.training and inputs.get("num_items_in_batch") is not None: - num_items_in_batch = inputs["num_items_in_batch"].clamp(min=1.0) / self.accelerator.num_processes - - # Route the whole loss (backbone + `lm_head` projection + JSD) through the DDP/FSDP wrapper via - # `_forward_redirection`, so DDP.forward() fires `prepare_for_backward()` and FSDP/DeepSpeed keep the student's - # sharded parameters (including the `lm_head`) materialized for the projection. - unwrapped_student = self.accelerator.unwrap_model(model) - loss, entropy_sum, num_valid_tokens = self._forward_redirection( - model, unwrapped_student, self._compute_loss, unwrapped_student, inputs, num_items_in_batch - ) - - # Log the mean per-token student entropy (in nats). The reduction runs here, after `_forward_redirection` - # returns, so the `gather_for_metrics` collective does not run inside the DDP/FSDP-wrapped forward (a hang/ - # ordering risk). Mirrors `SFTTrainer.compute_loss`. - mode = "train" if self.model.training else "eval" - num_valid_tokens = self.accelerator.gather_for_metrics(num_valid_tokens).sum() - entropy_sum = self.accelerator.gather_for_metrics(entropy_sum).sum() - entropy = (entropy_sum / num_valid_tokens).item() if num_valid_tokens > 0 else 0.0 - self._metrics[mode]["entropy"].append(entropy) - - return (loss, None) if return_outputs else loss - - def _compute_loss(self, unwrapped_student, inputs, num_items_in_batch): - # Chunked JSD path: project the teacher/student hidden states to vocab logits one chunk at a time (never - # materializing the full `(B, C, V)` logits), so the teacher's dense distribution can be matched without - # buffering it. Runs inside the student wrapper's forward (see `compute_loss`). - input_ids = torch.cat([inputs["prompt_ids"], inputs["completion_ids"]], dim=1) - attention_mask = torch.cat([inputs["prompt_mask"], inputs["completion_mask"]], dim=1) - completion_mask = inputs["completion_mask"] - # Apply tool_mask for loss computation in multi-turn training scenarios - loss_mask = completion_mask if "tool_mask" not in inputs else completion_mask * inputs["tool_mask"] - logits_to_keep = inputs["completion_ids"].size(1) # only the completion tokens are trained on - - # Multimodal (VLM) fields, extracted during generation and split to this micro-batch by `_prepare_inputs`. - multimodal_keys = ( - "pixel_values", - "image_grid_thw", - "pixel_attention_mask", - "spatial_shapes", - "image_sizes", - "token_type_ids", - "mm_token_type_ids", - "image_position_ids", - ) - multimodal_inputs = {k: inputs[k] for k in multimodal_keys if k in inputs} - - student_hidden_states = self._get_last_hidden_state( - unwrapped_student, input_ids, attention_mask, logits_to_keep, **multimodal_inputs - ) - - # Route the teacher backbone through its own wrapper via `_forward_redirection` too, so FSDP/DeepSpeed - # materialize its sharded parameters before the forward runs (the backbone call would otherwise see shards). - self.teacher_model.eval() - unwrapped_teacher = self.accelerator.unwrap_model(self.teacher_model) - with torch.no_grad(): - teacher_hidden_states = self._forward_redirection( - self.teacher_model, - unwrapped_teacher, - self._get_last_hidden_state, - unwrapped_teacher, - input_ids, - attention_mask, - logits_to_keep, - **multimodal_inputs, - ) - - student_lm_head = unwrapped_student.get_output_embeddings() - teacher_lm_head = unwrapped_teacher.get_output_embeddings() - - # On VLMs the logit post-processing lives on `text_config`, so read it through `get_text_config()`. - student_config = unwrapped_student.config.get_text_config() - teacher_config = unwrapped_teacher.config.get_text_config() - # `logit_scale` is None on models that don't scale (e.g. MPT); read that as unscaled (1.0). A real 0.0 is kept - # as-is. Muse Glimmer applies the same pre-softcap multiplier under the name `output_multiplier`. - student_logit_scale = getattr(student_config, "logit_scale", None) - if student_logit_scale is None: - student_logit_scale = getattr(student_config, "output_multiplier", None) - # Falcon-H1 multiplies the logits by `lm_head_multiplier`. Granite and MiniCPM3 divide them by - # `logits_scaling`, while HyperCLOVA X multiplies them by it. - if student_logit_scale is None: - student_logit_scale = getattr(student_config, "lm_head_multiplier", None) - if student_logit_scale is None and getattr(student_config, "logits_scaling", None) is not None: - logits_scaling = student_config.logits_scaling - student_logit_scale = logits_scaling if student_config.model_type == "hyperclovax" else 1 / logits_scaling - teacher_logit_scale = getattr(teacher_config, "logit_scale", None) - if teacher_logit_scale is None: - teacher_logit_scale = getattr(teacher_config, "output_multiplier", None) - # Falcon-H1 multiplies the logits by `lm_head_multiplier`. Granite and MiniCPM3 divide them by - # `logits_scaling`, while HyperCLOVA X multiplies them by it. - if teacher_logit_scale is None: - teacher_logit_scale = getattr(teacher_config, "lm_head_multiplier", None) - if teacher_logit_scale is None and getattr(teacher_config, "logits_scaling", None) is not None: - logits_scaling = teacher_config.logits_scaling - teacher_logit_scale = logits_scaling if teacher_config.model_type == "hyperclovax" else 1 / logits_scaling - student_logit_scale = 1.0 if student_logit_scale is None else student_logit_scale - teacher_logit_scale = 1.0 if teacher_logit_scale is None else teacher_logit_scale - loss, entropy_sum, n_valid = _chunked_divergence_loss( - student_hidden_states, - teacher_hidden_states, - student_lm_head.weight, - teacher_lm_head.weight, - loss_mask, - self.beta, - num_items_in_batch=num_items_in_batch, - student_lm_head_bias=student_lm_head.bias, - teacher_lm_head_bias=teacher_lm_head.bias, - student_logit_scale=student_logit_scale, - teacher_logit_scale=teacher_logit_scale, - student_final_logit_softcapping=getattr(student_config, "final_logit_softcapping", None), - teacher_final_logit_softcapping=getattr(teacher_config, "final_logit_softcapping", None), - temperature=self.temperature, - ) - # Return the raw entropy sum and valid-token count for `compute_loss` to aggregate and log after the forward - # returns (see there). Detached: the metric is gradient-free. - return loss, entropy_sum.detach(), n_valid - - def training_step(self, model, inputs, num_items_in_batch): - time_before = time.perf_counter() - output = super().training_step(model, inputs, num_items_in_batch) - self._step += 1 - time_after = time.perf_counter() - self._current_train_step_time += time_after - time_before - if self._step % self.current_gradient_accumulation_steps == 0: - self._metrics["train"]["step_time"].append(self._current_train_step_time) - self._current_train_step_time = 0.0 - return output - - # During eval, Trainer calls prediction_step. If no labels are present in the inputs, it only runs forward and - # returns logits. We override prediction_step to force compute_loss, because this trainer doesn't involve labels. - def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys: list[str] | None = None): - inputs = self._prepare_inputs(inputs) - with torch.no_grad(): - with self.compute_loss_context_manager(): - loss = self.compute_loss(model, inputs) - loss = loss.mean().detach() - return loss, None, None - - def log(self, logs: dict[str, float], start_time: float | None = None) -> None: - mode = "train" if self.model.training else "eval" - # Average the metrics - metrics = {} - for key, val in self._metrics[mode].items(): - # Filter out NaN values before averaging. With logging_steps > 1, a naive sum()/len() would let a single - # NaN contaminate valid data from other batches. Only return None when no valid values remain (e.g. JSON - # loggers crash on float NaN). - valid = [v for v in val if not math.isnan(v)] - metrics[key] = sum(valid) / len(valid) if valid else None - - # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` - # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. - if mode == "eval": - metrics = {f"eval_{key}": val for key, val in metrics.items()} - - logs.update(metrics) - super().log(logs, start_time) - self._metrics[mode].clear() - - if self.accelerator.is_main_process and self.log_completions: - if is_rich_available(): - print_prompt_completions_sample( - self._logs["prompt"], - self._logs["completion"], - {}, - None, - self.state.global_step, - self.num_completions_to_print, - ) - - logging_backends = [] - if self.args.report_to and "wandb" in self.args.report_to and wandb.run is not None: - logging_backends.append(wandb) - if self.args.report_to and "trackio" in self.args.report_to: - logging_backends.append(trackio) - - import pandas as pd - - table = { - "step": [self.state.global_step] * len(self._logs["prompt"]), - "prompt": self._logs["prompt"], - "completion": self._logs["completion"], - } - df_base = pd.DataFrame(table) - df_base.to_parquet( - os.path.join( - self.args.output_dir, - "completions", - f"completions_{self.state.global_step:05d}.parquet", - ) - ) - - for logging_backend in logging_backends: - df = df_base - if self.log_unique_prompts: - df = df.drop_duplicates(subset=["prompt"]) - logging_backend.log({"completions": logging_backend.Table(dataframe=df)}) - - # Ensure the model card is saved along with the checkpoint - def _save_checkpoint(self, model, trial): - if self.args.hub_model_id is None: - model_name = Path(self.args.output_dir).name - else: - model_name = self.args.hub_model_id.split("/")[-1] - self.create_model_card(model_name=model_name) - super()._save_checkpoint(model, trial) -class UnslothDistillationTrainer(_UnslothDistillationTrainer): - """ - - Trainer for knowledge distillation. The student is trained on-policy — it generates the completions itself — to - match the teacher's next-token distribution under a generalized Jensen-Shannon divergence (interpolating forward - KL, reverse KL, and JSD via `beta`), as introduced in [On-Policy Distillation of Language - Models](https://huggingface.co/papers/2306.13649). - - Example: - - ```python - >>> from trl import DistillationTrainer - >>> from datasets import load_dataset - - >>> dataset = load_dataset("trl-lib/tldr", split="train") - - >>> trainer = DistillationTrainer( - ... model="Qwen/Qwen2.5-0.5B-Instruct", - ... teacher_model="Qwen/Qwen2.5-1.5B-Instruct", - ... train_dataset=dataset, - ... ) - >>> trainer.train() - ``` - - Args: - model (`str` or [`~transformers.PreTrainedModel`] or [`~peft.PeftModel`]): - Model to be trained. Can be either: - - - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a - path to a *directory* containing model weights saved using - [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded - using `.from_pretrained` (where `` is derived from the model - config) with the keyword arguments in `args.model_init_kwargs`. If `dtype` is not specified in - `args.model_init_kwargs`, it defaults to `float32`. This differs from - [`~transformers.PreTrainedModel.from_pretrained`], where (since Transformers v5) the dtype is inferred - from the model config. - - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. - - A [`~peft.PeftModel`] object. Only causal language models are supported. - teacher_model (`str` or [`~transformers.PreTrainedModel`], *optional*): - Teacher model whose next-token distribution the student is trained to match. Can be a *model id* / path - (loaded like `model`, using `args.teacher_model_init_kwargs`) or an instantiated - [`~transformers.PreTrainedModel`]. It must share the student's vocabulary. May be omitted by subclasses - that supply the teacher another way (e.g. a remote server). - args ([`DistillationConfig`], *optional*): - Configuration for this trainer. If `None`, a default configuration is used. - train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): - Dataset to use for training. It must include a column `"prompt"`. Any additional columns in the dataset is - ignored. The format of the samples can be either: - - - [Standard](dataset_formats#standard): Each sample contains plain text. - - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role - and content). - - When `train_dataset` is an [`~datasets.IterableDataset`] (e.g. a streaming dataset), `max_steps` must be - set in the training arguments, since its length cannot be inferred and the total number of training steps - is required to bound the training loop and configure the learning rate scheduler. - eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`], [`~datasets.DatasetDict`], [`~datasets.IterableDatasetDict`] or `dict[str, Dataset | IterableDataset]`): - Dataset to use for evaluation. It must meet the same requirements as `train_dataset`. - processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.ProcessorMixin`], *optional*): - Processing class used to process the data. The padding side must be set to "left". If `None`, the - processing class is loaded from the model's name with [`~transformers.AutoProcessor.from_pretrained`]. A - padding token, `tokenizer.pad_token`, must be set. If the processing class has not set a padding token, - `tokenizer.eos_token` will be used as the default. - callbacks (list of [`~transformers.TrainerCallback`], *optional*): - List of callbacks to customize the training loop. Will add those to the list of default callbacks detailed - in [here](https://huggingface.co/docs/transformers/main_classes/callback). - - If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`] - method. - optimizers (`tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None]`, *optional*, defaults to `(None, None)`): - A tuple containing the optimizer and the scheduler to use. Will default to an instance of `AdamW` on your - model and a scheduler given by [`~transformers.get_linear_schedule_with_warmup`] controlled by `args`. - quantization_config ([`~transformers.BitsAndBytesConfig`], *optional*): - Quantization configuration used when loading the model from a model identifier. Combine with `peft_config` - for QLoRA training. Ignored if the model is already instantiated. - peft_config ([`~peft.PeftConfig`], *optional*): - PEFT configuration used to wrap the model. If `None`, the model is not wrapped. - tools (list of `Callable`, *optional*): - A list of callable tool functions that the model can invoke during generation. Each tool should be a - standard Python function with properly type-hinted arguments and return values, and a Google-style - docstring describing its purpose, arguments, and return value. For more details, see: - https://huggingface.co/docs/transformers/en/chat_extras#passing-tools. The model uses the function's name, - type hints, and docstring to determine how to call it. Ensure that the model's chat template supports tool - use and that it has been fine-tuned for tool calling. - - """ - def __init__( - self, - model, - teacher_model = None, - args = None, - train_dataset = None, - eval_dataset = None, - processing_class = None, - callbacks = None, - quantization_config = None, - peft_config = None, - tools = None, - **kwargs - ): - if args is None: args = UnslothDistillationConfig() - _unsloth_precision_requested = tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) - use_bf16 = getattr(args, 'bf16', False) - if type(use_bf16) is not bool: use_bf16 = False - use_fp16 = getattr(args, 'fp16', False) - if type(use_fp16) is not bool: use_fp16 = False - force_float32 = False - try: - from unsloth_zoo.device_type import device_is_bf16_supported as _bf16_supported - except Exception: - _bf16_supported = torch.cuda.is_bf16_supported - try: - from unsloth_zoo.device_type import DEVICE_TYPE as _unsloth_device_type - except Exception: - _unsloth_device_type = None - full_finetuning = getattr(model, '_unsloth_full_finetuning', None) - if full_finetuning is None: full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1' - model_forced_float32 = getattr(model, '_unsloth_forced_float32', None) - if model_forced_float32 is None: model_forced_float32 = os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1' - if model_forced_float32 and not (full_finetuning and _bf16_supported()): - print('Unsloth: Switching to float32 training since model cannot work with float16') - force_float32 = True - mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') - dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None) - if dtype is None: dtype = model.get_input_embeddings().weight.dtype - from unsloth_zoo.utils import _get_dtype - dtype = _get_dtype(dtype) - float16 = dtype == torch.float16 - bfloat16 = dtype == torch.bfloat16 - float32 = dtype == torch.float32 - user_float32 = bool(getattr(model, '_unsloth_user_float32', False)) - if full_finetuning: - if bfloat16 and use_fp16: use_fp16 = False - if float16 and use_bf16: use_bf16 = False - if not force_float32 and bfloat16 and use_fp16 and _unsloth_device_type == 'xpu' and not _bf16_supported(): - print('Unsloth: The model was loaded in bfloat16 on purpose, so training uses bfloat16 mixed precision instead of float16.') - use_fp16 = False - use_bf16 = True - args.fp16 = False - args.bf16 = True - if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`') - if not force_float32 and (bfloat16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`') - if force_float32: - # Forced float32 training - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32' and float32 and user_float32 and not _bf16_supported(): - print('Unsloth: Model is in float32 and this GPU has no bfloat16 support, so training stays in float32. Pass fp16 = True to force float16 mixed precision instead.') - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32': - # Mixed precision training. bf16 only if the GPU supports it; V100/T4 use fp16. - use_bf16_amp = (not float16) and (_bf16_supported() or (bfloat16 and _unsloth_device_type == 'xpu')) - args.fp16 = not use_bf16_amp - args.bf16 = use_bf16_amp - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16_amp else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16_amp else 'fp16' - # args.mixed_precision is a new argument which needs to be set now - elif mixed_precision_dtype == 'bfloat16': - # Both False since bfloat16 full finetuning doesn't do any autocasting. - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif use_bf16 or use_fp16: - # transformers <5 exported this itself from fp16/bf16; 5.x dropped the write, so an - # explicit flag left it unset and GRPO readers defaulted to 'fp16', wrapping a - # bfloat16 model in a float16 autocast. See unslothai/unsloth#4891. - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16 else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16 else 'fp16' - - if getattr(args, 'eval_dataset', None) is not None and getattr(args, 'eval_strategy', 'no') == 'no': - args.eval_strategy = 'steps' - if getattr(args, 'eval_steps', None) is None: args.eval_steps = 0.1 - ga_steps = getattr(args, 'gradient_accumulation_steps', None) - if ga_steps is not None and ga_steps > 1: - from transformers import __version__ as transformers_version - if Version(transformers_version) <= Version('4.45.2'): - print('**** Unsloth: Please use our fixed gradient_accumulation_steps by updating transformers, TRL and Unsloth!\n' - '`pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth transformers trl unsloth_zoo`') - if getattr(args, 'eval_strategy', 'no') != 'no': - eval_bsz = getattr(args, 'per_device_eval_batch_size', 8) - eval_generations = getattr(args, 'num_generations_eval', None) or getattr(args, 'num_generations', None) or 1 - if eval_bsz == 8 and args.per_device_train_batch_size < eval_bsz and (args.per_device_train_batch_size * args.world_size) % eval_generations == 0: args.per_device_eval_batch_size = args.per_device_train_batch_size - if getattr(args, 'eval_accumulation_steps', None) is None and ga_steps is not None: args.eval_accumulation_steps = ga_steps - fp16_full_eval = getattr(args, 'fp16_full_eval', False) - if type(fp16_full_eval) is not bool: fp16_full_eval = False - bf16_full_eval = getattr(args, 'bf16_full_eval', False) - if type(bf16_full_eval) is not bool: bf16_full_eval = False - if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True - if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False - if force_float32: - args.bf16_full_eval = False - args.fp16_full_eval = False - elif os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16': - args.bf16_full_eval = True - args.fp16_full_eval = False - elif not bf16_full_eval and not fp16_full_eval: - args.bf16_full_eval = args.bf16 - args.fp16_full_eval = args.fp16 - if tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) != _unsloth_precision_requested: - from unsloth.models.loader_utils import sync_deepspeed_mixed_precision - sync_deepspeed_mixed_precision(args) - if model is not None: - from unsloth.models.loader_utils import apply_deepspeed_zero3_model_setup - apply_deepspeed_zero3_model_setup(model) - if model is not None: - _warnings_issued = getattr(model, 'warnings_issued', None) - if _warnings_issued is None: - model.warnings_issued = {} - elif not isinstance(_warnings_issued, dict): - try: - model.warnings_issued = dict(_warnings_issued) - except Exception: - model.warnings_issued = {} - if 'max_seq_length' not in locals() and not hasattr(args, 'max_seq_length'): - pass - else: - model_max_seq_length = getattr(model, 'max_seq_length', None) - args_max_seq_length = getattr(args, 'max_seq_length', None) - if args_max_seq_length is None and model_max_seq_length is not None: - max_seq_length = model.max_seq_length - if hasattr(args, 'max_seq_length'): args.max_seq_length = max_seq_length - elif args_max_seq_length is not None and model_max_seq_length is not None: - if args_max_seq_length > model_max_seq_length: - print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but ' - 'the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.') - args.max_seq_length = model_max_seq_length - if model is not None and hasattr(model, 'for_training'): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right' - if 'processing_class' in locals(): - if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right' - if hasattr(processing_class, 'tokenizer') and hasattr(processing_class.tokenizer, 'padding_side'): processing_class.tokenizer.padding_side = 'right' - other_metrics = [] - - from unsloth_zoo.logging_utils import PatchRLStatistics - PatchRLStatistics('distillation_trainer', other_metrics) - - # [TODO] Fix up DataParallel multiplying batch sizes - # [TODO] DDP works, but DP seems to not work? [TODO] - if getattr(args, "parallel_mode", None) == ParallelMode.NOT_DISTRIBUTED and args.n_gpu > 1: - if getattr(args, "_n_gpu", 1) != 1: - args._n_gpu = 1 - if "model" in locals(): - try: - from unsloth.models._gc_config_default import honour_gradient_checkpointing_off - honour_gradient_checkpointing_off(args, model) - except Exception: - pass - if "model" in locals() and hasattr(model, "for_training"): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if "model" in locals(): - try: - from unsloth.models._ddp_defaults import unsloth_default_ddp_find_unused_parameters - unsloth_default_ddp_find_unused_parameters(args, model) - except Exception: - pass - super().__init__( - model = model, - teacher_model = teacher_model, - args = args, - train_dataset = train_dataset, - eval_dataset = eval_dataset, - processing_class = processing_class, - callbacks = callbacks, - quantization_config = quantization_config, - peft_config = peft_config, - tools = tools,**kwargs) - # A tensor parallel model (tp_native.tp_trainer_setup) stays on its training path: for_inference would - # route the decoder layers onto decode paths outside the tensor parallel regions. - if "model" in locals() and hasattr(model, "for_inference") and not getattr(self, "_unsloth_tensor_parallel", False): - model.for_inference() - if hasattr(self, 'neftune_hook_handle'): - self.neftune_hook_handle.remove() - if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle - if getattr(args, 'neftune_noise_alpha', None) is not None: - model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha - pass - if hasattr(self, 'accelerator'): - scaler = self.accelerator.scaler - current_model = model - while hasattr(current_model, 'model'): - current_model.accelerator_scaler = scaler - current_model = current_model.model - current_model.accelerator_scaler = scaler - pass - if hasattr(self, 'train'): - self.train = MethodType(prepare_for_training_mode(self.__class__.train), self) - pass - if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'): - _vllm_tok = self.llm.get_tokenizer() - _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None) - if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None: - _vllm_tok.chat_template = _pc.chat_template - pass - if getattr(self, 'aux_loss_enabled', False) and hasattr(getattr(self, 'model', None), 'config'): - _text_config = self.model.config - if hasattr(_text_config, 'get_text_config'): _text_config = _text_config.get_text_config() - _n_experts = [getattr(_text_config, _k) for _k in ('num_local_experts', 'num_experts', 'n_routed_experts', 'moe_num_experts') if isinstance(getattr(_text_config, _k, None), int)] - if _n_experts and all(_n <= 1 for _n in _n_experts): - self.aux_loss_enabled = False - _text_config.output_router_logits = False - pass - -pass - - -if hasattr(logger, "addFilter"): - import logging - class HideLoggingMessage(logging.Filter): - def __init__(self, text): self.text = text - def filter(self, x): return not (self.text in x.getMessage()) - pass - logger.addFilter(HideLoggingMessage("`use_cache=True`")) - diff --git a/unsloth_compiled_cache/UnslothGKDTrainer.py b/unsloth_compiled_cache/UnslothGKDTrainer.py deleted file mode 100644 index edead96eeed78e912398ce14dcc92752af77582b..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/UnslothGKDTrainer.py +++ /dev/null @@ -1,2172 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from unsloth_zoo.temporary_patches.common import torch_compile -from unsloth_zoo.temporary_patches.common import _maybe_compile -import functools -import transformers -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from trl.experimental.gkd.gkd_trainer import (Any, AutoModelForCausalLM, BaseImageProcessor, Callable, DataCollator, DataCollatorForChatML, Dataset, EvalPrediction, F, FeatureExtractionMixin, GKDConfig, GKDTrainer, GenerationConfig, PeftConfig, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, SFTTrainer, TrainerCallback, disable_dropout_in_model, empty_cache, logger, nn, prepare_deepspeed, random, textwrap, torch, unwrap_model_for_generation, AutoModelForCausalLM, BaseImageProcessor, Callable, DataCollator, DataCollatorForChatML, Dataset, EvalPrediction, F, FeatureExtractionMixin, GKDConfig, GKDTrainer, GenerationConfig, PeftConfig, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, SFTTrainer, TrainerCallback, disable_dropout_in_model, logger, nn, prepare_deepspeed, torch, F, empty_cache, torch, Any, nn, random, torch, unwrap_model_for_generation) - - -import os -import math -import logging -from typing import * -from dataclasses import dataclass, field -from packaging.version import Version -import torch -import numpy as np -from contextlib import nullcontext -from torch.nn import functional as F -import inspect -from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling -from transformers.training_args import ParallelMode -from unsloth_zoo.device_type import DEVICE_TYPE, DEVICE_TYPE_TORCH, device_synchronize - -# Wrap trainer with padding to right and enable training mode -import functools -from types import MethodType -try: - from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers -except: - def reset_unsloth_gradient_checkpointing_buffers(): pass -# Canonical reset lives in unsloth.models._utils so the SFT auto-packing wrapper and the plain -# Trainer loop can import the same helper; fall back to a no-op only if it can't be imported. -try: - from unsloth.models._utils import _unsloth_reset_stray_compile_cache -except Exception: - def _unsloth_reset_stray_compile_cache(self): pass -try: - from unsloth.models._utils import _unsloth_dataset_column_names -except Exception: - def _unsloth_dataset_column_names(dataset): return dataset.column_names -# Drops/renames config arguments the installed TRL no longer accepts, so a -# script pinned to an older TRL keeps working after an upgrade. Falls back to -# the historical raw passthrough so this can never break trainer construction. -try: - from unsloth.models.rl_config_compat import filter_config_init_kwargs as _unsloth_filter_config_init_kwargs - # A cache file generated here can be imported by an older Unsloth whose filter - # predates `mirrored_from`, so drop the argument rather than raise TypeError. - if "mirrored_from" not in inspect.signature(_unsloth_filter_config_init_kwargs).parameters: - _unsloth_filter_config_init_kwargs_old = _unsloth_filter_config_init_kwargs - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): - return _unsloth_filter_config_init_kwargs_old(config_class, kwargs) -except Exception: - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): return kwargs -def prepare_for_training_mode(f): - @functools.wraps(f) - def wrapper(self, *args, **kwargs): - # Drop any torch.compile graph cache poisoned by a stray pre-train forward. - try: - _unsloth_reset_stray_compile_cache(self) - except Exception: - pass - # Finish the previous W&B run if this is a subsequent train() call. - # We do this at the START of train() (not the end) so that - # evaluate() / log() still work after train() completes. - # HF's WandbCallback.setup() will call wandb.init() for the new run. - # See: https://github.com/unslothai/unsloth/issues/3954 - if getattr(self, '_unsloth_training_completed', False): - try: - import wandb - if wandb.run is not None: - wandb.finish() - # Reset HF's WandbCallback so it calls wandb.init() for the new run - for cb in self.callback_handler.callbacks: - if type(cb).__name__ == 'WandbCallback': - cb._initialized = False - break - except: - pass - # Enable training mode - _was_training = None - # Restore the GC mode the model was configured with at setup; fall back to - # the training args only when it wasn't recorded (issue #4735). Use hasattr, - # not a None sentinel, so a deliberately-recorded None is restored verbatim. - _model = getattr(self, 'model', None) - if hasattr(_model, '_unsloth_gradient_checkpointing'): - use_gc = _model._unsloth_gradient_checkpointing - else: - use_gc = getattr(self.args, 'gradient_checkpointing', True) - if hasattr(self, 'model') and hasattr(self.model, "training"): - _was_training = self.model.training - if hasattr(self, 'model') and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - output = f(self, *args, **kwargs) - # Restore previous mode when possible - if hasattr(self, 'model') and hasattr(self.model, "for_inference"): - if _was_training is False: - self.model.for_inference() - elif _was_training is True and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - # Reset gradient checkpointing buffers to free memory while staying ready for next run - try: - reset_unsloth_gradient_checkpointing_buffers() - except: - pass - # Mark that training completed so the next train() call can - # finish this W&B run before starting a new one - self._unsloth_training_completed = True - return output - return wrapper -pass - -torch_compile_options = { - "epilogue_fusion" : True, - "max_autotune" : False, - "shape_padding" : True, - "trace.enabled" : False, - "triton.cudagraphs" : False, -} - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_hidden_states_selective_log_softmax( - hidden_states: torch.Tensor, - lm_head: torch.Tensor, - index: torch.Tensor, - chunks: int = 4, - logit_scale_multiply: float = 0.0, - logit_scale_divide: float = 0.0, - logit_softcapping: float = 0.0, - temperature: float = 1.0, - # Rows per chunk cap. Read HERE, in the default, not in the body: the body - # is traced with fullgraph = True, and `os.environ` is an unsupported op - # there on torch 2.4 -- Dynamo raises - # torch._dynamo.exc.Unsupported: const method call bytes.decode - # from os._Environ.__getitem__, which the eager fallback does not catch - # (it only catches recompile-limit and disabled-hook breaks), so the very - # first call would die even with the variable unset. A default is evaluated - # once when the def runs, which is import time, outside any traced region. - # It is also a plain int argument, so Dynamo guards on it instead of - # constant-folding an unguarded read (2.7+ never notice a later change). - # A non-numeric value is ignored rather than raised on, so a typo cannot - # break the import. 0 keeps the previous chunk boundaries exactly. - max_rows_per_chunk: int = ( - int(os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip()) - if os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip().isdigit() - else 0 - ), - lm_head_bias: Optional[torch.Tensor] = None, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, - # Read in the default for the same Dynamo reason as max_rows_per_chunk. Pairs with vLLM's - # head_dtype=float32 that load_vllm requests under the same flag. - fp32_lm_head: bool = ( - os.environ.get("UNSLOTH_VLLM_FP32_LM_HEAD", "0").strip().lower() in ("1", "true", "on", "yes") - ), - # Older Unsloth scores old / ref log-probs here and the policy with the fused kernel: one scorer keeps ratio at 1. - fused_linear: bool = ( - os.environ.get("UNSLOTH_FUSED_LINEAR", "1") != "0" - and os.environ.get("UNSLOTH_GRPO_CHUNKED_LOGPROBS", "1") != "0" - ), -) -> torch.Tensor: - # All Unsloth Zoo code licensed under AGPL3 - from unsloth_zoo.fused_losses import chunked_logprob as _chunked_logprob - if (fused_linear and _chunked_logprob._FUSED_LINEAR and not fp32_lm_head - and support_top_k == 0 and support_top_p == 1.0 and support_min_p == 0.0 - and support_excluded_ids is None): - return _chunked_logprob.unsloth_chunked_logprobs( - hidden_states, lm_head, index, - temperature = temperature, - logit_scale_multiply = logit_scale_multiply, - logit_scale_divide = logit_scale_divide, - logit_softcapping = logit_softcapping, - lm_head_bias = lm_head_bias, - max_rows_per_chunk = max_rows_per_chunk if max_rows_per_chunk > 0 else None, - ).reshape((hidden_states.shape[0], hidden_states.shape[1])) - # Reshape on this tensor's own last dim: a no-op, so a wrong-width caller - # cannot have its row count silently rewritten and instead fails at the - # matmul below, which prints both operands. Do not swap in a bare - # torch._check: it reports only "Expected cond to be True", naming neither - # operand, and Dynamo rejects a message-carrying one. Callers dispatch on - # the width first -- see `compute_logprobs_chunk`, the packed path and - # `_pg_grad_forward`. - flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) - flat_index = index.reshape(-1) - - # Each chunk materialises rows x vocab logits and then a float32 copy of - # them, all on the device holding the output head. With a large vocabulary - # and a fixed chunk count that grows with the batch, so the peak scales with - # the batch rather than staying bounded. max_rows_per_chunk caps the rows - # per chunk instead, which is pure loop splitting: more, smaller chunks, - # same concatenated result. 0 (the default) keeps the previous chunk - # boundaries exactly. - if max_rows_per_chunk > 0: - n_rows = flat_hidden_states.shape[0] - chunks = max(chunks, -(-n_rows // max_rows_per_chunk)) - chunks = min(chunks, max(n_rows, 1)) - - # A one-row chunk's head gradient is a K=1 matmul Inductor mis-lowers under dynamic shapes (illegal - # memory access or NaN head gradients, torch 2.13), so keep chunks at 2+ rows when the head trains. - # Exactly `chunks` row ranges, so the unrolled loop length never depends on T: torch.chunk - # returns ceil(T / ceil(T / chunks)) pieces and recompiled per new count; tensor_split guards - # on T % chunks. At most rows // 2 chunks when the head trains keeps every piece at 2+ rows. - lm_head_grad = torch.is_grad_enabled() and lm_head.requires_grad - if lm_head_grad and flat_hidden_states.shape[0] < 2 * chunks: - chunks = max(flat_hidden_states.shape[0] // 2, 1) - - n_rows = flat_hidden_states.shape[0] - chunks = max(1, min(chunks, n_rows)) - bounds = [(i * n_rows) // chunks for i in range(chunks + 1)] - chunked_hidden_states = [flat_hidden_states[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - - # Cast once, outside the loop, so every chunk shares one fp32 head copy (also for autograd). - if fp32_lm_head: - lm_head = lm_head.to(torch.float32) - if lm_head_bias is not None: - lm_head_bias = lm_head_bias.to(torch.float32) - - all_per_token_logps = [] - - for chunk_hidden_states, chunk_index in zip(chunked_hidden_states, chunked_index): - # When the model is dispatched over several devices, the output head can - # sit on a different one from the hidden states, because accelerate - # places the tail of the model on the last device it fills. Co-locate on - # the head's device before the matmul, otherwise this raises - # Unhandled FakeTensor Device Propagation for aten.mm.default, - # found two different devices cuda:0, cuda:1 - # On a single device every .to() here is a no-op and the result is - # bit-identical to before. - chunk_hidden_states = chunk_hidden_states.to(device = lm_head.device, dtype = lm_head.dtype) - chunk_index = chunk_index.to(lm_head.device) - if fp32_lm_head: - # The GRPO loss runs this under bf16 / fp16 autocast, which would downcast the matmul. - with torch.autocast(device_type = chunk_hidden_states.device.type, enabled = False): - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - elif lm_head_bias is None: - chunk_logits = chunk_hidden_states @ lm_head.t() - else: - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - - if logit_scale_multiply != 0.0: - chunk_logits = chunk_logits * logit_scale_multiply - if logit_scale_divide != 0.0: - chunk_logits = chunk_logits / logit_scale_divide - if logit_softcapping != 0.0: - chunk_logits = logit_softcapping * torch.tanh(chunk_logits / logit_softcapping) - - chunk_logits = chunk_logits.to(torch.float32) - - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - per_token_logps = _truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - ) - else: - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - # Return to the caller's device so the concatenation below and every - # downstream consumer see the device they started on. - all_per_token_logps.append(per_token_logps.to(hidden_states.device)) - - all_per_token_logps = torch.concat(all_per_token_logps) - - all_per_token_logps = all_per_token_logps.reshape((hidden_states.shape[0], hidden_states.shape[1])) - return all_per_token_logps - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_selective_log_softmax( - logits, - index, - temperature: float = 1.0, - chunks: int = 4, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, -): - # Fixed `chunks` row ranges; see chunked_hidden_states_selective_log_softmax. - flat_logits = logits.reshape(-1, logits.shape[-1]) - flat_index = index.reshape(-1) - # Fewer rows than chunks: one row per range, as torch.chunk did, never empty ranges. - chunks = max(1, min(chunks, flat_logits.shape[0])) - bounds = [(i * flat_logits.shape[0]) // chunks for i in range(chunks + 1)] - chunked_logits = [flat_logits[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - all_per_token_logps = [] - # Per-chunk selective_log_softmax. - for chunk_logits, chunk_index in zip(chunked_logits, chunked_index): - chunk_logits = chunk_logits.to(torch.float32) - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - all_per_token_logps.append(_truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - )) - continue - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - all_per_token_logps.append(per_token_logps) - pass - all_per_token_logps = torch.concat(all_per_token_logps) - all_per_token_logps = all_per_token_logps.reshape((logits.shape[0], logits.shape[1])) - return all_per_token_logps - -def calculate_pad_tokens_in_prompt( - input_ids: torch.Tensor, - logits_to_keep: int, - pad_token_id: int -) -> torch.Tensor: - """Count left-padded tokens per sequence, e.g. [pad, pad, pad, cat] -> 3.""" - if logits_to_keep >= input_ids.shape[1]: - raise ValueError("logits_to_keep must be smaller than the sequence length.") - - prompt_section = input_ids[:, :-logits_to_keep] - - padding_mask = (prompt_section == pad_token_id) - - pad_token_counts = padding_mask.sum(dim=1) - - return pad_token_counts - -def create_completion_attention_mask( - completion_input_ids: torch.Tensor, - left_pad_tokens_per_prompt: torch.Tensor, - max_left_pad: int, - pad_token_id: int -) -> torch.Tensor: - """Build a completion mask that zeros leading prompt and trailing pad tokens. - - For [p,p,p,c,c,c,pad,pad,pad] (p=sliced prompt, c=completion, pad=padding) - this returns [0,0,0,1,1,1,0,0,0]. - """ - batch_size, completion_len = completion_input_ids.shape - device = completion_input_ids.device - - num_tokens_to_mask = max_left_pad - left_pad_tokens_per_prompt - - indices = torch.arange(completion_len, device=device).unsqueeze(0) - shift_mask = indices >= num_tokens_to_mask.unsqueeze(1) - - non_padding_mask = (completion_input_ids != pad_token_id) - - final_mask = shift_mask & non_padding_mask - - return final_mask - -def left_pack_padding(tensor: torch.Tensor, pad_id: int) -> torch.Tensor: - """Move all padding tokens in each sequence to the right.""" - mask = (tensor != pad_id) - # stable=True since the binary mask is unordered. - sorted_indices = torch.argsort(mask, dim=1, descending=True, stable=True) - packed_tensor = torch.gather(tensor, 1, sorted_indices) - return packed_tensor - -def align_logprobs_with_mask( - logprob_tensor: torch.Tensor, - attention_mask: torch.Tensor, - pad_value: float = 0.0 -) -> torch.Tensor: - """Align a log probability tensor with a given attention mask.""" - - device = logprob_tensor.device - batch_size, logprob_seq_len = logprob_tensor.shape - mask_seq_len = attention_mask.shape[1] - - padded_logprobs = torch.full( - attention_mask.shape, - fill_value=pad_value, - dtype=logprob_tensor.dtype, - device=device - ) - - left_pad_counts = torch.argmax(attention_mask, dim=1) - - cols = torch.arange(logprob_seq_len, device=device) - dest_indices = left_pad_counts.unsqueeze(1) + cols - - # Destination row indices, shape [batch_size, logprob_seq_len]. - row_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand_as(dest_indices) - - # Keep only in-bounds destinations, then scatter via advanced indexing. - valid_mask = dest_indices < mask_seq_len - valid_rows = row_indices[valid_mask] - valid_cols = dest_indices[valid_mask] - valid_vals = logprob_tensor[valid_mask] - padded_logprobs[valid_rows, valid_cols] = valid_vals - - return padded_logprobs - -def align_completion_tool_mask( - tool_mask: torch.Tensor, - completion_mask: torch.Tensor, -) -> torch.Tensor: - """Align a raw completion-length tool/env mask with Unsloth's repacked loss mask.""" - if tool_mask is None: - return completion_mask - if tool_mask.shape[0] != completion_mask.shape[0]: - raise ValueError("tool_mask batch size must match completion_mask batch size.") - - tool_mask = tool_mask.to(device=completion_mask.device) - if tool_mask.shape == completion_mask.shape: - aligned_tool_mask = tool_mask - else: - aligned_tool_mask = align_logprobs_with_mask( - tool_mask, - completion_mask, - pad_value=0, - ) - return completion_mask * aligned_tool_mask.to(dtype=completion_mask.dtype) - -def autotune_batch_and_chunks( - total_input_rows, - seq_len, - hidden_size, - vocab_size, - dtype_bytes=16, - multiplier=None -): - if multiplier is None: - final_m = max(4, seq_len // 4096) - else: - final_m = multiplier - - if torch.cuda.is_available(): - free_bytes, _ = torch.cuda.mem_get_info() - limit_gb = (free_bytes / (1024**3))*.80 - elif hasattr(torch, "xpu") and torch.xpu.is_available(): - # XPU: estimate free memory as total - reserved. - total_mem = torch.xpu.get_device_properties(0).total_memory - reserved_mem = torch.xpu.memory_reserved() - free_bytes = total_mem - reserved_mem - limit_gb = (free_bytes / (1024**3)) * 0.80 - else: - # Fallback: assume 8GB available. - limit_gb = 8.0 - - bytes_to_gb = 1024**3 - - b_vals = torch.arange(total_input_rows, 0, -1, device='cpu', dtype=torch.float32) - - hidden_gb = (b_vals * seq_len * hidden_size * dtype_bytes) / bytes_to_gb - - base_logits = ((b_vals/total_input_rows) * b_vals * seq_len * vocab_size * dtype_bytes) / bytes_to_gb - logits_gb = base_logits / final_m - - total_mem_gb = hidden_gb + logits_gb - - valid_mask = total_mem_gb <= limit_gb - valid_indices = torch.nonzero(valid_mask, as_tuple=False) - - if valid_indices.shape[0] == 0: - #This means your GPU will OOM - # Capped at the row count: unsloth's no-grad pass divides rows by this without max(1, ...). - return max(1, min(4, total_input_rows)), final_m - - best_idx = valid_indices[0].item() - final_b = int(b_vals[best_idx].item()) - - return final_b, final_m - -def sanitize_logprob(logprob): - """Local port of trl.scripts.vllm_serve.sanitize_logprob. - Filters NaN logprobs from vLLM outputs.""" - value = logprob.logprob - if math.isnan(value): - logging.getLogger(__name__).warning( - f"Generated NaN logprob, token logprob '{logprob}' will be ignored" - ) - return None - return value -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_text_configs(config): - """``config`` and its text sub-config, the two places a transform can live: Gemma-4-style configs keep it on ``config.text_config``, T5Gemma-style ones only reach it via ``config.get_text_config()``.""" - if config is None: - return [] - holders = [config] - 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: - holders.append(text_cfg) - return holders - -def _unsloth_get_final_logit_softcapping(model): - """The soft cap the loss applies, under any of its three spellings: ``final_logit_softcapping`` (Gemma), ``logits_soft_cap`` (RecurrentGemma), ``output_logit_soft_cap`` (xLSTM). Returns 0 if unset.""" - config = _unsloth_get_model_config(model) - if config is None: - return 0 - for holder in _unsloth_text_configs(config): - for name in ("final_logit_softcapping", "logits_soft_cap", "output_logit_soft_cap"): - softcap = getattr(holder, name, None) - if softcap: - return softcap - return 0 - -def _unsloth_resolve_logit_scales(model_config): - """``(multiply, divide)`` for the logits, read with the same field table and per-family overrides ``detect_logit_transforms`` uses. - - Only reached when the installed unsloth_zoo predates that helper. Must stay in step with ``resolve_logit_transforms`` in unsloth/models/llama.py: if the forward applies a transform GRPO does not, the policy log-probabilities come from different logits than the ones generated, which silently shifts every importance ratio. - """ - # Granite divides, HyperCLOVA X multiplies, MiniCPM3 scales hidden states (not a logit transform). - overrides = { - ("logits_scaling", "hyperclovax"): "multiply", - ("logits_scaling", "minicpm3"): None, - } - found = {"multiply": 0, "divide": 0} - for holder in _unsloth_text_configs(model_config): - model_type = getattr(holder, "model_type", "") or "" - for bucket, names in ( - ("multiply", ("logit_scale", "lm_head_multiplier", "output_multiplier")), - ("divide", ("logits_scaling",)), - ): - for name in names: - target = overrides.get((name, model_type), bucket) - if target is None or found[target]: - continue - value = getattr(holder, name, 0) or 0 - if value: - found[target] = value - break - return found["multiply"], found["divide"] - -def _unsloth_grpo_returns_hidden_states(model, tensor, lm_head): - """Does ``tensor`` (a forward's ``.logits``) carry hidden states or real logits? - - ``_get_per_token_logps_and_entropies`` sets ``UNSLOTH_RETURN_HIDDEN_STATES=1``, but only a forward that honours the name hands hidden states back as ``.logits``; any other forward returns a real ``[.., vocab]`` tensor that must not reach the ``lm_head`` matmul. - - The primary test is an explicit signal that the forward honours the flag, both set outside this file: ``__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__`` on the generated class, written by ``unsloth_zoo.compiler.create_standalone_class`` exactly when ``apply_fused_lm_head`` gave that forward its own branch; and ``_unsloth_grpo_hidden_states_forward_wrapped``, set by ``_install_grpo_hidden_states_forward_wrapper`` in ``unsloth/models/rl.py`` for models the compiler did not rewrite. That wrapper degrades to real logits when the model cannot produce hidden states and records whether it did so in ``_unsloth_grpo_hidden_states_degraded`` before returning, so reading the pair after a forward describes the call that just finished; degradation is per call, not per model. - - The width comparison stays as the fallback, for an ``unsloth_zoo`` old enough that it never writes the marker. It is decisive whenever ``vocab_size != hidden_size``, and the signal may only overrule it when it is not, which is the one case the shape cannot answer. - """ - if tensor.shape[-1] != lm_head.shape[1]: - return False # vocab-wide: real logits, whatever any signal claims - if lm_head.shape[0] != lm_head.shape[1]: - return True # hidden-wide and vocab_size != hidden_size: hidden states - return _unsloth_grpo_hidden_states_signal(model) is not False - -def _unsloth_grpo_hidden_states_signal(model): - """``True``/``False`` if the forward honours ``UNSLOTH_RETURN_HIDDEN_STATES``, ``None`` when neither marker is present. See ``_unsloth_grpo_returns_hidden_states`` for where each marker is set. Walks the wrapper chain because the markers are set on whichever object the trainer saw, which may be the DDP module or the PEFT base model rather than the object handed to the logprob loop.""" - candidates = [] - pending = [model] - while pending and len(candidates) < 8: - candidate = pending.pop(0) - if candidate is None or any(candidate is seen for seen in candidates): - continue - candidates.append(candidate) - get_base_model = getattr(candidate, "get_base_model", None) - if callable(get_base_model): - try: - pending.append(get_base_model()) - except Exception: - pass - for _attr in ("module", "base_model", "model"): - child = getattr(candidate, _attr, None) - if child is not None and hasattr(child, "forward"): - pending.append(child) - for candidate in candidates: - if getattr(candidate, "__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__", False): - return True - if getattr(type(candidate), "__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__", False): - return True - if any( - getattr(candidate, "_unsloth_grpo_hidden_states_forward_wrapped", False) - for candidate in candidates - ): - # The wrapper honours the flag unless it recorded that this call could not. - if any( - hasattr(candidate, "_unsloth_grpo_hidden_states_degraded") for candidate in candidates - ): - return not any( - getattr(candidate, "_unsloth_grpo_hidden_states_degraded", False) - for candidate in candidates - ) - # An unsloth/models/rl.py predating the per-call attribute set only the warn-once flag; it is the best signal such a wrapper offers. - return not any( - getattr(candidate, "_unsloth_grpo_hidden_states_warning_issued", False) - for candidate in candidates - ) - return None - -def _unsloth_gkd_canonical(source): - """``ast.unparse`` without docstrings/comments so checks survive reformatting; ``None`` if unparsable.""" - import ast as _ast - import textwrap as _textwrap - - try: - tree = _ast.parse(_textwrap.dedent(source)) - except Exception: - return None - for node in _ast.walk(tree): - if isinstance(node, (_ast.FunctionDef, _ast.AsyncFunctionDef)) and node.body: - first = node.body[0] - if ( - isinstance(first, _ast.Expr) - and isinstance(getattr(first, "value", None), _ast.Constant) - and isinstance(first.value.value, str) - ): - node.body = node.body[1:] or [_ast.Pass()] - try: - return _ast.unparse(tree) - except Exception: - return None - -def _unsloth_gkd_jsd_supported(trainer_class): - """Is ``trainer_class.generalized_jsd_loss`` one of TRL's own generalized JSDs, unchanged? Overrides keep the dense loss.""" - cache = trainer_class.__dict__.get("_unsloth_gkd_jsd_supported_cache", None) - if cache is not None: - return cache - required = ( - "student_logits = student_logits / temperature", - "teacher_logits = teacher_logits / temperature", - "student_log_probs = F.log_softmax(student_logits, dim=-1)", - "teacher_log_probs = F.log_softmax(teacher_logits, dim=-1)", - "if beta == 0:", - "jsd = F.kl_div(student_log_probs, teacher_log_probs, reduction='none', log_target=True)", - "elif beta == 1:", - "jsd = F.kl_div(teacher_log_probs, student_log_probs, reduction='none', log_target=True)", - "kl_teacher = F.kl_div(mixture_log_probs, teacher_log_probs, reduction='none', log_target=True)", - "kl_student = F.kl_div(mixture_log_probs, student_log_probs, reduction='none', log_target=True)", - "jsd = beta * kl_teacher + (1 - beta) * kl_student", - "mask = labels != -100", - "jsd = jsd[mask]", - "if reduction == 'batchmean':", - ) - known = set(required) | { - "@staticmethod", - "def generalized_jsd_loss(student_logits, teacher_logits, labels=None, beta=0.5, temperature=1.0, reduction='batchmean'):", - "def generalized_jsd_loss(student_logits, teacher_logits, labels=None, beta=0.5, temperature=1.0, reduction='batchmean', num_items_in_batch=None):", - "else:", - "if labels is not None:", - "beta = torch.tensor(beta, dtype=student_log_probs.dtype)", - "beta = torch.tensor(beta, dtype=student_log_probs.dtype, device=student_log_probs.device)", - "mixture_log_probs = torch.logsumexp(torch.stack([student_log_probs + torch.log(1 - beta), teacher_log_probs + torch.log(beta)]), dim=0)", - "mixture_log_probs = torch.logsumexp(torch.stack([student_log_probs + torch.log1p(-beta), teacher_log_probs + torch.log(beta)]), dim=0)", - "if num_items_in_batch is not None:", - "jsd_sum = jsd.sum()", - "if isinstance(num_items_in_batch, torch.Tensor):", - "num_items_in_batch = num_items_in_batch.to(jsd_sum.device)", - "return jsd_sum / num_items_in_batch", - "return jsd.sum() / mask.sum() if labels is not None else jsd.sum() / (jsd.size(0) * jsd.size(1))", - "return jsd.sum() / mask.sum() if labels is not None else jsd.sum() / jsd.size(0)", - "denom = mask.sum().clamp_min(1) if labels is not None else max(jsd.size(0), 1)", - "return jsd.sum() / denom", - "elif reduction == 'sum':", - "return jsd.sum()", - "elif reduction == 'mean':", - "return jsd.mean()", - "return jsd", - } - ok = False - try: - source = _unsloth_gkd_canonical(inspect.getsource(trainer_class.generalized_jsd_loss)) - lines = [line.strip() for line in source.splitlines() if line.strip()] - signature = next((line for line in lines if line.startswith("def ")), "") - ok = ( - all(line in known for line in lines) - and all(line in lines for line in required) - # a num_items_in_batch parameter must come with its reduction, and vice versa - and ( - ("num_items_in_batch=None" in signature) - == ("return jsd_sum / num_items_in_batch" in lines) - ) - ) - except Exception: - ok = False - try: - setattr(trainer_class, "_unsloth_gkd_jsd_supported_cache", ok) - except Exception: - pass - return ok - -def _unsloth_gkd_note_fallback(trainer, reason): - """Count why the chunked path declined (``trainer._unsloth_gkd_chunked_fallbacks``), log each reason once.""" - try: - fallbacks = trainer.__dict__.setdefault("_unsloth_gkd_chunked_fallbacks", {}) - except Exception: - return None - if reason not in fallbacks: - try: - from unsloth_zoo.log import logger as _gkd_logger - _gkd_logger.info( - f"Unsloth: GKD chunked JSD not used ({reason}); using TRL's dense loss." - ) - except Exception: - pass - fallbacks[reason] = fallbacks.get(reason, 0) + 1 - return None - -def _unsloth_gkd_dense_head(model): - """The output head when it is a plain dense ``[vocab, hidden]`` projection the chunked loss can read directly. - Rejects bnb heads (``nn.Linear`` subclasses with packed weights), PEFT heads and ZeRO-3 partitioned weights. - """ - get_output_embeddings = getattr(model, "get_output_embeddings", None) - if not callable(get_output_embeddings): - return None - try: - head = get_output_embeddings() - except Exception: - return None - if type(head) is not torch.nn.Linear: - return None - # An lm_head that is not the output embeddings runs before them (ModernBERT decoder, RoBERTa-style heads). - lm_head = getattr(model, "lm_head", None) - if isinstance(lm_head, torch.nn.Module) and lm_head is not head: - return None - weight = getattr(head, "weight", None) - if not isinstance(weight, torch.Tensor) or weight.dim() != 2 or weight.numel() == 0: - return None - if weight.dtype not in (torch.float16, torch.bfloat16, torch.float32, torch.float64): - return None - if type(weight) not in (torch.Tensor, torch.nn.Parameter): - return None - bias = getattr(head, "bias", None) - if bias is not None and (bias.dim() != 1 or bias.shape[0] != weight.shape[0]): - return None - if weight.is_meta or (bias is not None and bias.is_meta): - return None - return head - -def _unsloth_gkd_logit_transforms(model): - """``(scale, softcap)`` exactly as the model's own forward applies them to its logits.""" - model_config = _unsloth_get_model_config(model) - if detect_logit_transforms is not None: - _transforms = detect_logit_transforms(model_config) - logit_softcapping = _transforms["logit_softcapping"] - logit_scale_multiply = _transforms["logit_scale_multiply"] - logit_scale_divide = _transforms["logit_scale_divide"] - else: - logit_softcapping = _unsloth_get_final_logit_softcapping(model) - logit_scale_multiply, logit_scale_divide = _unsloth_resolve_logit_scales(model_config) - scale = 1.0 - if logit_scale_multiply: - scale = scale * float(logit_scale_multiply) - if logit_scale_divide: - scale = scale / float(logit_scale_divide) - return scale, float(logit_softcapping or 0.0) - -def _unsloth_gkd_project(hidden_states, head, scale, softcap): - """Full logits from hidden states, for the rare call where only one of the two forwards honoured the flag.""" - hidden_states = hidden_states.to(device = head.weight.device, dtype = head.weight.dtype) - logits = torch.nn.functional.linear(hidden_states, head.weight, head.bias) - if scale != 1.0: - logits = logits * scale - if softcap: - logits = softcap * torch.tanh(logits / softcap) - return logits - -def _unsloth_gkd_chunk_size(vocab_size): - """Rows per chunk: ``UNSLOTH_GKD_CHUNK_SIZE`` when set, else about 2**26 logits per chunk, bounded to [64, 1024].""" - requested = os.environ.get("UNSLOTH_GKD_CHUNK_SIZE", "") - if requested.strip().isdigit() and int(requested) > 0: - return int(requested) - rows = (1 << 26) // max(int(vocab_size), 1) - return int(min(1024, max(64, 1 << (max(rows, 1).bit_length() - 1)))) - -def _unsloth_gkd_right_align( - inputs, - layout, - liger = False, -): - """Roll left-padded rows (TRL's ChatML collator and generate both left-pad) so they start at column 0: Unsloth's training forward drops the 2D mask (#11885). - The prompt layout's ``[:, P:]`` slice becomes ``labels[:, :P] = -100`` plus a one-column ``prompts``; TRL's Liger branch scores every label with no slice, so it only rolls. - """ - try: - attention_mask = inputs["attention_mask"] - left_pad = (attention_mask.cumsum(dim = 1) == 0).sum(dim = 1, keepdim = True) - except Exception: - return inputs - if not bool(left_pad.any()): - return inputs - aligned = dict(inputs) - labels = inputs["labels"] - if layout["shift"] == "prompt" and not liger: - labels = labels.clone() - labels[:, : inputs["prompts"].shape[1]] = -100 - aligned["prompts"] = inputs["prompts"][:, :1] - width = attention_mask.shape[1] - index = (torch.arange(width, device = attention_mask.device).unsqueeze(0) + left_pad) % width - for key, value in ( - ("input_ids", inputs["input_ids"]), - ("attention_mask", attention_mask), - ("labels", labels), - ): - aligned[key] = value.gather(1, index) - return aligned - -def _unsloth_gkd_chunked_loss(self, model, inputs, num_items_in_batch, layout): - """TRL's GKD loss without either full logits tensor, or ``None`` to run TRL's own ``compute_loss``.""" - if layout is None: - return None - if distillation_chunked_jsd is None and _unsloth_fused_linear_jsd is None: - return _unsloth_gkd_note_fallback(self, "unsloth_zoo has no distillation_chunked_jsd") - if os.environ.get("UNSLOTH_GKD_CHUNKED", "1").lower() in ("0", "false", "no", "off"): - return _unsloth_gkd_note_fallback(self, "UNSLOTH_GKD_CHUNKED=0") - # Read per call like the other fused lm_head routes: UNSLOTH_FUSED_LINEAR=0 restores TRL's dense loss. - if os.environ.get("UNSLOTH_FUSED_LINEAR", "1") == "0": - return _unsloth_gkd_note_fallback(self, "UNSLOTH_FUSED_LINEAR=0") - use_fused = _unsloth_fused_linear_jsd is not None - if getattr(self, "use_liger_gkd_loss", False): - return _unsloth_gkd_note_fallback(self, "use_liger_gkd_loss") - if getattr(self, "is_fsdp_enabled", False) or getattr(self, "is_deepspeed_enabled", False): - return _unsloth_gkd_note_fallback(self, "FSDP / DeepSpeed") - jsd_kwargs = {} - scaler = getattr(self.accelerator, "scaler", None) - if scaler is not None and scaler.is_enabled() and not use_fused: - if "scaler" not in inspect.signature(distillation_chunked_jsd).parameters: - return _unsloth_gkd_note_fallback(self, "unsloth_zoo JSD has no GradScaler support") - jsd_kwargs["scaler"] = scaler - try: - missing = any(k not in inputs for k in ("input_ids", "attention_mask", "labels")) - except Exception: - missing = True - if missing: - return _unsloth_gkd_note_fallback(self, "inputs") - if layout["shift"] == "prompt" and "prompts" not in inputs: - return _unsloth_gkd_note_fallback(self, "inputs") - teacher_model = getattr(self, "teacher_model", None) - if teacher_model is None: - return _unsloth_gkd_note_fallback(self, "no teacher_model") - if not _unsloth_gkd_jsd_supported(type(self)): - return _unsloth_gkd_note_fallback(self, "generalized_jsd_loss is not TRL's") - try: - unwrapped_student = self.accelerator.unwrap_model(model) - unwrapped_teacher = self.accelerator.unwrap_model(teacher_model) - except Exception: - return _unsloth_gkd_note_fallback(self, "unwrap_model") - student_head = _unsloth_gkd_dense_head(unwrapped_student) - teacher_head = _unsloth_gkd_dense_head(unwrapped_teacher) - if student_head is None or teacher_head is None: - return _unsloth_gkd_note_fallback(self, "output head is not a dense nn.Linear") - if student_head.weight.shape[0] != teacher_head.weight.shape[0]: - return _unsloth_gkd_note_fallback(self, "vocab mismatch") - # MiniCPM3 divides hidden states by logits_scaling before lm_head; a wrapped forward's hidden_states[-1] predates it. - for unwrapped in (unwrapped_student, unwrapped_teacher): - for config in _unsloth_text_configs(_unsloth_get_model_config(unwrapped)): - if getattr(config, "model_type", None) == "minicpm3": - return _unsloth_gkd_note_fallback( - self, "minicpm3 scales hidden states before the head" - ) - # DDP(find_unused_parameters=True) marks a head skipped in forward as unused, then its grad hook fires twice. - if getattr(model, "find_unused_parameters", False) and any( - p is not None and p.requires_grad for p in (student_head.weight, student_head.bias) - ): - return _unsloth_gkd_note_fallback(self, "DDP find_unused_parameters with a trainable head") - - prior_hidden_states = os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES") - os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" - try: - student_outputs = model( - input_ids = inputs["input_ids"], - attention_mask = inputs["attention_mask"], - ) - teacher_model.eval() - with torch.no_grad(): - teacher_outputs = teacher_model( - input_ids = inputs["input_ids"], - attention_mask = inputs["attention_mask"], - ) - finally: - if prior_hidden_states is None: - os.environ.pop("UNSLOTH_RETURN_HIDDEN_STATES", None) - else: - os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = prior_hidden_states - - student_states = student_outputs.logits - teacher_states = teacher_outputs.logits - if layout["shift"] == "prompt": - prompt_lengths = inputs["prompts"].shape[1] - student_states = student_states[:, prompt_lengths - 1 : -1, :] - teacher_states = teacher_states[:, prompt_lengths - 1 : -1, :] - shifted_labels = inputs["labels"][:, prompt_lengths:] - else: - student_states = student_states[:, :-1, :] - teacher_states = teacher_states[:, :-1, :] - shifted_labels = inputs["labels"][:, 1:] - - student_scale, student_softcap = _unsloth_gkd_logit_transforms(unwrapped_student) - teacher_scale, teacher_softcap = _unsloth_gkd_logit_transforms(unwrapped_teacher) - items = num_items_in_batch if layout["num_items_in_batch"] else None - - student_hidden = _unsloth_grpo_returns_hidden_states( - unwrapped_student, student_states, student_head.weight - ) - teacher_hidden = _unsloth_grpo_returns_hidden_states( - unwrapped_teacher, teacher_states, teacher_head.weight - ) - if not (student_hidden and teacher_hidden): - if student_hidden: - student_states = _unsloth_gkd_project( - student_states, student_head, student_scale, student_softcap - ) - if teacher_hidden: - with torch.no_grad(): - teacher_states = _unsloth_gkd_project( - teacher_states, teacher_head, teacher_scale, teacher_softcap - ) - _unsloth_gkd_note_fallback(self, "a forward returned logits, not hidden states") - # TRL's dense path gets both logits on the input device (accelerate's top-level hook), beside the labels. - student_states = student_states.to(shifted_labels.device) - teacher_states = teacher_states.to(shifted_labels.device) - extra = {} - if layout["num_items_in_batch"]: - extra["num_items_in_batch"] = num_items_in_batch - return self.generalized_jsd_loss( - student_logits = student_states, - teacher_logits = teacher_states, - labels = shifted_labels, - beta = self.beta, - **extra, - ) - - # A dispatched model can return hidden states off its head's device; the zoo's accumulators follow the states. - student_states = student_states.to(student_head.weight.device) - teacher_states = teacher_states.to(teacher_head.weight.device) - if use_fused: - # Both heads one row chunk at a time, student gradient taken in forward (zoo fused_linear). - rows = os.environ.get("UNSLOTH_GKD_CHUNK_SIZE", "") - loss = _unsloth_fused_linear_jsd( - student_states, - student_head.weight, - teacher_states.detach(), - teacher_head.weight.detach(), - labels = shifted_labels, - beta = float(self.beta), - # No TRL release passes a temperature to generalized_jsd_loss (GKDConfig.temperature is for sampling). - temperature = 1.0, - student_bias = student_head.bias, - teacher_bias = None if teacher_head.bias is None else teacher_head.bias.detach(), - student_logit_scale = student_scale, - student_logit_softcapping = student_softcap, - teacher_logit_scale = teacher_scale, - teacher_logit_softcapping = teacher_softcap, - num_items_in_batch = items, - scaling = scaler.get_scale() if scaler is not None and scaler.is_enabled() else None, - rows_per_chunk = int(rows) if rows.strip().isdigit() else None, - ) - try: - self._unsloth_gkd_chunked_calls = getattr(self, "_unsloth_gkd_chunked_calls", 0) + 1 - self._unsloth_gkd_fused_calls = getattr(self, "_unsloth_gkd_fused_calls", 0) + 1 - except Exception: - pass - return loss - loss, _entropy_sum, _n_valid = distillation_chunked_jsd( - student_states, - teacher_states.detach(), - student_head.weight, - teacher_head.weight.detach(), - shifted_labels != -100, - beta = float(self.beta), - chunk_size = _unsloth_gkd_chunk_size(student_head.weight.shape[0]), - num_items_in_batch = items, - student_lm_head_bias = student_head.bias, - teacher_lm_head_bias = None if teacher_head.bias is None else teacher_head.bias.detach(), - student_logit_scale = student_scale, - teacher_logit_scale = teacher_scale, - student_final_logit_softcapping = student_softcap, - teacher_final_logit_softcapping = teacher_softcap, - # No TRL release passes a temperature to generalized_jsd_loss (GKDConfig.temperature is for sampling). - temperature = 1.0, - **jsd_kwargs, - ) - try: - self._unsloth_gkd_chunked_calls = getattr(self, "_unsloth_gkd_chunked_calls", 0) + 1 - except Exception: - pass - return loss - -import re -import inspect -try: - from unsloth_zoo.device_map_planner import detect_logit_transforms -except Exception: - detect_logit_transforms = None -try: - from unsloth_zoo.rl_replacements import distillation_chunked_jsd -except Exception: - distillation_chunked_jsd = None -try: - from unsloth_zoo.fused_losses.fused_linear import fused_linear_jsd as _unsloth_fused_linear_jsd -except Exception: - _unsloth_fused_linear_jsd = None -@dataclass -class UnslothGKDConfig(GKDConfig): - """ - - Configuration class for [`experimental.gkd.GKDTrainer`]. - - This class includes only the parameters that are specific to GKD training. For a full list of training arguments, - please refer to the [`~transformers.TrainingArguments`] and [`SFTConfig`] documentation. - - Args: - temperature (`float`, *optional*, defaults to `0.9`): - Temperature for sampling. The higher the temperature, the more random the completions. - lmbda (`float`, *optional*, defaults to `0.5`): - Lambda parameter that controls the student data fraction (i.e., the proportion of on-policy - student-generated outputs). - beta (`float`, *optional*, defaults to `0.5`): - Interpolation coefficient between `0.0` and `1.0` of the Generalized Jensen-Shannon Divergence loss. When - beta is `0.0`, the loss is the KL divergence. When beta is `1.0`, the loss is the Inverse KL Divergence. - max_new_tokens (`int`, *optional*, defaults to `128`): - Maximum number of tokens to generate per completion. - teacher_model_name_or_path (`str`, *optional*): - Model name or path of the teacher model. If `None`, the teacher model will be the same as the model being - trained. - teacher_model_init_kwargs (`dict[str, Any]`, *optional*): - Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the teacher model - from a string. - disable_dropout (`bool`, *optional*, defaults to `True`): - Whether to disable dropout in the model. - seq_kd (`bool`, *optional*, defaults to `False`): - Seq_kd parameter that controls whether to perform Sequence-Level KD (can be viewed as supervised FT on - teacher-generated output). - - """ - _unsloth_accepts_adamw_flash = True - vllm_sampling_params: Optional[Any] = field( - default = None, - metadata = {'help': 'vLLM SamplingParams'}, - ) - unsloth_num_chunks : Optional[int] = field( - default = -1, - metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'}, - ) - unsloth_logit_chunk_multiplier : Optional[int] = field( - default = None, - metadata = {'help': 'Multiplier for chunked logit computations.'}, - ) - unsloth_grpo_mini_batch : Optional[int] = field( - default = None, - metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'}, - ) - max_seq_length : Optional[int] = field( - default = None, - metadata = {'help': 'Maximum sequence length to truncate to.'}, - ) - def __init__( - self, - output_dir = None, - per_device_train_batch_size = 4, - num_train_epochs = 3.0, - max_steps = -1, - learning_rate = 5e-05, - lr_scheduler_type = 'linear', - lr_scheduler_kwargs = None, - warmup_steps = 0.1, - optim = 'adamw_8bit', - optim_args = None, - weight_decay = 0.001, - adam_beta1 = 0.9, - adam_beta2 = 0.999, - adam_epsilon = 1e-08, - optim_target_modules = None, - gradient_accumulation_steps = 2, - average_tokens_across_devices = True, - max_grad_norm = 1.0, - label_smoothing_factor = 0.0, - bf16 = False, - fp16 = False, - bf16_full_eval = False, - fp16_full_eval = False, - tf32 = None, - gradient_checkpointing = True, - gradient_checkpointing_kwargs = None, - torch_compile = False, - torch_compile_backend = None, - torch_compile_mode = None, - use_liger_kernel = False, - liger_kernel_config = None, - use_cache = False, - neftune_noise_alpha = None, - torch_empty_cache_steps = 250, - auto_find_batch_size = False, - logging_strategy = 'steps', - logging_steps = 1, - logging_first_step = False, - log_on_each_node = True, - logging_nan_inf_filter = False, - include_num_input_tokens_seen = False, - log_level = 'passive', - log_level_replica = 'warning', - disable_tqdm = None, - report_to = 'none', - run_name = None, - project = 'huggingface', - trackio_space_id = 'trackio', - eval_strategy = 'no', - eval_steps = None, - eval_delay = 0, - per_device_eval_batch_size = 4, - prediction_loss_only = False, - eval_on_start = False, - eval_do_concat_batches = True, - eval_use_gather_object = False, - eval_accumulation_steps = 2, - batch_eval_metrics = False, - save_only_model = False, - save_strategy = 'steps', - save_steps = 500, - save_on_each_node = False, - save_total_limit = None, - enable_jit_checkpoint = False, - push_to_hub = False, - hub_token = None, - hub_private_repo = None, - hub_model_id = None, - hub_strategy = 'every_save', - hub_always_push = False, - hub_revision = None, - load_best_model_at_end = False, - metric_for_best_model = None, - greater_is_better = None, - ignore_data_skip = False, - restore_callback_states_from_checkpoint = False, - full_determinism = False, - seed = 3407, - data_seed = None, - use_cpu = False, - accelerator_config = None, - parallelism_config = None, - dataloader_drop_last = False, - dataloader_num_workers = 0, - dataloader_pin_memory = True, - dataloader_persistent_workers = False, - dataloader_prefetch_factor = None, - remove_unused_columns = True, - label_names = None, - train_sampling_strategy = 'random', - length_column_name = 'length', - ddp_find_unused_parameters = None, - ddp_bucket_cap_mb = None, - ddp_broadcast_buffers = None, - ddp_backend = None, - ddp_timeout = 1800, - fsdp = None, - fsdp_config = None, - deepspeed = None, - debug = '', - skip_memory_metrics = True, - do_train = False, - do_eval = False, - do_predict = False, - resume_from_checkpoint = None, - warmup_ratio = None, - logging_dir = None, - local_rank = -1, - model_init_kwargs = None, - trust_remote_code = False, - router_aux_loss_coef = None, - chat_template_path = None, - dataset_text_field = 'text', - dataset_kwargs = None, - dataset_num_proc = None, - eos_token = None, - max_length = 1024, - truncation_mode = 'keep_start', - shuffle_dataset = False, - packing = False, - packing_strategy = 'bfd', - padding_free = None, - pad_to_multiple_of = None, - eval_packing = None, - completion_only_loss = None, - assistant_only_loss = False, - loss_type = 'nll', - activation_offloading = False, - pad_token = None, - temperature = 0.9, - lmbda = 0.5, - beta = 0.5, - max_new_tokens = 128, - teacher_model_name_or_path = None, - teacher_model_init_kwargs = None, - disable_dropout = True, - seq_kd = False, - vllm_sampling_params = None, - unsloth_num_chunks = -1, - unsloth_logit_chunk_multiplier = None, - unsloth_grpo_mini_batch = None, - max_seq_length = None, - **kwargs, - ): - if learning_rate < 1e-7: print(f'Unsloth: Your learning rate of `{learning_rate}` is too small and less than 1e-7! Consider increasing it, otherwise gradient updates will be close to 0!') - if learning_rate > 1: print(f'Unsloth: Your learning rate of `{learning_rate}` is way too larger > 1! Consider decreasing it to 1e-1, otherwise gradient updates will explode!') - _unsloth_fsdp2 = os.environ.get('FSDP_VERSION', '').strip() == '2' - _unsloth_fsdp_cfg = fsdp_config - if fsdp and isinstance(_unsloth_fsdp_cfg, str) and os.path.isfile(_unsloth_fsdp_cfg): - try: - import json as _unsloth_json - with open(_unsloth_fsdp_cfg, encoding = 'utf-8') as _unsloth_f: - _unsloth_fsdp_cfg = _unsloth_json.load(_unsloth_f) - except Exception: - _unsloth_fsdp_cfg = None - if fsdp and _unsloth_fsdp_cfg is None: - _unsloth_fsdp_cfg = {} - if fsdp and isinstance(_unsloth_fsdp_cfg, dict): - _unsloth_fsdp2 = _unsloth_fsdp2 or str(_unsloth_fsdp_cfg.get('fsdp_version', _unsloth_fsdp_cfg.get('version', 1))) == '2' - _unsloth_flash_adamw = str(getattr(optim, 'value', optim)).strip().lower() == 'adamw_flash' - if _unsloth_flash_adamw: optim = 'adamw_torch' - if _unsloth_fsdp2 and str(getattr(optim, 'value', optim)) == 'adamw_8bit': - print("Unsloth: FSDP2 shards parameters as DTensors, which bitsandbytes' 8-bit optimizers cannot step. Switching optim = 'adamw_8bit' to 'adamw_torch_fused'.") - optim = 'adamw_torch_fused' - if num_train_epochs is None: - num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override - if output_dir is None and save_strategy == 'steps' and save_steps == 500: - output_dir = 'unsloth_training_checkpoints' - save_strategy = 'no' - try: - from unsloth_zoo.dataset_num_proc import get_dataset_num_proc as _unsloth_get_dataset_num_proc - except Exception: - try: - from unsloth.dataset_num_proc import get_dataset_num_proc as _unsloth_get_dataset_num_proc - except Exception: - _unsloth_get_dataset_num_proc = None - if _unsloth_get_dataset_num_proc is not None: - dataset_num_proc = _unsloth_get_dataset_num_proc(dataset_num_proc, serial_as_none = True) - if os.environ.get('UNSLOTH_ENABLE_FLEX_ATTENTION', '0') == '1': - from unsloth_zoo.flex_attention import HAS_FLEX_ATTENTION - if HAS_FLEX_ATTENTION and pad_to_multiple_of is None: - from unsloth_zoo.flex_attention import FLEX_ATTENTION_BLOCK_SIZE - pad_to_multiple_of = FLEX_ATTENTION_BLOCK_SIZE - - if eval_steps is not None and eval_strategy != 'steps': - print(f'Unsloth: `eval_steps = {eval_steps}` is ignored because `eval_strategy` is {getattr(eval_strategy, "value", eval_strategy)!r}. Set `eval_strategy = "steps"` to evaluate every `eval_steps` steps.') - - if temperature <= 0: - raise ValueError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.') - elif temperature >= 10: - raise ValueError('Unsloth: Please set a positive non-zero temperature less than 10, since sampling will be quite erratic.') - - - # One dict so the filter sees the mirrored parameters AND `**kwargs`: - # filtering kwargs alone would double-bind any argument TRL renamed, - # since the new name is itself a mirrored parameter. - _unsloth_config_arguments = dict( - output_dir = output_dir, - per_device_train_batch_size = per_device_train_batch_size, - num_train_epochs = num_train_epochs, - max_steps = max_steps, - learning_rate = learning_rate, - lr_scheduler_type = lr_scheduler_type, - lr_scheduler_kwargs = lr_scheduler_kwargs, - warmup_steps = warmup_steps, - optim = optim, - optim_args = optim_args, - weight_decay = weight_decay, - adam_beta1 = adam_beta1, - adam_beta2 = adam_beta2, - adam_epsilon = adam_epsilon, - optim_target_modules = optim_target_modules, - gradient_accumulation_steps = gradient_accumulation_steps, - average_tokens_across_devices = average_tokens_across_devices, - max_grad_norm = max_grad_norm, - label_smoothing_factor = label_smoothing_factor, - bf16 = bf16, - fp16 = fp16, - bf16_full_eval = bf16_full_eval, - fp16_full_eval = fp16_full_eval, - tf32 = tf32, - gradient_checkpointing = gradient_checkpointing, - gradient_checkpointing_kwargs = gradient_checkpointing_kwargs, - torch_compile = torch_compile, - torch_compile_backend = torch_compile_backend, - torch_compile_mode = torch_compile_mode, - use_liger_kernel = use_liger_kernel, - liger_kernel_config = liger_kernel_config, - use_cache = use_cache, - neftune_noise_alpha = neftune_noise_alpha, - torch_empty_cache_steps = torch_empty_cache_steps, - auto_find_batch_size = auto_find_batch_size, - logging_strategy = logging_strategy, - logging_steps = logging_steps, - logging_first_step = logging_first_step, - log_on_each_node = log_on_each_node, - logging_nan_inf_filter = logging_nan_inf_filter, - include_num_input_tokens_seen = include_num_input_tokens_seen, - log_level = log_level, - log_level_replica = log_level_replica, - disable_tqdm = disable_tqdm, - report_to = report_to, - run_name = run_name, - project = project, - trackio_space_id = trackio_space_id, - eval_strategy = eval_strategy, - eval_steps = eval_steps, - eval_delay = eval_delay, - per_device_eval_batch_size = per_device_eval_batch_size, - prediction_loss_only = prediction_loss_only, - eval_on_start = eval_on_start, - eval_do_concat_batches = eval_do_concat_batches, - eval_use_gather_object = eval_use_gather_object, - eval_accumulation_steps = eval_accumulation_steps, - batch_eval_metrics = batch_eval_metrics, - save_only_model = save_only_model, - save_strategy = save_strategy, - save_steps = save_steps, - save_on_each_node = save_on_each_node, - save_total_limit = save_total_limit, - enable_jit_checkpoint = enable_jit_checkpoint, - push_to_hub = push_to_hub, - hub_token = hub_token, - hub_private_repo = hub_private_repo, - hub_model_id = hub_model_id, - hub_strategy = hub_strategy, - hub_always_push = hub_always_push, - hub_revision = hub_revision, - load_best_model_at_end = load_best_model_at_end, - metric_for_best_model = metric_for_best_model, - greater_is_better = greater_is_better, - ignore_data_skip = ignore_data_skip, - restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint, - full_determinism = full_determinism, - seed = seed, - data_seed = data_seed, - use_cpu = use_cpu, - accelerator_config = accelerator_config, - parallelism_config = parallelism_config, - dataloader_drop_last = dataloader_drop_last, - dataloader_num_workers = dataloader_num_workers, - dataloader_pin_memory = dataloader_pin_memory, - dataloader_persistent_workers = dataloader_persistent_workers, - dataloader_prefetch_factor = dataloader_prefetch_factor, - remove_unused_columns = remove_unused_columns, - label_names = label_names, - train_sampling_strategy = train_sampling_strategy, - length_column_name = length_column_name, - ddp_find_unused_parameters = ddp_find_unused_parameters, - ddp_bucket_cap_mb = ddp_bucket_cap_mb, - ddp_broadcast_buffers = ddp_broadcast_buffers, - ddp_backend = ddp_backend, - ddp_timeout = ddp_timeout, - fsdp = fsdp, - fsdp_config = fsdp_config, - deepspeed = deepspeed, - debug = debug, - skip_memory_metrics = skip_memory_metrics, - do_train = do_train, - do_eval = do_eval, - do_predict = do_predict, - resume_from_checkpoint = resume_from_checkpoint, - warmup_ratio = warmup_ratio, - logging_dir = logging_dir, - local_rank = local_rank, - model_init_kwargs = model_init_kwargs, - trust_remote_code = trust_remote_code, - router_aux_loss_coef = router_aux_loss_coef, - chat_template_path = chat_template_path, - dataset_text_field = dataset_text_field, - dataset_kwargs = dataset_kwargs, - dataset_num_proc = dataset_num_proc, - eos_token = eos_token, - max_length = max_length, - truncation_mode = truncation_mode, - shuffle_dataset = shuffle_dataset, - packing = packing, - packing_strategy = packing_strategy, - padding_free = padding_free, - pad_to_multiple_of = pad_to_multiple_of, - eval_packing = eval_packing, - completion_only_loss = completion_only_loss, - assistant_only_loss = assistant_only_loss, - loss_type = loss_type, - activation_offloading = activation_offloading, - pad_token = pad_token, - temperature = temperature, - lmbda = lmbda, - beta = beta, - max_new_tokens = max_new_tokens, - teacher_model_name_or_path = teacher_model_name_or_path, - teacher_model_init_kwargs = teacher_model_init_kwargs, - disable_dropout = disable_dropout, - seq_kd = seq_kd,**kwargs) - super().__init__(**_unsloth_filter_config_init_kwargs(GKDConfig, _unsloth_config_arguments, mirrored_from = __class__)) - self.vllm_sampling_params = vllm_sampling_params - self.unsloth_num_chunks = unsloth_num_chunks - if unsloth_grpo_mini_batch is not None: - if self.generation_batch_size >= unsloth_grpo_mini_batch: - self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch - else: - raise ValueError( - f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, " - f"which is self.per_device_train_batch_size * gradient_accumulation_steps." - ) - self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier - self.max_seq_length = max_seq_length - # Unsloth: keep the reentrant checkpoint path - if getattr(self, 'gradient_checkpointing', False): - _gc_kwargs = getattr(self, 'gradient_checkpointing_kwargs', None) or {} - if _gc_kwargs.get('context_fn') is None and not _gc_kwargs.get('debug', False): - _gc_kwargs['use_reentrant'] = True - self.gradient_checkpointing_kwargs = _gc_kwargs - from unsloth.optimizers.flash_adamw import is_flash_adamw_marked, mark_flash_adamw, unmark_flash_adamw - if _unsloth_flash_adamw or is_flash_adamw_marked(self): - mark_flash_adamw(self) - else: - unmark_flash_adamw(self) - -pass - -class _UnslothGKDTrainer(SFTTrainer): - """""" - - _tag_names = ["trl", "gkd"] - _name = "GKD" - _paper = { - "title": "On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes", - "id": "2306.13649", - # docstyle-ignore - "citation": textwrap.dedent("""\ - @inproceedings{agarwal2024on-policy, - title = {{On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes}}, - author = {Rishabh Agarwal and Nino Vieillard and Yongchao Zhou and Piotr Stanczyk and Sabela Ramos Garea and Matthieu Geist and Olivier Bachem}, - year = 2024, - booktitle = {The Twelfth International Conference on Learning Representations, {ICLR} 2024, Vienna, Austria, May 7-11, 2024}, - publisher = {OpenReview.net}, - url = {https://openreview.net/forum?id=3zKtaqxLhW}, - }"""), - } - - def __init__( - self, - model: PreTrainedModel | nn.Module | str | None = None, - teacher_model: PreTrainedModel | nn.Module | str = None, - args: GKDConfig | None = None, - data_collator: DataCollator | None = None, # type: ignore - train_dataset: Dataset | None = None, - eval_dataset: Dataset | dict[str, Dataset] | None = None, - processing_class: PreTrainedTokenizerBase - | BaseImageProcessor - | FeatureExtractionMixin - | ProcessorMixin - | None = None, - compute_metrics: Callable[[EvalPrediction], dict] | None = None, - callbacks: list[TrainerCallback] | None = None, - optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), - preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, - peft_config: "PeftConfig | None" = None, - formatting_func: Callable | None = None, - ): - # Ensure Trainer does not drop non-signature columns used by the collator [e.g., "prompts"] - args.remove_unused_columns = False - # Respect a user-provided data_collator; otherwise, provide a ChatML collator that - if data_collator is None: - data_collator = DataCollatorForChatML(tokenizer=processing_class, max_length=args.max_length) - - # Ensure SFTTrainer does not pre-process the dataset when using a ChatML collator, - # so that raw conversational fields [e.g., "messages"] remain available to the collator. - if args.dataset_kwargs is None: - args.dataset_kwargs = {"skip_prepare_dataset": True} - else: - args.dataset_kwargs["skip_prepare_dataset"] = True - - super().__init__( - model, - args=args, - data_collator=data_collator, - train_dataset=train_dataset, - eval_dataset=eval_dataset, - processing_class=processing_class, - compute_metrics=compute_metrics, - callbacks=callbacks, - optimizers=optimizers, - preprocess_logits_for_metrics=preprocess_logits_for_metrics, - peft_config=peft_config, - formatting_func=formatting_func, - ) - - if args.teacher_model_init_kwargs is None: - teacher_model_init_kwargs = {} - elif not isinstance(teacher_model, str): - raise ValueError( - "You passed teacher_model_init_kwargs to the GKDConfig, but your teacher_model is already instantiated." - ) - else: - teacher_model_init_kwargs = args.teacher_model_init_kwargs - teacher_model_init_kwargs["dtype"] = ( - teacher_model_init_kwargs["dtype"] - if teacher_model_init_kwargs["dtype"] in ["auto", None] - else getattr(torch, teacher_model_init_kwargs["dtype"]) - ) - - if isinstance(teacher_model, str): - teacher_model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) - teacher_model = AutoModelForCausalLM.from_pretrained(teacher_model, **teacher_model_init_kwargs) - - student_vocab_size = self.model.config.get_text_config().vocab_size - teacher_vocab_size = teacher_model.config.get_text_config().vocab_size - if student_vocab_size != teacher_vocab_size: - raise ValueError( - f"The student model has vocab_size {student_vocab_size} but the teacher model has " - f"vocab_size {teacher_vocab_size}. GKD compares the teacher's full next-token " - f"distribution, which requires a shared vocabulary. Use a teacher with the same vocab_size, or " - f"GOLD for cross-tokenizer distillation." - ) - - # Disable dropout in the model - if args.disable_dropout: - disable_dropout_in_model(self.model) - - if self.is_deepspeed_enabled: - self.teacher_model = prepare_deepspeed(teacher_model, self.accelerator) - else: - self.teacher_model = self.accelerator.prepare_model(teacher_model, evaluation_mode=True) - - self.lmbda = args.lmbda - self.beta = args.beta - self.temperature = args.temperature - self.seq_kd = args.seq_kd - - # With `lmbda=1.0` training is fully on-policy and `seq_kd` is never reached, and with `temperature=1.0` the - # sampling temperature matches the one `DistillationTrainer` also applies to the divergence. In that setting, - # `DistillationTrainer` covers this run and additionally supports vLLM generation and a chunked loss that - # never materializes the full logits. - if self.lmbda == 1.0 and self.temperature == 1.0: - logger.warning( - "This GKD configuration (`lmbda=1.0`, `temperature=1.0`) is fully covered by `DistillationTrainer`, " - "which is maintained in the main codebase and supports vLLM generation and a memory-efficient " - "chunked loss. Consider migrating: replace `GKDConfig`/`GKDTrainer` with " - "`DistillationConfig`/`DistillationTrainer`, `max_new_tokens` with `max_completion_length`, and set " - f"`beta={self.beta}` explicitly (`beta` means the same in both, but defaults to 0.5 here and 1.0 " - "there)." - ) - - generation_kwargs = { - "max_new_tokens": args.max_new_tokens, - "temperature": args.temperature, - "do_sample": True, - "top_k": 0, - "top_p": 1.0, - "use_cache": False if args.gradient_checkpointing else True, - "pad_token_id": self.processing_class.pad_token_id, - } - self.generation_config = GenerationConfig(**generation_kwargs) - # Keep training-specific generation kwargs to overwrite model's original generation config - self.generation_kwargs = generation_kwargs - # Set custom EOS tokens if they are specified by the model's generation - # config. This is important for models with the Llama 3 chat template, - # which use special tokens <|eot_id|> and <|eom_id|> to mark the end of - # turns or messages. - if ( - hasattr(self.model.generation_config, "eos_token_id") - and self.model.generation_config.eos_token_id is not None - ): - self.generation_config.eos_token_id = self.model.generation_config.eos_token_id - - @staticmethod - def generalized_jsd_loss( - student_logits, - teacher_logits, - labels=None, - beta=0.5, - temperature=1.0, - reduction="batchmean", - num_items_in_batch=None, - ): - """ - Compute the generalized Jensen-Shannon Divergence loss for knowledge distillation using F.kl_div. See Eq. (1) - of https://huggingface.co/papers/2306.13649 for the definition. - - Args: - student_logits: - Tensor of shape (batch_size, sequence_length, vocab_size) - teacher_logits: - Tensor of shape (batch_size, sequence_length, vocab_size) - labels: - Tensor of shape (batch_size, sequence_length) with -100 for padding tokens to ignore when computing - loss - beta: - Interpolation coefficient between 0 and 1 (default: 0.5) - temperature: - Softmax temperature (default: 1.0) - reduction: - Specifies the reduction to apply to the output (default: 'batchmean') - - Returns: - loss: Scalar tensor with the generalized JSD loss - """ - - # Apply temperature scaling - student_logits = student_logits / temperature - teacher_logits = teacher_logits / temperature - - # Compute log probabilities for student and probabilities for teacher - student_log_probs = F.log_softmax(student_logits, dim=-1) - teacher_log_probs = F.log_softmax(teacher_logits, dim=-1) - - if beta == 0: - jsd = F.kl_div(student_log_probs, teacher_log_probs, reduction="none", log_target=True) - elif beta == 1: - jsd = F.kl_div(teacher_log_probs, student_log_probs, reduction="none", log_target=True) - else: - # Compute the log of the mixture distribution - # log(a + b) = log(exp(log(a)) + exp(log(b))) -> for mixture - beta = torch.tensor(beta, dtype=student_log_probs.dtype, device=student_log_probs.device) - mixture_log_probs = torch.logsumexp( - torch.stack([student_log_probs + torch.log1p(-beta), teacher_log_probs + torch.log(beta)]), - dim=0, - ) - - # Compute KL divergences using F.kl_div - # PyTorch differs from the standard mathematical definition, so the order of the probability distributions is swapped compared to that defined in the paper. - kl_teacher = F.kl_div(mixture_log_probs, teacher_log_probs, reduction="none", log_target=True) - kl_student = F.kl_div(mixture_log_probs, student_log_probs, reduction="none", log_target=True) - - # Compute the Generalized Jensen-Shannon Divergence - jsd = beta * kl_teacher + (1 - beta) * kl_student - - # Masking - if labels is not None: - mask = labels != -100 - jsd = jsd[mask] - - # Apply reduction - if num_items_in_batch is not None: - # Normalize by the global number of valid tokens for gradient-accumulation-correct loss (see issue #4719). - jsd_sum = jsd.sum() - if isinstance(num_items_in_batch, torch.Tensor): - num_items_in_batch = num_items_in_batch.to(jsd_sum.device) - return jsd_sum / num_items_in_batch - if reduction == "batchmean": - # clamp_min(1) avoids 0/0 -> nan when a sample has no unmasked positions - # (e.g. completion fully truncated). jsd[mask] is empty -> jsd.sum() == 0, - # so 0/1 == 0 with a valid grad path. - denom = mask.sum().clamp_min(1) if labels is not None else max(jsd.size(0), 1) - return jsd.sum() / denom - elif reduction == "sum": - return jsd.sum() - elif reduction == "mean": - return jsd.mean() - else: - return jsd - - def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): - # Unsloth: chunked generalized JSD over hidden states; TRL's own loss below is the fallback. - if not return_outputs: - inputs = _unsloth_gkd_right_align( - inputs, {'shift': 'shift', 'num_items_in_batch': True}, getattr(self, 'use_liger_gkd_loss', False), - ) - loss = _unsloth_gkd_chunked_loss( - self, model, inputs, num_items_in_batch, {'shift': 'shift', 'num_items_in_batch': True}, - ) - if loss is not None: - return loss - return self._unsloth_trl_compute_loss( - model, inputs, return_outputs=return_outputs, num_items_in_batch=num_items_in_batch, - ) - - def _unsloth_trl_compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): - # compute student output - student_outputs = model( - input_ids=inputs["input_ids"], - attention_mask=inputs["attention_mask"], - ) - - # compute teacher output in eval mode - self.teacher_model.eval() - with torch.no_grad(): - teacher_outputs = self.teacher_model( - input_ids=inputs["input_ids"], - attention_mask=inputs["attention_mask"], - ) - - # Standard causal shift: logits at position i predict the token at i + 1. The `labels != -100` mask - # inside `generalized_jsd_loss` already excludes prompt (and padding) positions, so we do not slice by - # prompt length. Slicing by `inputs["prompts"].shape[1]` (the batch-max prompt width) would drop real - # completion tokens for samples whose prompt is shorter than the batch maximum, since `labels` is padded - # to the full-sequence width independently of `prompts`. - shifted_student_logits = student_outputs.logits[:, :-1, :] - shifted_teacher_logits = teacher_outputs.logits[:, :-1, :] - shifted_labels = inputs["labels"][:, 1:] - - # compute loss - loss = self.generalized_jsd_loss( - student_logits=shifted_student_logits, - teacher_logits=shifted_teacher_logits, - labels=shifted_labels, - beta=self.beta, - num_items_in_batch=num_items_in_batch, - ) - - # empty cache - empty_cache() - - # Return loss - return (loss, student_outputs) if return_outputs else loss - - @staticmethod - def generate_on_policy_outputs(model, inputs, generation_config): - # Generate output with respect to the prompt-only - generated_outputs = model.generate( - input_ids=inputs["prompts"], - attention_mask=inputs["prompt_attention_mask"], - generation_config=generation_config, - return_dict_in_generate=True, - ) - - # Get the generated token IDs - generated_tokens = generated_outputs.sequences - device = generated_tokens.device - prompt_mask = inputs["prompt_attention_mask"] - prompt_length = inputs["prompts"].shape[1] - completion_ids = generated_tokens[:, prompt_length:] - - # Mask everything after the first EOS token. eos_token_id can be a list of stop tokens (Llama 3), - # so match with torch.isin; with no eos id nothing stops early, so keep the whole completion. - if generation_config.eos_token_id is None: - completion_mask = torch.ones_like(completion_ids) - else: - is_eos = torch.isin(completion_ids, torch.tensor(generation_config.eos_token_id, device=device)) - eos_idx = torch.full((is_eos.size(0),), is_eos.size(1), dtype=torch.long, device=device) - eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)] - sequence_indices = torch.arange(is_eos.size(1), device=device).expand(is_eos.size(0), -1) - completion_mask = (sequence_indices <= eos_idx.unsqueeze(1)).int() - - # Pad ids can appear inside real prompt text, so use the prompt mask instead of matching ids - new_attention_mask = torch.ones_like(generated_tokens) - new_attention_mask[:, prompt_length:] = completion_mask - new_attention_mask[:, :prompt_length] = prompt_mask - - new_labels = generated_tokens.clone() - new_labels[new_attention_mask == 0] = -100 - - # Mask the prompt so only the generated completion contributes to the loss. `generate` echoes - # the prompt back as the first `prompt_length` columns, so masking them with -100 matches the - # collator convention (`labels[:len(prompt)] = -100`) that `compute_loss` relies on. - new_labels[:, :prompt_length] = -100 - - return generated_tokens, new_attention_mask, new_labels - - def training_step( - self, model: nn.Module, inputs: dict[str, torch.Tensor | Any], num_items_in_batch: int | None = None - ) -> torch.Tensor: - """ - Perform a training step for the Generalized Knowledge Distillation (GKD) model. - - This method implements the on-policy learning approach described in the GKD paper. With probability - `self.lmbda`, it generates new responses using the student model, which are then used for training instead of - the original inputs. - """ - pending = getattr(self, "_unsloth_gkd_pending_batches", None) - if pending is not None: - batches, batch_args, batch_kwargs = pending - self._unsloth_gkd_pending_batches = None - - def prepared_batches(): - for batch in batches: - self._unsloth_prepare_gkd_batch(model, batch) - yield batch - - _, count = super().get_batch_samples( - prepared_batches(), len(batches), *batch_args, **batch_kwargs - ) - # TRL 1.7-1.8 mask EOS when it is the pad token, so an all-EOS window counts 0: sum / 0 is NaN. - if torch.is_tensor(count): - count = count.clamp_min(1) - elif count is not None: - count = max(count, 1) - for batch in batches: - batch["_unsloth_gkd_num_items"] = count - - if "_unsloth_gkd_num_items" in inputs: - num_items_in_batch = inputs.pop("_unsloth_gkd_num_items") - else: - self._unsloth_prepare_gkd_batch(model, inputs) - loss = super().training_step(model, inputs, num_items_in_batch) - return loss - - def _unsloth_prepare_gkd_batch(self, model, inputs): - if random.random() <= self.lmbda: - with ( - unwrap_model_for_generation( - model, - self.accelerator, - generation_kwargs=self.generation_kwargs, # Override model.generation_config with generation_kwargs to fix transformers#42762 - ) as unwrapped_model - ): - new_input_ids, new_attention_mask, new_labels = self.generate_on_policy_outputs( - unwrapped_model, inputs, self.generation_config - ) - inputs["input_ids"] = new_input_ids - inputs["attention_mask"] = new_attention_mask - inputs["labels"] = new_labels - elif self.seq_kd: - with ( - unwrap_model_for_generation( - self.teacher_model, - self.accelerator, - generation_kwargs=self.generation_kwargs, # Override model.generation_config with generation_kwargs to fix transformers#42762 - ) as unwrapped_model - ): - new_input_ids, new_attention_mask, new_labels = self.generate_on_policy_outputs( - unwrapped_model, inputs, self.generation_config - ) - inputs["input_ids"] = new_input_ids - inputs["attention_mask"] = new_attention_mask - inputs["labels"] = new_labels - def get_batch_samples(self, epoch_iterator, num_batches, *args, **kwargs): - batches, count = super().get_batch_samples(epoch_iterator, num_batches, *args, **kwargs) - if self.lmbda > 0 or self.seq_kd: - self._unsloth_gkd_pending_batches = (batches, args, kwargs) - return batches, count -class UnslothGKDTrainer(_UnslothGKDTrainer): - """ - Trainer for Generalized Knowledge Distillation (GKD) of language models. - - For details on GKD, see the paper: [On-Policy Distillation of Language Models: Learning from Self-Generated - Mistakes](https://huggingface.co/papers/2306.13649). - - Args: - model ([`~transformers.PreTrainedModel`] or `torch.nn.Module` or `str`, *optional*): - Model to be trained, or the string identifier of the model to be instantiated from a pretrained model. - teacher_model ([`~transformers.PreTrainedModel`] or `torch.nn.Module` or `str`, *optional*): - Teacher model for knowledge distillation, or the string identifier of the model to be instantiated from a - pretrained model. - args ([`experimental.gkd.GKDConfig`], *optional*): - Training arguments. - data_collator ([`~transformers.DataCollator`], *optional*): - Data collator to batch samples from the dataset. It defaults to a - [`experimental.utils.DataCollatorForChatML`] using the `processing_class`. - train_dataset ([`~datasets.Dataset`], *optional*): - Dataset for training. - eval_dataset ([`~datasets.Dataset`] or `dict` of [`~datasets.Dataset`], *optional*): - Dataset for evaluation. - processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.BaseImageProcessor`], [`~transformers.FeatureExtractionMixin`] or [`~transformers.ProcessorMixin`], *optional*): - Class to process the data. - compute_metrics (`Callable`, *optional*): - Function to compute metrics at evaluation. Must take in an [`~transformers.EvalPrediction`] and return a - dictionary string to float. - callbacks (`list` of [`~transformers.TrainerCallback`], *optional*): - Callbacks to use during training. - optimizers (`tuple` of `torch.optim.Optimizer` and `torch.optim.lr_scheduler.LambdaLR`, *optional*, defaults to `(None, None)`): - Tuple containing the optimizer and the learning rate scheduler to use for training. - preprocess_logits_for_metrics (`Callable`, *optional*): - Function to preprocess the logits before computing the metrics. Must take in the `logits` and `labels` and - return the logits to be used for metrics computation. - peft_config ([`~peft.PeftConfig`], *optional*): - PEFT configuration to use PEFT for training. If `None`, PEFT is not used. If provided, the `model` will be - wrapped with the specified PEFT adapter. - formatting_func (`Callable`, *optional*): - Function to format the dataset. Must take in an example and return an example. - - """ - def __init__( - self, - model = None, - teacher_model = None, - args = None, - data_collator = None, - train_dataset = None, - eval_dataset = None, - processing_class = None, - compute_metrics = None, - callbacks = None, - preprocess_logits_for_metrics = None, - peft_config = None, - formatting_func = None, - **kwargs - ): - if args is None: args = UnslothGKDConfig() - _unsloth_precision_requested = tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) - use_bf16 = getattr(args, 'bf16', False) - if type(use_bf16) is not bool: use_bf16 = False - use_fp16 = getattr(args, 'fp16', False) - if type(use_fp16) is not bool: use_fp16 = False - force_float32 = False - try: - from unsloth_zoo.device_type import device_is_bf16_supported as _bf16_supported - except Exception: - _bf16_supported = torch.cuda.is_bf16_supported - try: - from unsloth_zoo.device_type import DEVICE_TYPE as _unsloth_device_type - except Exception: - _unsloth_device_type = None - full_finetuning = getattr(model, '_unsloth_full_finetuning', None) - if full_finetuning is None: full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1' - model_forced_float32 = getattr(model, '_unsloth_forced_float32', None) - if model_forced_float32 is None: model_forced_float32 = os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1' - if model_forced_float32 and not (full_finetuning and _bf16_supported()): - print('Unsloth: Switching to float32 training since model cannot work with float16') - force_float32 = True - mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') - dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None) - if dtype is None: dtype = model.get_input_embeddings().weight.dtype - from unsloth_zoo.utils import _get_dtype - dtype = _get_dtype(dtype) - float16 = dtype == torch.float16 - bfloat16 = dtype == torch.bfloat16 - float32 = dtype == torch.float32 - user_float32 = bool(getattr(model, '_unsloth_user_float32', False)) - if full_finetuning: - if bfloat16 and use_fp16: use_fp16 = False - if float16 and use_bf16: use_bf16 = False - if not force_float32 and bfloat16 and use_fp16 and _unsloth_device_type == 'xpu' and not _bf16_supported(): - print('Unsloth: The model was loaded in bfloat16 on purpose, so training uses bfloat16 mixed precision instead of float16.') - use_fp16 = False - use_bf16 = True - args.fp16 = False - args.bf16 = True - if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`') - if not force_float32 and (bfloat16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`') - if force_float32: - # Forced float32 training - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32' and float32 and user_float32 and not _bf16_supported(): - print('Unsloth: Model is in float32 and this GPU has no bfloat16 support, so training stays in float32. Pass fp16 = True to force float16 mixed precision instead.') - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32': - # Mixed precision training. bf16 only if the GPU supports it; V100/T4 use fp16. - use_bf16_amp = (not float16) and (_bf16_supported() or (bfloat16 and _unsloth_device_type == 'xpu')) - args.fp16 = not use_bf16_amp - args.bf16 = use_bf16_amp - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16_amp else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16_amp else 'fp16' - # args.mixed_precision is a new argument which needs to be set now - elif mixed_precision_dtype == 'bfloat16': - # Both False since bfloat16 full finetuning doesn't do any autocasting. - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif use_bf16 or use_fp16: - # transformers <5 exported this itself from fp16/bf16; 5.x dropped the write, so an - # explicit flag left it unset and GRPO readers defaulted to 'fp16', wrapping a - # bfloat16 model in a float16 autocast. See unslothai/unsloth#4891. - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16 else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16 else 'fp16' - - if getattr(args, 'eval_dataset', None) is not None and getattr(args, 'eval_strategy', 'no') == 'no': - args.eval_strategy = 'steps' - if getattr(args, 'eval_steps', None) is None: args.eval_steps = 0.1 - ga_steps = getattr(args, 'gradient_accumulation_steps', None) - if ga_steps is not None and ga_steps > 1: - from transformers import __version__ as transformers_version - if Version(transformers_version) <= Version('4.45.2'): - print('**** Unsloth: Please use our fixed gradient_accumulation_steps by updating transformers, TRL and Unsloth!\n' - '`pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth transformers trl unsloth_zoo`') - if getattr(args, 'eval_strategy', 'no') != 'no': - eval_bsz = getattr(args, 'per_device_eval_batch_size', 8) - eval_generations = getattr(args, 'num_generations_eval', None) or getattr(args, 'num_generations', None) or 1 - if eval_bsz == 8 and args.per_device_train_batch_size < eval_bsz and (args.per_device_train_batch_size * args.world_size) % eval_generations == 0: args.per_device_eval_batch_size = args.per_device_train_batch_size - if getattr(args, 'eval_accumulation_steps', None) is None and ga_steps is not None: args.eval_accumulation_steps = ga_steps - fp16_full_eval = getattr(args, 'fp16_full_eval', False) - if type(fp16_full_eval) is not bool: fp16_full_eval = False - bf16_full_eval = getattr(args, 'bf16_full_eval', False) - if type(bf16_full_eval) is not bool: bf16_full_eval = False - if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True - if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False - if force_float32: - args.bf16_full_eval = False - args.fp16_full_eval = False - elif os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16': - args.bf16_full_eval = True - args.fp16_full_eval = False - elif not bf16_full_eval and not fp16_full_eval: - args.bf16_full_eval = args.bf16 - args.fp16_full_eval = args.fp16 - if tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) != _unsloth_precision_requested: - from unsloth.models.loader_utils import sync_deepspeed_mixed_precision - sync_deepspeed_mixed_precision(args) - if model is not None: - from unsloth.models.loader_utils import apply_deepspeed_zero3_model_setup - apply_deepspeed_zero3_model_setup(model) - if model is not None: - _warnings_issued = getattr(model, 'warnings_issued', None) - if _warnings_issued is None: - model.warnings_issued = {} - elif not isinstance(_warnings_issued, dict): - try: - model.warnings_issued = dict(_warnings_issued) - except Exception: - model.warnings_issued = {} - if 'max_seq_length' not in locals() and not hasattr(args, 'max_seq_length'): - pass - else: - model_max_seq_length = getattr(model, 'max_seq_length', None) - args_max_seq_length = getattr(args, 'max_seq_length', None) - if args_max_seq_length is None and model_max_seq_length is not None: - max_seq_length = model.max_seq_length - if hasattr(args, 'max_seq_length'): args.max_seq_length = max_seq_length - elif args_max_seq_length is not None and model_max_seq_length is not None: - if args_max_seq_length > model_max_seq_length: - print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but ' - 'the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.') - args.max_seq_length = model_max_seq_length - if model is not None and hasattr(model, 'for_training'): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right' - if 'processing_class' in locals(): - if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right' - if hasattr(processing_class, 'tokenizer') and hasattr(processing_class.tokenizer, 'padding_side'): processing_class.tokenizer.padding_side = 'right' - __tokenizer = processing_class if 'processing_class' in locals() else tokenizer - from unsloth_zoo.vision_utils import UnslothVisionDataCollator - if not isinstance(data_collator, UnslothVisionDataCollator): - if isinstance(data_collator, DataCollatorForSeq2Seq) and 'labels' not in _unsloth_dataset_column_names(train_dataset): - data_collator = TransformersDataCollatorForLanguageModeling( - __tokenizer, - mlm = False, - mlm_probability = 0.0, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling) and 'labels' in _unsloth_dataset_column_names(train_dataset): - data_collator = DataCollatorForSeq2Seq( - __tokenizer, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - else: - if hasattr(args, 'remove_unused_columns'): args.remove_unused_columns = False - if hasattr(args, 'dataset_text_field'): args.dataset_text_field = '' - if hasattr(args, 'dataset_kwargs'): args.dataset_kwargs = {'skip_prepare_dataset': True} - if not isinstance(data_collator, UnslothVisionDataCollator): - if not hasattr(__tokenizer, 'pad') and hasattr(__tokenizer, 'tokenizer'): - if isinstance(data_collator, DataCollatorForSeq2Seq): - data_collator = DataCollatorForSeq2Seq( - __tokenizer.tokenizer, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling): - data_collator = TransformersDataCollatorForLanguageModeling( - __tokenizer.tokenizer, - mlm = False, - mlm_probability = 0.0, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - other_metrics = [] - - from unsloth_zoo.logging_utils import PatchRLStatistics - PatchRLStatistics('gkd_trainer', other_metrics) - - # [TODO] Fix up DataParallel multiplying batch sizes - # [TODO] DDP works, but DP seems to not work? [TODO] - if getattr(args, "parallel_mode", None) == ParallelMode.NOT_DISTRIBUTED and args.n_gpu > 1: - if getattr(args, "_n_gpu", 1) != 1: - args._n_gpu = 1 - if "model" in locals(): - try: - from unsloth.models._gc_config_default import honour_gradient_checkpointing_off - honour_gradient_checkpointing_off(args, model) - except Exception: - pass - if "model" in locals() and hasattr(model, "for_training"): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if "model" in locals(): - try: - from unsloth.models._ddp_defaults import unsloth_default_ddp_find_unused_parameters - unsloth_default_ddp_find_unused_parameters(args, model) - except Exception: - pass - super().__init__( - model = model, - teacher_model = teacher_model, - args = args, - data_collator = data_collator, - train_dataset = train_dataset, - eval_dataset = eval_dataset, - processing_class = processing_class, - compute_metrics = compute_metrics, - callbacks = callbacks, - preprocess_logits_for_metrics = preprocess_logits_for_metrics, - peft_config = peft_config, - formatting_func = formatting_func,**kwargs) - # A tensor parallel model (tp_native.tp_trainer_setup) stays on its training path: for_inference would - # route the decoder layers onto decode paths outside the tensor parallel regions. - if "model" in locals() and hasattr(model, "for_inference") and not getattr(self, "_unsloth_tensor_parallel", False): - model.for_inference() - if hasattr(self, 'neftune_hook_handle'): - self.neftune_hook_handle.remove() - if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle - if getattr(args, 'neftune_noise_alpha', None) is not None: - model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha - pass - if hasattr(self, 'accelerator'): - scaler = self.accelerator.scaler - current_model = model - while hasattr(current_model, 'model'): - current_model.accelerator_scaler = scaler - current_model = current_model.model - current_model.accelerator_scaler = scaler - pass - if hasattr(self, 'train'): - self.train = MethodType(prepare_for_training_mode(self.__class__.train), self) - pass - if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'): - _vllm_tok = self.llm.get_tokenizer() - _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None) - if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None: - _vllm_tok.chat_template = _pc.chat_template - pass - if getattr(self, 'aux_loss_enabled', False) and hasattr(getattr(self, 'model', None), 'config'): - _text_config = self.model.config - if hasattr(_text_config, 'get_text_config'): _text_config = _text_config.get_text_config() - _n_experts = [getattr(_text_config, _k) for _k in ('num_local_experts', 'num_experts', 'n_routed_experts', 'moe_num_experts') if isinstance(getattr(_text_config, _k, None), int)] - if _n_experts and all(_n <= 1 for _n in _n_experts): - self.aux_loss_enabled = False - _text_config.output_router_logits = False - pass - if getattr(self, 'liger_jsd_loss', None) is not None: - self.liger_jsd_loss.weight_hard_loss = 0.0 - self.liger_jsd_loss.weight_soft_loss = 1.0 - self.liger_jsd_loss.temperature = 1.0 - -pass - - -if hasattr(logger, "addFilter"): - import logging - class HideLoggingMessage(logging.Filter): - def __init__(self, text): self.text = text - def filter(self, x): return not (self.text in x.getMessage()) - pass - logger.addFilter(HideLoggingMessage("`use_cache=True`")) - diff --git a/unsloth_compiled_cache/UnslothGOLDTrainer.py b/unsloth_compiled_cache/UnslothGOLDTrainer.py deleted file mode 100644 index 4e04f9f2be38585a21c3e2a6cd92de7ab71943d1..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/UnslothGOLDTrainer.py +++ /dev/null @@ -1,3423 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from unsloth_zoo.temporary_patches.common import torch_compile -from unsloth_zoo.temporary_patches.common import _maybe_compile -import functools -import transformers -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from trl.experimental.gold.gold_trainer import (Any, AutoConfig, AutoProcessor, AutoTokenizer, BaseImageProcessor, Callable, DataCollator, DataCollatorForChatML, DataCollatorForVisionLanguageChatML, DataLoader, Dataset, DistributedType, EvalPrediction, F, FeatureExtractionMixin, GOLDConfig, GOLDTrainer, GenerationConfig, IterableDataset, Optional, PeftConfig, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, RepeatSampler, SFTTrainer, Table, TrainerCallback, ULDLoss, VLLMGeneration, Version, broadcast_object_list, build_teacher_inputs_from_texts, create_model_from_path, defaultdict, deque, disable_dropout_in_model, dist, empty_cache, encode_with_byte_offsets, gather_object, get_config_model_id, global_then_local_main_first, identity, is_conversational, is_datasets_available, is_rich_available, is_vllm_available, maybe_convert_to_chatml, nn, pack_dataset, pad, pad_byte_offsets, partial, piece_byte_len, prepare_deepspeed, prepare_multimodal_messages, print_prompt_completions_sample_uld, profiling_decorator, random, seed_worker, split_tensor_dict, takewhile, textwrap, torch, transformers, unwrap_model_for_generation, warnings, AutoConfig, AutoProcessor, AutoTokenizer, BaseImageProcessor, Callable, DataCollator, DataCollatorForChatML, DataCollatorForVisionLanguageChatML, Dataset, EvalPrediction, F, FeatureExtractionMixin, GOLDConfig, GenerationConfig, Optional, PeftConfig, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, SFTTrainer, TrainerCallback, ULDLoss, VLLMGeneration, create_model_from_path, defaultdict, deque, disable_dropout_in_model, dist, get_config_model_id, identity, is_vllm_available, nn, pad, prepare_deepspeed, torch, warnings) - - -import os -import math -import logging -from typing import * -from dataclasses import dataclass, field -from packaging.version import Version -import torch -import numpy as np -from contextlib import nullcontext -from torch.nn import functional as F -import inspect -from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling -from transformers.training_args import ParallelMode -from unsloth_zoo.device_type import DEVICE_TYPE, DEVICE_TYPE_TORCH, device_synchronize - -# Wrap trainer with padding to right and enable training mode -import functools -from types import MethodType -try: - from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers -except: - def reset_unsloth_gradient_checkpointing_buffers(): pass -# Canonical reset lives in unsloth.models._utils so the SFT auto-packing wrapper and the plain -# Trainer loop can import the same helper; fall back to a no-op only if it can't be imported. -try: - from unsloth.models._utils import _unsloth_reset_stray_compile_cache -except Exception: - def _unsloth_reset_stray_compile_cache(self): pass -try: - from unsloth.models._utils import _unsloth_dataset_column_names -except Exception: - def _unsloth_dataset_column_names(dataset): return dataset.column_names -# Drops/renames config arguments the installed TRL no longer accepts, so a -# script pinned to an older TRL keeps working after an upgrade. Falls back to -# the historical raw passthrough so this can never break trainer construction. -try: - from unsloth.models.rl_config_compat import filter_config_init_kwargs as _unsloth_filter_config_init_kwargs - # A cache file generated here can be imported by an older Unsloth whose filter - # predates `mirrored_from`, so drop the argument rather than raise TypeError. - if "mirrored_from" not in inspect.signature(_unsloth_filter_config_init_kwargs).parameters: - _unsloth_filter_config_init_kwargs_old = _unsloth_filter_config_init_kwargs - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): - return _unsloth_filter_config_init_kwargs_old(config_class, kwargs) -except Exception: - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): return kwargs -def prepare_for_training_mode(f): - @functools.wraps(f) - def wrapper(self, *args, **kwargs): - # Drop any torch.compile graph cache poisoned by a stray pre-train forward. - try: - _unsloth_reset_stray_compile_cache(self) - except Exception: - pass - # Finish the previous W&B run if this is a subsequent train() call. - # We do this at the START of train() (not the end) so that - # evaluate() / log() still work after train() completes. - # HF's WandbCallback.setup() will call wandb.init() for the new run. - # See: https://github.com/unslothai/unsloth/issues/3954 - if getattr(self, '_unsloth_training_completed', False): - try: - import wandb - if wandb.run is not None: - wandb.finish() - # Reset HF's WandbCallback so it calls wandb.init() for the new run - for cb in self.callback_handler.callbacks: - if type(cb).__name__ == 'WandbCallback': - cb._initialized = False - break - except: - pass - # Enable training mode - _was_training = None - # Restore the GC mode the model was configured with at setup; fall back to - # the training args only when it wasn't recorded (issue #4735). Use hasattr, - # not a None sentinel, so a deliberately-recorded None is restored verbatim. - _model = getattr(self, 'model', None) - if hasattr(_model, '_unsloth_gradient_checkpointing'): - use_gc = _model._unsloth_gradient_checkpointing - else: - use_gc = getattr(self.args, 'gradient_checkpointing', True) - if hasattr(self, 'model') and hasattr(self.model, "training"): - _was_training = self.model.training - if hasattr(self, 'model') and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - output = f(self, *args, **kwargs) - # Restore previous mode when possible - if hasattr(self, 'model') and hasattr(self.model, "for_inference"): - if _was_training is False: - self.model.for_inference() - elif _was_training is True and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - # Reset gradient checkpointing buffers to free memory while staying ready for next run - try: - reset_unsloth_gradient_checkpointing_buffers() - except: - pass - # Mark that training completed so the next train() call can - # finish this W&B run before starting a new one - self._unsloth_training_completed = True - return output - return wrapper -pass - -torch_compile_options = { - "epilogue_fusion" : True, - "max_autotune" : False, - "shape_padding" : True, - "trace.enabled" : False, - "triton.cudagraphs" : False, -} - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_hidden_states_selective_log_softmax( - hidden_states: torch.Tensor, - lm_head: torch.Tensor, - index: torch.Tensor, - chunks: int = 4, - logit_scale_multiply: float = 0.0, - logit_scale_divide: float = 0.0, - logit_softcapping: float = 0.0, - temperature: float = 1.0, - # Rows per chunk cap. Read HERE, in the default, not in the body: the body - # is traced with fullgraph = True, and `os.environ` is an unsupported op - # there on torch 2.4 -- Dynamo raises - # torch._dynamo.exc.Unsupported: const method call bytes.decode - # from os._Environ.__getitem__, which the eager fallback does not catch - # (it only catches recompile-limit and disabled-hook breaks), so the very - # first call would die even with the variable unset. A default is evaluated - # once when the def runs, which is import time, outside any traced region. - # It is also a plain int argument, so Dynamo guards on it instead of - # constant-folding an unguarded read (2.7+ never notice a later change). - # A non-numeric value is ignored rather than raised on, so a typo cannot - # break the import. 0 keeps the previous chunk boundaries exactly. - max_rows_per_chunk: int = ( - int(os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip()) - if os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip().isdigit() - else 0 - ), - lm_head_bias: Optional[torch.Tensor] = None, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, - # Read in the default for the same Dynamo reason as max_rows_per_chunk. Pairs with vLLM's - # head_dtype=float32 that load_vllm requests under the same flag. - fp32_lm_head: bool = ( - os.environ.get("UNSLOTH_VLLM_FP32_LM_HEAD", "0").strip().lower() in ("1", "true", "on", "yes") - ), - # Older Unsloth scores old / ref log-probs here and the policy with the fused kernel: one scorer keeps ratio at 1. - fused_linear: bool = ( - os.environ.get("UNSLOTH_FUSED_LINEAR", "1") != "0" - and os.environ.get("UNSLOTH_GRPO_CHUNKED_LOGPROBS", "1") != "0" - ), -) -> torch.Tensor: - # All Unsloth Zoo code licensed under AGPL3 - from unsloth_zoo.fused_losses import chunked_logprob as _chunked_logprob - if (fused_linear and _chunked_logprob._FUSED_LINEAR and not fp32_lm_head - and support_top_k == 0 and support_top_p == 1.0 and support_min_p == 0.0 - and support_excluded_ids is None): - return _chunked_logprob.unsloth_chunked_logprobs( - hidden_states, lm_head, index, - temperature = temperature, - logit_scale_multiply = logit_scale_multiply, - logit_scale_divide = logit_scale_divide, - logit_softcapping = logit_softcapping, - lm_head_bias = lm_head_bias, - max_rows_per_chunk = max_rows_per_chunk if max_rows_per_chunk > 0 else None, - ).reshape((hidden_states.shape[0], hidden_states.shape[1])) - # Reshape on this tensor's own last dim: a no-op, so a wrong-width caller - # cannot have its row count silently rewritten and instead fails at the - # matmul below, which prints both operands. Do not swap in a bare - # torch._check: it reports only "Expected cond to be True", naming neither - # operand, and Dynamo rejects a message-carrying one. Callers dispatch on - # the width first -- see `compute_logprobs_chunk`, the packed path and - # `_pg_grad_forward`. - flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) - flat_index = index.reshape(-1) - - # Each chunk materialises rows x vocab logits and then a float32 copy of - # them, all on the device holding the output head. With a large vocabulary - # and a fixed chunk count that grows with the batch, so the peak scales with - # the batch rather than staying bounded. max_rows_per_chunk caps the rows - # per chunk instead, which is pure loop splitting: more, smaller chunks, - # same concatenated result. 0 (the default) keeps the previous chunk - # boundaries exactly. - if max_rows_per_chunk > 0: - n_rows = flat_hidden_states.shape[0] - chunks = max(chunks, -(-n_rows // max_rows_per_chunk)) - chunks = min(chunks, max(n_rows, 1)) - - # A one-row chunk's head gradient is a K=1 matmul Inductor mis-lowers under dynamic shapes (illegal - # memory access or NaN head gradients, torch 2.13), so keep chunks at 2+ rows when the head trains. - # Exactly `chunks` row ranges, so the unrolled loop length never depends on T: torch.chunk - # returns ceil(T / ceil(T / chunks)) pieces and recompiled per new count; tensor_split guards - # on T % chunks. At most rows // 2 chunks when the head trains keeps every piece at 2+ rows. - lm_head_grad = torch.is_grad_enabled() and lm_head.requires_grad - if lm_head_grad and flat_hidden_states.shape[0] < 2 * chunks: - chunks = max(flat_hidden_states.shape[0] // 2, 1) - - n_rows = flat_hidden_states.shape[0] - chunks = max(1, min(chunks, n_rows)) - bounds = [(i * n_rows) // chunks for i in range(chunks + 1)] - chunked_hidden_states = [flat_hidden_states[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - - # Cast once, outside the loop, so every chunk shares one fp32 head copy (also for autograd). - if fp32_lm_head: - lm_head = lm_head.to(torch.float32) - if lm_head_bias is not None: - lm_head_bias = lm_head_bias.to(torch.float32) - - all_per_token_logps = [] - - for chunk_hidden_states, chunk_index in zip(chunked_hidden_states, chunked_index): - # When the model is dispatched over several devices, the output head can - # sit on a different one from the hidden states, because accelerate - # places the tail of the model on the last device it fills. Co-locate on - # the head's device before the matmul, otherwise this raises - # Unhandled FakeTensor Device Propagation for aten.mm.default, - # found two different devices cuda:0, cuda:1 - # On a single device every .to() here is a no-op and the result is - # bit-identical to before. - chunk_hidden_states = chunk_hidden_states.to(device = lm_head.device, dtype = lm_head.dtype) - chunk_index = chunk_index.to(lm_head.device) - if fp32_lm_head: - # The GRPO loss runs this under bf16 / fp16 autocast, which would downcast the matmul. - with torch.autocast(device_type = chunk_hidden_states.device.type, enabled = False): - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - elif lm_head_bias is None: - chunk_logits = chunk_hidden_states @ lm_head.t() - else: - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - - if logit_scale_multiply != 0.0: - chunk_logits = chunk_logits * logit_scale_multiply - if logit_scale_divide != 0.0: - chunk_logits = chunk_logits / logit_scale_divide - if logit_softcapping != 0.0: - chunk_logits = logit_softcapping * torch.tanh(chunk_logits / logit_softcapping) - - chunk_logits = chunk_logits.to(torch.float32) - - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - per_token_logps = _truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - ) - else: - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - # Return to the caller's device so the concatenation below and every - # downstream consumer see the device they started on. - all_per_token_logps.append(per_token_logps.to(hidden_states.device)) - - all_per_token_logps = torch.concat(all_per_token_logps) - - all_per_token_logps = all_per_token_logps.reshape((hidden_states.shape[0], hidden_states.shape[1])) - return all_per_token_logps - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_selective_log_softmax( - logits, - index, - temperature: float = 1.0, - chunks: int = 4, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, -): - # Fixed `chunks` row ranges; see chunked_hidden_states_selective_log_softmax. - flat_logits = logits.reshape(-1, logits.shape[-1]) - flat_index = index.reshape(-1) - # Fewer rows than chunks: one row per range, as torch.chunk did, never empty ranges. - chunks = max(1, min(chunks, flat_logits.shape[0])) - bounds = [(i * flat_logits.shape[0]) // chunks for i in range(chunks + 1)] - chunked_logits = [flat_logits[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - all_per_token_logps = [] - # Per-chunk selective_log_softmax. - for chunk_logits, chunk_index in zip(chunked_logits, chunked_index): - chunk_logits = chunk_logits.to(torch.float32) - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - all_per_token_logps.append(_truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - )) - continue - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - all_per_token_logps.append(per_token_logps) - pass - all_per_token_logps = torch.concat(all_per_token_logps) - all_per_token_logps = all_per_token_logps.reshape((logits.shape[0], logits.shape[1])) - return all_per_token_logps - -def calculate_pad_tokens_in_prompt( - input_ids: torch.Tensor, - logits_to_keep: int, - pad_token_id: int -) -> torch.Tensor: - """Count left-padded tokens per sequence, e.g. [pad, pad, pad, cat] -> 3.""" - if logits_to_keep >= input_ids.shape[1]: - raise ValueError("logits_to_keep must be smaller than the sequence length.") - - prompt_section = input_ids[:, :-logits_to_keep] - - padding_mask = (prompt_section == pad_token_id) - - pad_token_counts = padding_mask.sum(dim=1) - - return pad_token_counts - -def create_completion_attention_mask( - completion_input_ids: torch.Tensor, - left_pad_tokens_per_prompt: torch.Tensor, - max_left_pad: int, - pad_token_id: int -) -> torch.Tensor: - """Build a completion mask that zeros leading prompt and trailing pad tokens. - - For [p,p,p,c,c,c,pad,pad,pad] (p=sliced prompt, c=completion, pad=padding) - this returns [0,0,0,1,1,1,0,0,0]. - """ - batch_size, completion_len = completion_input_ids.shape - device = completion_input_ids.device - - num_tokens_to_mask = max_left_pad - left_pad_tokens_per_prompt - - indices = torch.arange(completion_len, device=device).unsqueeze(0) - shift_mask = indices >= num_tokens_to_mask.unsqueeze(1) - - non_padding_mask = (completion_input_ids != pad_token_id) - - final_mask = shift_mask & non_padding_mask - - return final_mask - -def left_pack_padding(tensor: torch.Tensor, pad_id: int) -> torch.Tensor: - """Move all padding tokens in each sequence to the right.""" - mask = (tensor != pad_id) - # stable=True since the binary mask is unordered. - sorted_indices = torch.argsort(mask, dim=1, descending=True, stable=True) - packed_tensor = torch.gather(tensor, 1, sorted_indices) - return packed_tensor - -def align_logprobs_with_mask( - logprob_tensor: torch.Tensor, - attention_mask: torch.Tensor, - pad_value: float = 0.0 -) -> torch.Tensor: - """Align a log probability tensor with a given attention mask.""" - - device = logprob_tensor.device - batch_size, logprob_seq_len = logprob_tensor.shape - mask_seq_len = attention_mask.shape[1] - - padded_logprobs = torch.full( - attention_mask.shape, - fill_value=pad_value, - dtype=logprob_tensor.dtype, - device=device - ) - - left_pad_counts = torch.argmax(attention_mask, dim=1) - - cols = torch.arange(logprob_seq_len, device=device) - dest_indices = left_pad_counts.unsqueeze(1) + cols - - # Destination row indices, shape [batch_size, logprob_seq_len]. - row_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand_as(dest_indices) - - # Keep only in-bounds destinations, then scatter via advanced indexing. - valid_mask = dest_indices < mask_seq_len - valid_rows = row_indices[valid_mask] - valid_cols = dest_indices[valid_mask] - valid_vals = logprob_tensor[valid_mask] - padded_logprobs[valid_rows, valid_cols] = valid_vals - - return padded_logprobs - -def align_completion_tool_mask( - tool_mask: torch.Tensor, - completion_mask: torch.Tensor, -) -> torch.Tensor: - """Align a raw completion-length tool/env mask with Unsloth's repacked loss mask.""" - if tool_mask is None: - return completion_mask - if tool_mask.shape[0] != completion_mask.shape[0]: - raise ValueError("tool_mask batch size must match completion_mask batch size.") - - tool_mask = tool_mask.to(device=completion_mask.device) - if tool_mask.shape == completion_mask.shape: - aligned_tool_mask = tool_mask - else: - aligned_tool_mask = align_logprobs_with_mask( - tool_mask, - completion_mask, - pad_value=0, - ) - return completion_mask * aligned_tool_mask.to(dtype=completion_mask.dtype) - -def autotune_batch_and_chunks( - total_input_rows, - seq_len, - hidden_size, - vocab_size, - dtype_bytes=16, - multiplier=None -): - if multiplier is None: - final_m = max(4, seq_len // 4096) - else: - final_m = multiplier - - if torch.cuda.is_available(): - free_bytes, _ = torch.cuda.mem_get_info() - limit_gb = (free_bytes / (1024**3))*.80 - elif hasattr(torch, "xpu") and torch.xpu.is_available(): - # XPU: estimate free memory as total - reserved. - total_mem = torch.xpu.get_device_properties(0).total_memory - reserved_mem = torch.xpu.memory_reserved() - free_bytes = total_mem - reserved_mem - limit_gb = (free_bytes / (1024**3)) * 0.80 - else: - # Fallback: assume 8GB available. - limit_gb = 8.0 - - bytes_to_gb = 1024**3 - - b_vals = torch.arange(total_input_rows, 0, -1, device='cpu', dtype=torch.float32) - - hidden_gb = (b_vals * seq_len * hidden_size * dtype_bytes) / bytes_to_gb - - base_logits = ((b_vals/total_input_rows) * b_vals * seq_len * vocab_size * dtype_bytes) / bytes_to_gb - logits_gb = base_logits / final_m - - total_mem_gb = hidden_gb + logits_gb - - valid_mask = total_mem_gb <= limit_gb - valid_indices = torch.nonzero(valid_mask, as_tuple=False) - - if valid_indices.shape[0] == 0: - #This means your GPU will OOM - # Capped at the row count: unsloth's no-grad pass divides rows by this without max(1, ...). - return max(1, min(4, total_input_rows)), final_m - - best_idx = valid_indices[0].item() - final_b = int(b_vals[best_idx].item()) - - return final_b, final_m - -def sanitize_logprob(logprob): - """Local port of trl.scripts.vllm_serve.sanitize_logprob. - Filters NaN logprobs from vLLM outputs.""" - value = logprob.logprob - if math.isnan(value): - logging.getLogger(__name__).warning( - f"Generated NaN logprob, token logprob '{logprob}' will be ignored" - ) - return None - return value -@dataclass -class UnslothGOLDConfig(GOLDConfig): - """ - - Configuration class for [`experimental.gold.GOLDTrainer`]. - - This class includes only the parameters that are specific to GOLD training. For a full list of training arguments, - please refer to the [`~transformers.TrainingArguments`] and [`SFTConfig`] documentation. - - Args: - > Parameters that control generation and the training loop - - temperature (`float`, *optional*, defaults to `0.9`): - Temperature for sampling. The higher the temperature, the more random the completions. - top_p (`float`, *optional*, defaults to `0.95`): - If set to float < 1, only the smallest set of most probable tokens with probabilities that add up to - `top_p` or higher are kept for generation. - top_k (`int`, *optional*, defaults to `0`): - Number of highest probability vocabulary tokens to keep for top-k-filtering. If `0`, top-k-filtering is - disabled and all tokens are considered. - lmbda (`float`, *optional*, defaults to `0.5`): - Lambda parameter that controls the student data fraction (i.e., the proportion of on-policy - student-generated outputs). - beta (`float`, *optional*, defaults to `0.5`): - Interpolation coefficient between `0.0` and `1.0` of the Generalized Jensen-Shannon Divergence loss. When - beta is `0.0`, the loss is the KL divergence. When beta is `1.0`, the loss is the Inverse KL Divergence. - max_completion_length (`int`, *optional*, defaults to `128`): - Maximum number of tokens to generate per completion. - teacher_model_name_or_path (`str`, *optional*): - Model name or path of the teacher model. If `None`, the teacher model will be the same as the model being - trained. - teacher_model_revision (`str` or `None`, *optional*, defaults to `None`): - Model revision of the teacher model (e.g., branch name, tag, or commit hash). If `None`, the default - revision is used. - teacher_model_init_kwargs (`dict[str, Any]`, *optional*): - Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the teacher model - from a string. The `revision` value is also used when loading the teacher processing class. - teacher_tokenizer_name_or_path (`str`, *optional*): - Tokenizer name or path for the teacher model. If None when using ULD loss, will use the same tokenizer as - the student model (not recommended for cross-tokenizer distillation). - disable_dropout (`bool`, *optional*, defaults to `True`): - Whether to disable dropout in the model. - seq_kd (`bool`, *optional*, defaults to `False`): - Seq_kd parameter that controls whether to perform Sequence-Level KD (can be viewed as supervised FT on - teacher-generated output). - num_generations (`int`, *optional*, defaults to `1`): - Number of generations per prompt. Each prompt is repeated this many times in the generation batch. - generation_batch_size (`int` or `None`, *optional*, defaults to `None`): - Number of unique prompts per worker per optimizer step. If `None`, it is computed from - `(per_device_train_batch_size * gradient_accumulation_steps) // num_generations`. - > Parameters that control the ULD loss - - use_uld_loss (`bool`, *optional*, defaults to `False`): - Whether to use Universal Logit Distillation (ULD) loss instead of Generalized Jensen-Shannon Divergence - loss. - use_extended_uld (`bool`, *optional*, defaults to `True`): - Whether to enable extended ULD alignment that uses tokenizers to align and merge token probabilities across - student and teacher tokenizations. When `True`, the trainer will compute token mappings and merge - probabilities for split tokens; when `False`, ULD will use simple positional truncation like in the - original ULD paper. Set to `False` for SentencePiece or other non-ByteLevel tokenizers. - uld_token_merge_strategy (`str`, *optional*, defaults to `"observed"`): - Strategy used to align answer logits and merge token probabilities in the ULD loss. With `"observed"`, the - answer logits are sliced at the answer positions and split tokens (when `use_extended_uld=True`) are merged - by multiplying the marginal distribution at the first position by the scalar conditional probabilities of - the actual later tokens. With `"bayesian"`, the answer-logit slice is shifted one position earlier so that - `probs[k]` predicts `token_ids[k]` (chain rule), and split tokens are merged using the last position's full - distribution, conditioned on the actual prefix tokens, multiplied by the scalar probabilities of the - earlier tokens. The logit-alignment shift applies whether or not `use_extended_uld` is enabled. - uld_use_hybrid_loss (`bool`, *optional*, defaults to `False`): - Whether to use a hybrid loss that combines ULD loss and JSD loss. When `True`, the final loss is a - combination of JSD for known token mappings and ULD for unknown token mappings. - uld_hybrid_matched_weight (`float` or `None`, *optional*): - Weight for the matched token loss component when using hybrid ULD + JSD loss. This weight scales the JSD - loss computed over tokens that have a direct mapping between student and teacher tokenizations. If `None`, - uses adaptive weighting based on vocabulary overlap. Must be set together with - `uld_hybrid_unmatched_weight` (both `None` or both `float`). - uld_hybrid_unmatched_weight (`float` or `None`, *optional*): - Weight for the unmatched token loss component when using hybrid ULD + JSD loss. This weight scales the ULD - loss computed over tokens that do not have a direct mapping between student and teacher tokenizations. If - `None`, uses adaptive weighting based on vocabulary overlap. Must be set together with - `uld_hybrid_matched_weight` (both `None` or both `float`). - uld_crossentropy_weight (`float`, *optional*, defaults to `0.0`): - Weight for the cross-entropy loss component in ULD loss. If 0, only ULD distillation loss is used. - uld_distillation_weight (`float`, *optional*, defaults to `1.0`): - Weight for the distillation loss component in ULD loss. - uld_student_temperature (`float`, *optional*, defaults to `1.0`): - Temperature for student logits in ULD loss computation. - uld_teacher_temperature (`float`, *optional*, defaults to `1.0`): - Temperature for teacher logits in ULD loss computation. - uld_skip_student_eos (`bool`, *optional*, defaults to `True`): - Whether to skip EOS token for student in ULD loss computation. - uld_skip_teacher_eos (`bool`, *optional*, defaults to `True`): - Whether to skip EOS token for teacher in ULD loss computation. - > Parameters that control vLLM integration - - use_vllm (`bool`, *optional*, defaults to `False`): - Whether to use vLLM for generating completions from the student model. Requires `vllm` to be installed. - vllm_mode (`str`, *optional*, defaults to `"colocate"`): - Mode for student vLLM integration. Either `"server"` (connect to a running vLLM server) or `"colocate"` - (run vLLM in the same process). - vllm_server_host (`str`, *optional*, defaults to `"0.0.0.0"`): - Host of the vLLM server for the student model (if `vllm_mode="server"`). - vllm_server_port (`int`, *optional*, defaults to `8001`): - Port of the vLLM server for the student model (if `vllm_mode="server"`). - vllm_server_timeout (`float`, *optional*, defaults to `240.0`): - Timeout for connecting to the student vLLM server (if `vllm_mode="server"`). - vllm_gpu_memory_utilization (`float`, *optional*, defaults to `0.9`): - GPU memory utilization for the colocated student vLLM engine (if `vllm_mode="colocate"`). It is recommended - to set this to a low value if the student and teacher models share the same GPU. - vllm_tensor_parallel_size (`int`, *optional*, defaults to `1`): - Tensor parallel size for the colocated student vLLM engine (if `vllm_mode="colocate"`). - vllm_structured_outputs_regex (`str`, *optional*): - Regex for vLLM structured outputs for the student model. - vllm_server_base_url (`str`, *optional*): - Base URL for the vLLM server (e.g., `"http://localhost:8001"`). If provided, `vllm_server_host` and - `vllm_server_port` are ignored. - vllm_group_port (`int`, *optional*, defaults to `51216`): - Port for the vLLM weight-update group (NCCL communicator). Unless the port is occupied, there is no need to - change it. - vllm_max_model_length (`int`, *optional*): - Maximum model sequence length for the colocated vLLM engine when `vllm_mode="colocate"`. Defaults to the - model's maximum context length. - vllm_model_impl (`str`, *optional*, defaults to `"vllm"`): - Model implementation backend to use in vLLM. Use `"vllm"` (default) or `"transformers"`. - vllm_sync_frequency (`int`, *optional*, defaults to `1`): - Frequency (in training steps) to synchronize student model weights to vLLM engine. Set to 1 to sync after - every step. - vllm_enable_sleep_mode (`bool`, *optional*, defaults to `False`): - Enable vLLM sleep mode to offload student weights/cache during the optimizer step. Keeps GPU memory usage - low, but waking the engine adds host–device transfer latency. - - > Parameters that control logging - - log_completions (`bool`, *optional*, defaults to `False`): - Whether to log a sample of (prompt, completion) pairs every `logging_steps` steps. If `rich` is installed, - it prints the sample. If `wandb` logging is enabled, it logs it to `wandb`. - log_completions_steps (`int`, *optional*, defaults to `100`): - Number of steps between logging (prompt, completion) pairs. Only used if `log_completions` is set to - `True`. - num_completions_to_print (`int` or `None`, *optional*): - Number of completions to print with `rich`. If `None`, all completions are logged. - wandb_log_unique_prompts (`bool`, *optional*, defaults to `True`): - Whether to log the unique prompts to wandb. This will create a new run for each unique prompt. - callbacks (`list[str]`, *optional*, defaults to `[]`): - The callbacks to run during training. - - > [!NOTE] - > These parameters have default values different from [`~transformers.TrainingArguments`]: - > - `learning_rate`: Defaults to `1e-7` instead of `5e-5`. - - """ - _unsloth_accepts_adamw_flash = True - vllm_sampling_params: Optional[Any] = field( - default = None, - metadata = {'help': 'vLLM SamplingParams'}, - ) - unsloth_num_chunks : Optional[int] = field( - default = -1, - metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'}, - ) - unsloth_logit_chunk_multiplier : Optional[int] = field( - default = None, - metadata = {'help': 'Multiplier for chunked logit computations.'}, - ) - unsloth_grpo_mini_batch : Optional[int] = field( - default = None, - metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'}, - ) - max_seq_length : Optional[int] = field( - default = None, - metadata = {'help': 'Maximum sequence length to truncate to.'}, - ) - def __init__( - self, - output_dir = None, - per_device_train_batch_size = 4, - num_train_epochs = 3.0, - max_steps = -1, - learning_rate = 5e-05, - lr_scheduler_type = 'linear', - lr_scheduler_kwargs = None, - warmup_steps = 0.1, - optim = 'adamw_8bit', - optim_args = None, - weight_decay = 0.001, - adam_beta1 = 0.9, - adam_beta2 = 0.999, - adam_epsilon = 1e-08, - optim_target_modules = None, - gradient_accumulation_steps = 2, - average_tokens_across_devices = True, - max_grad_norm = 1.0, - label_smoothing_factor = 0.0, - bf16 = False, - fp16 = False, - bf16_full_eval = False, - fp16_full_eval = False, - tf32 = None, - gradient_checkpointing = True, - gradient_checkpointing_kwargs = None, - torch_compile = False, - torch_compile_backend = None, - torch_compile_mode = None, - use_liger_kernel = False, - liger_kernel_config = None, - use_cache = False, - neftune_noise_alpha = None, - torch_empty_cache_steps = 250, - auto_find_batch_size = False, - logging_strategy = 'steps', - logging_steps = 1, - logging_first_step = False, - log_on_each_node = True, - logging_nan_inf_filter = False, - include_num_input_tokens_seen = False, - log_level = 'passive', - log_level_replica = 'warning', - disable_tqdm = None, - report_to = 'none', - run_name = None, - project = 'huggingface', - trackio_space_id = 'trackio', - eval_strategy = 'no', - eval_steps = None, - eval_delay = 0, - per_device_eval_batch_size = 4, - prediction_loss_only = False, - eval_on_start = False, - eval_do_concat_batches = True, - eval_use_gather_object = False, - eval_accumulation_steps = 2, - batch_eval_metrics = False, - save_only_model = False, - save_strategy = 'steps', - save_steps = 500, - save_on_each_node = False, - save_total_limit = None, - enable_jit_checkpoint = False, - push_to_hub = False, - hub_token = None, - hub_private_repo = None, - hub_model_id = None, - hub_strategy = 'every_save', - hub_always_push = False, - hub_revision = None, - load_best_model_at_end = False, - metric_for_best_model = None, - greater_is_better = None, - ignore_data_skip = False, - restore_callback_states_from_checkpoint = False, - full_determinism = False, - seed = 3407, - data_seed = None, - use_cpu = False, - accelerator_config = None, - parallelism_config = None, - dataloader_drop_last = False, - dataloader_num_workers = 0, - dataloader_pin_memory = True, - dataloader_persistent_workers = False, - dataloader_prefetch_factor = None, - remove_unused_columns = True, - label_names = None, - train_sampling_strategy = 'random', - length_column_name = 'length', - ddp_find_unused_parameters = None, - ddp_bucket_cap_mb = None, - ddp_broadcast_buffers = None, - ddp_backend = None, - ddp_timeout = 1800, - fsdp = None, - fsdp_config = None, - deepspeed = None, - debug = '', - skip_memory_metrics = True, - do_train = False, - do_eval = False, - do_predict = False, - resume_from_checkpoint = None, - warmup_ratio = None, - logging_dir = None, - local_rank = -1, - model_init_kwargs = None, - trust_remote_code = False, - router_aux_loss_coef = None, - chat_template_path = None, - dataset_text_field = 'text', - dataset_kwargs = None, - dataset_num_proc = None, - eos_token = None, - max_length = 1024, - truncation_mode = 'keep_start', - shuffle_dataset = False, - packing = False, - packing_strategy = 'bfd', - padding_free = None, - pad_to_multiple_of = None, - eval_packing = None, - completion_only_loss = None, - assistant_only_loss = False, - loss_type = 'nll', - activation_offloading = False, - pad_token = None, - temperature = 0.9, - top_p = 0.95, - top_k = 0, - lmbda = 0.5, - beta = 0.5, - max_completion_length = 128, - teacher_model_name_or_path = None, - teacher_model_revision = None, - teacher_model_init_kwargs = None, - teacher_tokenizer_name_or_path = None, - disable_dropout = True, - seq_kd = False, - num_generations = 1, - generation_batch_size = None, - use_uld_loss = False, - uld_token_merge_strategy = 'observed', - use_extended_uld = True, - uld_use_hybrid_loss = False, - uld_hybrid_matched_weight = None, - uld_hybrid_unmatched_weight = None, - uld_crossentropy_weight = 0.0, - uld_distillation_weight = 1.0, - uld_student_temperature = 1.0, - uld_teacher_temperature = 1.0, - uld_skip_student_eos = True, - uld_skip_teacher_eos = True, - use_vllm = False, - vllm_mode = 'colocate', - vllm_server_base_url = None, - vllm_server_host = '0.0.0.0', - vllm_server_port = 8001, - vllm_server_timeout = 240.0, - vllm_group_port = 51216, - vllm_gpu_memory_utilization = 0.9, - vllm_tensor_parallel_size = 1, - vllm_max_model_length = None, - vllm_model_impl = 'vllm', - vllm_structured_outputs_regex = None, - vllm_sync_frequency = 1, - vllm_enable_sleep_mode = False, - log_completions = False, - log_completions_steps = 100, - num_completions_to_print = None, - wandb_log_unique_prompts = True, - vllm_sampling_params = None, - unsloth_num_chunks = -1, - unsloth_logit_chunk_multiplier = None, - unsloth_grpo_mini_batch = None, - max_seq_length = None, - **kwargs, - ): - if learning_rate < 1e-7: print(f'Unsloth: Your learning rate of `{learning_rate}` is too small and less than 1e-7! Consider increasing it, otherwise gradient updates will be close to 0!') - if learning_rate > 1: print(f'Unsloth: Your learning rate of `{learning_rate}` is way too larger > 1! Consider decreasing it to 1e-1, otherwise gradient updates will explode!') - _unsloth_fsdp2 = os.environ.get('FSDP_VERSION', '').strip() == '2' - _unsloth_fsdp_cfg = fsdp_config - if fsdp and isinstance(_unsloth_fsdp_cfg, str) and os.path.isfile(_unsloth_fsdp_cfg): - try: - import json as _unsloth_json - with open(_unsloth_fsdp_cfg, encoding = 'utf-8') as _unsloth_f: - _unsloth_fsdp_cfg = _unsloth_json.load(_unsloth_f) - except Exception: - _unsloth_fsdp_cfg = None - if fsdp and _unsloth_fsdp_cfg is None: - _unsloth_fsdp_cfg = {} - if fsdp and isinstance(_unsloth_fsdp_cfg, dict): - _unsloth_fsdp2 = _unsloth_fsdp2 or str(_unsloth_fsdp_cfg.get('fsdp_version', _unsloth_fsdp_cfg.get('version', 1))) == '2' - _unsloth_flash_adamw = str(getattr(optim, 'value', optim)).strip().lower() == 'adamw_flash' - if _unsloth_flash_adamw: optim = 'adamw_torch' - if _unsloth_fsdp2 and str(getattr(optim, 'value', optim)) == 'adamw_8bit': - print("Unsloth: FSDP2 shards parameters as DTensors, which bitsandbytes' 8-bit optimizers cannot step. Switching optim = 'adamw_8bit' to 'adamw_torch_fused'.") - optim = 'adamw_torch_fused' - if num_train_epochs is None: - num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override - if output_dir is None and save_strategy == 'steps' and save_steps == 500: - output_dir = 'unsloth_training_checkpoints' - save_strategy = 'no' - try: - from unsloth_zoo.dataset_num_proc import get_dataset_num_proc as _unsloth_get_dataset_num_proc - except Exception: - try: - from unsloth.dataset_num_proc import get_dataset_num_proc as _unsloth_get_dataset_num_proc - except Exception: - _unsloth_get_dataset_num_proc = None - if _unsloth_get_dataset_num_proc is not None: - dataset_num_proc = _unsloth_get_dataset_num_proc(dataset_num_proc, serial_as_none = True) - if os.environ.get('UNSLOTH_ENABLE_FLEX_ATTENTION', '0') == '1': - from unsloth_zoo.flex_attention import HAS_FLEX_ATTENTION - if HAS_FLEX_ATTENTION and pad_to_multiple_of is None: - from unsloth_zoo.flex_attention import FLEX_ATTENTION_BLOCK_SIZE - pad_to_multiple_of = FLEX_ATTENTION_BLOCK_SIZE - - if eval_steps is not None and eval_strategy != 'steps': - print(f'Unsloth: `eval_steps = {eval_steps}` is ignored because `eval_strategy` is {getattr(eval_strategy, "value", eval_strategy)!r}. Set `eval_strategy = "steps"` to evaluate every `eval_steps` steps.') - - if temperature <= 0: - raise ValueError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.') - elif temperature >= 10: - raise ValueError('Unsloth: Please set a positive non-zero temperature less than 10, since sampling will be quite erratic.') - - - # One dict so the filter sees the mirrored parameters AND `**kwargs`: - # filtering kwargs alone would double-bind any argument TRL renamed, - # since the new name is itself a mirrored parameter. - _unsloth_config_arguments = dict( - output_dir = output_dir, - per_device_train_batch_size = per_device_train_batch_size, - num_train_epochs = num_train_epochs, - max_steps = max_steps, - learning_rate = learning_rate, - lr_scheduler_type = lr_scheduler_type, - lr_scheduler_kwargs = lr_scheduler_kwargs, - warmup_steps = warmup_steps, - optim = optim, - optim_args = optim_args, - weight_decay = weight_decay, - adam_beta1 = adam_beta1, - adam_beta2 = adam_beta2, - adam_epsilon = adam_epsilon, - optim_target_modules = optim_target_modules, - gradient_accumulation_steps = gradient_accumulation_steps, - average_tokens_across_devices = average_tokens_across_devices, - max_grad_norm = max_grad_norm, - label_smoothing_factor = label_smoothing_factor, - bf16 = bf16, - fp16 = fp16, - bf16_full_eval = bf16_full_eval, - fp16_full_eval = fp16_full_eval, - tf32 = tf32, - gradient_checkpointing = gradient_checkpointing, - gradient_checkpointing_kwargs = gradient_checkpointing_kwargs, - torch_compile = torch_compile, - torch_compile_backend = torch_compile_backend, - torch_compile_mode = torch_compile_mode, - use_liger_kernel = use_liger_kernel, - liger_kernel_config = liger_kernel_config, - use_cache = use_cache, - neftune_noise_alpha = neftune_noise_alpha, - torch_empty_cache_steps = torch_empty_cache_steps, - auto_find_batch_size = auto_find_batch_size, - logging_strategy = logging_strategy, - logging_steps = logging_steps, - logging_first_step = logging_first_step, - log_on_each_node = log_on_each_node, - logging_nan_inf_filter = logging_nan_inf_filter, - include_num_input_tokens_seen = include_num_input_tokens_seen, - log_level = log_level, - log_level_replica = log_level_replica, - disable_tqdm = disable_tqdm, - report_to = report_to, - run_name = run_name, - project = project, - trackio_space_id = trackio_space_id, - eval_strategy = eval_strategy, - eval_steps = eval_steps, - eval_delay = eval_delay, - per_device_eval_batch_size = per_device_eval_batch_size, - prediction_loss_only = prediction_loss_only, - eval_on_start = eval_on_start, - eval_do_concat_batches = eval_do_concat_batches, - eval_use_gather_object = eval_use_gather_object, - eval_accumulation_steps = eval_accumulation_steps, - batch_eval_metrics = batch_eval_metrics, - save_only_model = save_only_model, - save_strategy = save_strategy, - save_steps = save_steps, - save_on_each_node = save_on_each_node, - save_total_limit = save_total_limit, - enable_jit_checkpoint = enable_jit_checkpoint, - push_to_hub = push_to_hub, - hub_token = hub_token, - hub_private_repo = hub_private_repo, - hub_model_id = hub_model_id, - hub_strategy = hub_strategy, - hub_always_push = hub_always_push, - hub_revision = hub_revision, - load_best_model_at_end = load_best_model_at_end, - metric_for_best_model = metric_for_best_model, - greater_is_better = greater_is_better, - ignore_data_skip = ignore_data_skip, - restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint, - full_determinism = full_determinism, - seed = seed, - data_seed = data_seed, - use_cpu = use_cpu, - accelerator_config = accelerator_config, - parallelism_config = parallelism_config, - dataloader_drop_last = dataloader_drop_last, - dataloader_num_workers = dataloader_num_workers, - dataloader_pin_memory = dataloader_pin_memory, - dataloader_persistent_workers = dataloader_persistent_workers, - dataloader_prefetch_factor = dataloader_prefetch_factor, - remove_unused_columns = remove_unused_columns, - label_names = label_names, - train_sampling_strategy = train_sampling_strategy, - length_column_name = length_column_name, - ddp_find_unused_parameters = ddp_find_unused_parameters, - ddp_bucket_cap_mb = ddp_bucket_cap_mb, - ddp_broadcast_buffers = ddp_broadcast_buffers, - ddp_backend = ddp_backend, - ddp_timeout = ddp_timeout, - fsdp = fsdp, - fsdp_config = fsdp_config, - deepspeed = deepspeed, - debug = debug, - skip_memory_metrics = skip_memory_metrics, - do_train = do_train, - do_eval = do_eval, - do_predict = do_predict, - resume_from_checkpoint = resume_from_checkpoint, - warmup_ratio = warmup_ratio, - logging_dir = logging_dir, - local_rank = local_rank, - model_init_kwargs = model_init_kwargs, - trust_remote_code = trust_remote_code, - router_aux_loss_coef = router_aux_loss_coef, - chat_template_path = chat_template_path, - dataset_text_field = dataset_text_field, - dataset_kwargs = dataset_kwargs, - dataset_num_proc = dataset_num_proc, - eos_token = eos_token, - max_length = max_length, - truncation_mode = truncation_mode, - shuffle_dataset = shuffle_dataset, - packing = packing, - packing_strategy = packing_strategy, - padding_free = padding_free, - pad_to_multiple_of = pad_to_multiple_of, - eval_packing = eval_packing, - completion_only_loss = completion_only_loss, - assistant_only_loss = assistant_only_loss, - loss_type = loss_type, - activation_offloading = activation_offloading, - pad_token = pad_token, - temperature = temperature, - top_p = top_p, - top_k = top_k, - lmbda = lmbda, - beta = beta, - max_completion_length = max_completion_length, - teacher_model_name_or_path = teacher_model_name_or_path, - teacher_model_revision = teacher_model_revision, - teacher_model_init_kwargs = teacher_model_init_kwargs, - teacher_tokenizer_name_or_path = teacher_tokenizer_name_or_path, - disable_dropout = disable_dropout, - seq_kd = seq_kd, - num_generations = num_generations, - generation_batch_size = generation_batch_size, - use_uld_loss = use_uld_loss, - uld_token_merge_strategy = uld_token_merge_strategy, - use_extended_uld = use_extended_uld, - uld_use_hybrid_loss = uld_use_hybrid_loss, - uld_hybrid_matched_weight = uld_hybrid_matched_weight, - uld_hybrid_unmatched_weight = uld_hybrid_unmatched_weight, - uld_crossentropy_weight = uld_crossentropy_weight, - uld_distillation_weight = uld_distillation_weight, - uld_student_temperature = uld_student_temperature, - uld_teacher_temperature = uld_teacher_temperature, - uld_skip_student_eos = uld_skip_student_eos, - uld_skip_teacher_eos = uld_skip_teacher_eos, - use_vllm = use_vllm, - vllm_mode = vllm_mode, - vllm_server_base_url = vllm_server_base_url, - vllm_server_host = vllm_server_host, - vllm_server_port = vllm_server_port, - vllm_server_timeout = vllm_server_timeout, - vllm_group_port = vllm_group_port, - vllm_gpu_memory_utilization = vllm_gpu_memory_utilization, - vllm_tensor_parallel_size = vllm_tensor_parallel_size, - vllm_max_model_length = vllm_max_model_length, - vllm_model_impl = vllm_model_impl, - vllm_structured_outputs_regex = vllm_structured_outputs_regex, - vllm_sync_frequency = vllm_sync_frequency, - vllm_enable_sleep_mode = vllm_enable_sleep_mode, - log_completions = log_completions, - log_completions_steps = log_completions_steps, - num_completions_to_print = num_completions_to_print, - wandb_log_unique_prompts = wandb_log_unique_prompts,**kwargs) - super().__init__(**_unsloth_filter_config_init_kwargs(GOLDConfig, _unsloth_config_arguments, mirrored_from = __class__)) - self.vllm_sampling_params = vllm_sampling_params - self.unsloth_num_chunks = unsloth_num_chunks - if unsloth_grpo_mini_batch is not None: - if self.generation_batch_size >= unsloth_grpo_mini_batch: - self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch - else: - raise ValueError( - f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, " - f"which is self.per_device_train_batch_size * gradient_accumulation_steps." - ) - self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier - self.max_seq_length = max_seq_length - # Unsloth: keep the reentrant checkpoint path - if getattr(self, 'gradient_checkpointing', False): - _gc_kwargs = getattr(self, 'gradient_checkpointing_kwargs', None) or {} - if _gc_kwargs.get('context_fn') is None and not _gc_kwargs.get('debug', False): - _gc_kwargs['use_reentrant'] = True - self.gradient_checkpointing_kwargs = _gc_kwargs - from unsloth.optimizers.flash_adamw import is_flash_adamw_marked, mark_flash_adamw, unmark_flash_adamw - if _unsloth_flash_adamw or is_flash_adamw_marked(self): - mark_flash_adamw(self) - else: - unmark_flash_adamw(self) - -pass - -class _UnslothGOLDTrainer(SFTTrainer): - _tag_names = ["trl", "gold"] - _name = "GOLD" - _paper = { - "title": "Unlocking On-Policy Distillation for Any Model Family", - # docstyle-ignore - "citation": textwrap.dedent( - """\ - @misc{patino2025unlocking, - title = {{Unlocking On-Policy Distillation for Any Model Family}}, - author = {Carlos Miguel Patiño and Kashif Rasul and Quentin Gallouédec and Ben Burtenshaw and Sergio Paniego and Vaibhav Srivastav and Thibaud Frere and Ed Beeching and Lewis Tunstall and Leandro von Werra and Thomas Wolf}, - year = 2025, - url = {https://huggingface.co/spaces/HuggingFaceH4/general-on-policy-logit-distillation}, - }""" - ), - } - - def __init__( - self, - model: PreTrainedModel | nn.Module | str | None = None, - teacher_model: PreTrainedModel | nn.Module | str = None, - args: GOLDConfig | None = None, - data_collator: DataCollator | None = None, # type: ignore - train_dataset: Dataset | None = None, - eval_dataset: Dataset | dict[str, Dataset] | None = None, - processing_class: ( - PreTrainedTokenizerBase | BaseImageProcessor | FeatureExtractionMixin | ProcessorMixin | None - ) = None, - compute_metrics: Callable[[EvalPrediction], dict] | None = None, - callbacks: list[TrainerCallback] | None = None, - optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = ( - None, - None, - ), - preprocess_logits_for_metrics: (Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None) = None, - peft_config: Optional["PeftConfig"] = None, - ): - - if hasattr(model, 'vllm_engine') and hasattr(args, 'use_vllm'): - _unsloth_fsdp_cfg = getattr(args, 'fsdp_config', None) - if getattr(args, 'fsdp', None) and isinstance(_unsloth_fsdp_cfg, dict) and str(_unsloth_fsdp_cfg.get('version', _unsloth_fsdp_cfg.get('fsdp_version', 1))) == '2': - raise NotImplementedError('Unsloth: `fast_inference = True` (vLLM) is not supported with FSDP2.\nLoad the model with `fast_inference = False` to train with FSDP2, or launch without FSDP2 to use vLLM.') - if (getattr(args, 'use_vllm', False) == False): - args.use_vllm = True - if getattr(args, 'top_k', -1) is None or getattr(args, 'top_k', -1) == 0: - args.top_k = -1 - self.model_name_or_path = model if isinstance(model, str) else model.config._name_or_path - self.model_revision = (args.model_init_kwargs or {}).get("revision") if isinstance(model, str) else None - teacher_revision = ( - (args.teacher_model_init_kwargs or {}).get("revision", args.teacher_model_revision) - if isinstance(teacher_model, str) - else None - ) - dataset_sample = next(iter(train_dataset)) if train_dataset is not None else {} - if processing_class is None: - model_id = model if isinstance(model, str) else get_config_model_id(model.config) - processing_class = AutoProcessor.from_pretrained( - model_id, revision=self.model_revision, trust_remote_code=args.trust_remote_code - ) - # simplified logic from SFTTrainer - # Handle pad token for processors or tokenizers - if isinstance(processing_class, ProcessorMixin): - self._tokenizer = processing_class.tokenizer - self._is_vlm = True - if self._tokenizer.pad_token is None: - self._tokenizer.pad_token = self._tokenizer.eos_token - else: - self._tokenizer = processing_class - self._is_vlm = False - - self.pad_token_id = self._tokenizer.pad_token_id - - # VLM distillation: only VLM-to-VLM is supported. Both student and teacher must be - # VLMs so that both receive images and multimodal inputs. - self._teacher_processor = None - self._is_cross_architecture_vlm = False - if self._is_vlm: - if isinstance(teacher_model, str): - # Teacher not yet instantiated -- validate it's a VLM - teacher_proc = AutoProcessor.from_pretrained( - teacher_model, revision=teacher_revision, trust_remote_code=args.trust_remote_code - ) - if not isinstance(teacher_proc, ProcessorMixin): - raise ValueError( - "VLM distillation requires both student and teacher to be vision-language models. " - "The student has a `ProcessorMixin` but the teacher does not." - ) - teacher_model_type = AutoConfig.from_pretrained( - teacher_model, trust_remote_code=args.trust_remote_code - ).model_type - else: - # Teacher already instantiated — check if it looks like a VLM by checking for a vision config - if teacher_model.config.vision_config is None: - raise ValueError( - "VLM distillation requires both student and teacher to be vision-language models. " - "The student has a `ProcessorMixin` but the teacher model does not appear to be a VLM " - "(missing `vision_config`)." - ) - teacher_model_type = teacher_model.config.model_type - - # Check for cross-architecture VLM distillation - student_model_type = ( - AutoConfig.from_pretrained(model, trust_remote_code=args.trust_remote_code).model_type - if isinstance(model, str) - else model.config.model_type - ) - is_cross_architecture = student_model_type and teacher_model_type != student_model_type - self._is_cross_architecture_vlm = is_cross_architecture - if is_cross_architecture: - warnings.warn( - f"Cross-architecture VLM distillation detected: student is '{student_model_type}', " - f"teacher is '{teacher_model_type}'. Images will be processed separately through each " - "model's processor, which may increase memory usage and computation time." - ) - if is_cross_architecture or args.use_uld_loss: - self._teacher_processor = ( - teacher_proc - if isinstance(teacher_model, str) - else AutoProcessor.from_pretrained( - teacher_model.config._name_or_path, - revision=teacher_revision, - trust_remote_code=args.trust_remote_code, - ) - ) - if self._is_cross_architecture_vlm and not args.use_uld_loss: - raise ValueError( - "Cross-architecture VLM distillation (student and teacher have different `model_type`) is not " - "supported with the standard JSD loss because the models require different image token formats " - "and tokenizers. Please set `use_uld_loss=True` in your GOLDConfig to enable cross-tokenizer " - "alignment via ULD loss." - ) - self._is_vision_dataset = "image" in dataset_sample or "images" in dataset_sample - if self._is_vision_dataset and not self._is_vlm: - raise ValueError( - "The dataset appears to be vision-related (contains 'image' or 'images' keys), but the provided " - "model does not seem to be a vision-language model. Please check your model and dataset." - ) - if self._is_vlm and args.max_length is not None and args.truncation_mode == "keep_end": - raise ValueError( - "truncation_mode='keep_end' is not supported for vision-language models. Image tokens reside " - "inside the prompt portion of the sequence; depending on the example, keep_end may silently " - "drop them, causing pixel_values to be forwarded to the model with no corresponding visual " - "tokens in input_ids. Use truncation_mode='keep_start' (the default) or set max_length=None." - ) - - # Respect a user-provided data_collator for text; otherwise, pick the right collator based on modality. - # For VLMs, always use identity collator to preserve raw PIL images in the dataloader. - # Raw images are needed for: [1] vLLM generation, [2] cross-architecture teacher processing. - # A separate _vlm_collator is stored for on-the-fly collation inside _fill_buffer. - if self._is_vision_dataset and data_collator is not None: - raise ValueError( - "Passing a custom data collator is not supported for VLM training. GOLD manages its own collation " - "to preserve raw images for generation and teacher processing; leave `data_collator=None`." - ) - self._vlm_collator = None - if data_collator is None: - if self._is_vision_dataset: - self._vlm_collator = DataCollatorForVisionLanguageChatML( - processor=processing_class, - max_length=args.max_length, - ) - data_collator = identity - else: - data_collator = DataCollatorForChatML(tokenizer=self._tokenizer, max_length=args.max_length) - - if args.teacher_model_init_kwargs is None: - teacher_model_init_kwargs = {} - elif not isinstance(teacher_model, str): - raise ValueError( - "You passed teacher_model_init_kwargs to the GOLDConfig, but your teacher_model is already instantiated." - ) - else: - teacher_model_init_kwargs = args.teacher_model_init_kwargs - teacher_model_init_kwargs["dtype"] = ( - teacher_model_init_kwargs["dtype"] - if teacher_model_init_kwargs["dtype"] in ["auto", None] - else getattr(torch, teacher_model_init_kwargs["dtype"]) - ) - - if args.use_uld_loss and args.teacher_tokenizer_name_or_path is None: - if isinstance(teacher_model, str): - args.teacher_tokenizer_name_or_path = teacher_model - elif teacher_model.config._name_or_path: - args.teacher_tokenizer_name_or_path = teacher_model.config._name_or_path - else: - raise ValueError( - "`teacher_tokenizer_name_or_path` must be set when using ULD loss with a pre-instantiated teacher model." - ) - - # The teacher revision pins a commit in the teacher model's repo, so it only applies to a tokenizer served from - # that same repo. ULD's cross-tokenizer setup points `teacher_tokenizer_name_or_path` at a different repo, where - # that commit does not exist. - teacher_tokenizer_revision = teacher_revision if args.teacher_tokenizer_name_or_path == teacher_model else None - - if isinstance(teacher_model, str): - init_kwargs = dict(teacher_model_init_kwargs) - if args.teacher_model_revision is not None: - init_kwargs.setdefault("revision", args.teacher_model_revision) - init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) - # Distributed training requires device_map=None ["auto" fails] - if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: - init_kwargs["device_map"] = None - teacher_model = create_model_from_path(teacher_model, **init_kwargs) - self.use_uld_loss = args.use_uld_loss - self.teacher_tokenizer = None - if args.use_uld_loss and self._teacher_processor is not None: - self.teacher_tokenizer = self._teacher_processor.tokenizer - if self.teacher_tokenizer.pad_token is None: - self.teacher_tokenizer.pad_token = self.teacher_tokenizer.eos_token - elif args.use_uld_loss and args.teacher_tokenizer_name_or_path is not None: - self.teacher_tokenizer = AutoTokenizer.from_pretrained( - args.teacher_tokenizer_name_or_path, - revision=teacher_tokenizer_revision, - trust_remote_code=args.trust_remote_code, - ) - if self.teacher_tokenizer.pad_token is None: - self.teacher_tokenizer.pad_token = self.teacher_tokenizer.eos_token - - # Hybrid ULD loss configuration is handled in ULDLoss class - - super().__init__( - model, - args=args, - data_collator=data_collator, - train_dataset=train_dataset, - eval_dataset=eval_dataset, - processing_class=processing_class, - compute_metrics=compute_metrics, - callbacks=callbacks, - optimizers=optimizers, - preprocess_logits_for_metrics=preprocess_logits_for_metrics, - peft_config=peft_config, - ) - - if args.disable_dropout: - disable_dropout_in_model(self.model) - if not args.use_uld_loss: - teacher_model.resize_token_embeddings(self.model.config.get_text_config().vocab_size) - - if self.is_deepspeed_enabled: - self.teacher_model = prepare_deepspeed(teacher_model, self.accelerator) - else: - self.teacher_model = self.accelerator.prepare_model(teacher_model, evaluation_mode=True) - - self.lmbda = args.lmbda - self.beta = args.beta - self.temperature = args.temperature - self.top_p = args.top_p - self.seq_kd = args.seq_kd - self.num_generations = args.num_generations - - # Track per-step loss statistics for on/off-policy batches [used in logging] - self._on_policy_loss_total = 0.0 - self._off_policy_loss_total = 0.0 - self._on_policy_step_equiv = 0.0 - self._off_policy_step_equiv = 0.0 - - # Buffering for rollouts across gradient accumulation steps - self._buffered_inputs = None - self._buffered_on_policy = None - self._buffered_text_logs = None - self._step = 0 - - # Hybrid ULD matched/unmatched accumulators [logged every step when ULD hybrid is used] - self._matched_sum = 0.0 - self._unmatched_sum = 0.0 - self._matched_step_eq = 0.0 - self._unmatched_step_eq = 0.0 - - self.uld_loss_fn = None - if self.use_uld_loss: - self.uld_loss_fn = ULDLoss( - config=args, - student_tokenizer=self._tokenizer, - teacher_tokenizer=self.teacher_tokenizer, - device=self.accelerator.device, - ) - - generation_kwargs = { - "max_new_tokens": args.max_completion_length, - "temperature": args.temperature, - "top_p": args.top_p, - "do_sample": True, - "top_k": args.top_k, - "pad_token_id": self.pad_token_id, - } - self.generation_config = GenerationConfig(**generation_kwargs) - # Keep training-specific generation kwargs to overwrite model's original generation config - self.generation_kwargs = generation_kwargs - if ( - hasattr(self.model.generation_config, "eos_token_id") - and self.model.generation_config.eos_token_id is not None - ): - self.generation_config.eos_token_id = self.model.generation_config.eos_token_id - - # Initialize the metrics - self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} - self._total_train_tokens = 0 - self.log_completions = args.log_completions - self.log_completion_steps = args.log_completions_steps - self.wandb_log_unique_prompts = args.wandb_log_unique_prompts - self.num_completions_to_print = args.num_completions_to_print - # maxlen is set to the total number of forward passes per step. This value of `maxlen` ensures we log only the - # final optimization step. - maxlen = self.accelerator.num_processes * args.per_device_train_batch_size * args.gradient_accumulation_steps - self._textual_logs = { - "prompt": deque(maxlen=maxlen), - "completion": deque(maxlen=maxlen), - "rewards": defaultdict(lambda: deque(maxlen=maxlen)), - "advantages": deque(maxlen=maxlen), - } - - self.use_vllm = args.use_vllm - if self.use_vllm: - if not is_vllm_available(): - raise ImportError( - "vLLM is not available and use_vllm is set to True. Please install vLLM with " - "`pip install vllm` to use it." - ) - self.vllm_generation = VLLMGeneration( - model=self.model, - accelerator=self.accelerator, - processing_class=self.processing_class, - mode=args.vllm_mode, - structured_outputs_regex=args.vllm_structured_outputs_regex, - server_base_url=args.vllm_server_base_url, - server_host=args.vllm_server_host, - server_port=args.vllm_server_port, - group_port=args.vllm_group_port, - server_timeout=args.vllm_server_timeout, - tensor_parallel_size=args.vllm_tensor_parallel_size, - gpu_memory_utilization=args.vllm_gpu_memory_utilization, - max_model_length=args.vllm_max_model_length or args.max_length, - max_num_seqs=args.per_device_train_batch_size * args.gradient_accumulation_steps, - enable_sleep_mode=args.vllm_enable_sleep_mode, - model_impl=args.vllm_model_impl, - trust_remote_code=args.trust_remote_code, - repetition_penalty=getattr(args, "repetition_penalty", 1.0), - temperature=args.temperature, - top_p=args.top_p, - top_k=args.top_k, - min_p=getattr(args, "min_p", 0.0), - max_completion_length=args.max_completion_length, - logprobs=None, - ) - self.vllm_sync_frequency = args.vllm_sync_frequency - self._last_vllm_sync_step = -self.vllm_sync_frequency - - def _set_signature_columns_if_needed(self): - super()._set_signature_columns_if_needed() - required_columns = [ - "prompt", - "completion", - "prompts", - "prompt_attention_mask", - "messages", - "chat_template_kwargs", - "tools", - "original_prompt_text", - "original_completion_text", - "byte_offsets", - "completion_mask", - "images", - "image", - "pixel_values", - "image_grid_thw", - "image_position_ids", - "pixel_attention_mask", - "image_sizes", - "spatial_shapes", - "token_type_ids", - "mm_token_type_ids", - ] - if self._signature_columns is None: - self._signature_columns = required_columns - else: - for column in required_columns: - if column not in self._signature_columns: - self._signature_columns.append(column) - - def _get_train_sampler(self, dataset=None): - if dataset is None: - dataset = self.train_dataset - return RepeatSampler( - data_source=dataset, - mini_repeat_count=self.num_generations, - batch_size=self.args.generation_batch_size * self.accelerator.num_processes, - repeat_count=self.args.gradient_accumulation_steps, - shuffle=True, - seed=self.args.seed, - ) - - def get_train_dataloader(self): - """ - Override Trainer.get_train_dataloader to load one generation batch per optimizer window. - - The dataloader yields local batches of size `per_device_train_batch_size * gradient_accumulation_steps`. The - `RepeatSampler` (with `repeat_count=gradient_accumulation_steps`) ensures each generation batch is sampled - `gradient_accumulation_steps` times so Trainer's loop iterates the correct number of times. Only the first - batch in each window triggers `_fill_buffer`; the rest are ignored by `_prepare_inputs`. - """ - if self.train_dataset is None: - raise ValueError("Trainer: training requires a train_dataset.") - - train_dataset = self.train_dataset - data_collator = self.data_collator - if is_datasets_available() and isinstance(train_dataset, Dataset): - train_dataset = self._remove_unused_columns(train_dataset, description="training") - else: - data_collator = self._get_collator_with_removed_columns(data_collator, description="training") - - dataloader_params = { - "batch_size": self._train_batch_size * self.args.gradient_accumulation_steps, - "collate_fn": data_collator, - "num_workers": self.args.dataloader_num_workers, - "pin_memory": self.args.dataloader_pin_memory, - "persistent_workers": self.args.dataloader_persistent_workers, - } - - if not isinstance(train_dataset, torch.utils.data.IterableDataset): - dataloader_params["sampler"] = self._get_train_sampler() - dataloader_params["drop_last"] = self.args.dataloader_drop_last - dataloader_params["worker_init_fn"] = partial( - seed_worker, - num_workers=self.args.dataloader_num_workers, - rank=self.args.process_index, - ) - if self.args.dataloader_num_workers > 0: - dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor - - return self.accelerator.prepare(DataLoader(train_dataset, **dataloader_params)) - - @profiling_decorator - def _prepare_inputs(self, generation_batch: dict[str, torch.Tensor | Any]) -> dict[str, torch.Tensor | Any]: - if not self.model.training: - # Evaluation is off-policy (no generation): the student never samples, both models are forwarded over - # the dataset's ground-truth prompt+completion and the distillation loss is taken over the completion. - # For text the collated tensor dict is consumed directly. For VLMs the identity collator yields raw - # dicts (to preserve PIL images for the train-time on-policy path), so collate them here -- mirroring - # the off-policy slice construction in _fill_buffer -- including the raw images/prompts the cross-arch - # / ULD teacher processor needs. - if self._vlm_collator is not None: - pending_slice = {"_gold_vlm_lazy_examples": list(generation_batch)} - if self._teacher_processor is not None: - raw_images, raw_prompts = self._extract_images_and_prompts(list(generation_batch)) - pending_slice["_gold_vlm_raw_images"] = raw_images - pending_slice["_gold_vlm_raw_prompts"] = raw_prompts - return self._materialize_vlm_slice(pending_slice) - return generation_batch - - buffer_steps = self.args.gradient_accumulation_steps - if self._step % buffer_steps == 0 or self._buffered_inputs is None: - self._fill_buffer(generation_batch, buffer_steps) - - slice_idx = self._step % buffer_steps - inputs = self._buffered_inputs[slice_idx] - if isinstance(inputs, dict): - if "_gold_vlm_on_policy_raw_examples" in inputs: - inputs, text_logs = self._generate_on_policy_vlm_slice(inputs) - self._buffered_inputs[slice_idx] = inputs - self._buffered_text_logs[slice_idx] = text_logs - elif "_gold_vlm_lazy_examples" in inputs: - inputs = self._materialize_vlm_slice(inputs) - self._buffered_inputs[slice_idx] = inputs - self._step += 1 - return inputs - - def _generate_on_policy_vlm_slice( - self, pending_slice: dict[str, Any] - ) -> tuple[dict[str, torch.Tensor | Any], tuple[list[str], list[str]]]: - """Generate and collate one non-vLLM on-policy VLM slice immediately before it is consumed.""" - raw_examples = pending_slice["_gold_vlm_on_policy_raw_examples"] - generation_examples = [] - for example in raw_examples: - generation_example = dict(example) - completion = generation_example.get("completion") - generation_example["completion"] = ( - [{"role": "assistant", "content": [{"type": "text", "text": ""}]}] - if isinstance(completion, list) - else "" - ) - generation_examples.append(generation_example) - collated = self._vlm_collator(generation_examples) - collated = { - k: (v.to(self.accelerator.device) if isinstance(v, torch.Tensor) else v) for k, v in collated.items() - } - - with unwrap_model_for_generation( - self.model, self.accelerator, generation_kwargs=self.generation_kwargs - ) as unwrapped_model: - ( - new_input_ids, - new_attention_mask, - new_labels, - prompt_texts, - completion_texts, - ) = self.generate_on_policy_outputs( - unwrapped_model, - collated, - self.generation_config, - ) - - updated_slice = dict(collated) - updated_slice["input_ids"] = new_input_ids - updated_slice["attention_mask"] = new_attention_mask - updated_slice["labels"] = new_labels - # Rebuild sequence-length-dependent keys to match new input_ids shape - new_seq_len = new_input_ids.shape[1] - prompt_seq_len = collated["prompts"].shape[1] - for k in self._SEQUENCE_KEYS: - if k in updated_slice: - sequence_dtype = updated_slice[k].dtype - prompt_part = self._get_prompt_sequence_key(collated, k) - comp_part = torch.zeros( - new_input_ids.shape[0], - new_seq_len - prompt_seq_len, - dtype=sequence_dtype, - device=new_input_ids.device, - ) - updated_slice[k] = torch.cat([prompt_part, comp_part], dim=1) - if "original_prompt_text" not in updated_slice: - updated_slice["original_prompt_text"] = prompt_texts - # Keep special tokens (e.g. EOS) so the teacher inputs are built from text covering the same generated - # content as the supervised `labels != -100` tokens. The alignment itself remains byte-offset based. - updated_slice["original_completion_text"] = completion_texts - self._maybe_add_completion_byte_offsets(updated_slice) - if self._teacher_processor is not None: - updated_slice["_raw_images"] = pending_slice["_gold_vlm_raw_images"] - updated_slice["_raw_prompts"] = pending_slice["_gold_vlm_raw_prompts"] - - return updated_slice, (prompt_texts, completion_texts) - - def _materialize_vlm_slice(self, pending_slice: dict[str, Any]) -> dict[str, torch.Tensor | Any]: - """Collate one pending VLM slice immediately before it is consumed.""" - slice_inputs = self._vlm_collator([dict(example) for example in pending_slice["_gold_vlm_lazy_examples"]]) - slice_inputs = { - k: (v.to(self.accelerator.device) if isinstance(v, torch.Tensor) else v) for k, v in slice_inputs.items() - } - - if self.use_uld_loss and self.teacher_tokenizer is not None: - slice_inputs = self._ensure_original_text_fields(slice_inputs) - if "original_prompt_text" not in slice_inputs or "original_completion_text" not in slice_inputs: - raise ValueError( - "Off-policy batch missing 'original_prompt_text' or 'original_completion_text' fields. " - "When using ULD loss with cross-tokenizer alignment, datasets must be prepared with " - "_prepare_dataset_with_original_text(). Ensure your dataset includes these fields." - ) - - if self._teacher_processor is not None: - slice_inputs["_raw_images"] = pending_slice["_gold_vlm_raw_images"] - slice_inputs["_raw_prompts"] = pending_slice["_gold_vlm_raw_prompts"] - - return slice_inputs - - @staticmethod - def _build_sequence_batch( - new_input_ids: torch.Tensor, - prompt_lengths: torch.Tensor, - attention_mask: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Build attention mask and labels from full sequences and prompt lengths.""" - prompt_lengths = prompt_lengths.to(device=new_input_ids.device, dtype=torch.long) - positions = torch.arange(new_input_ids.shape[1], device=new_input_ids.device).unsqueeze(0) - completion_mask = positions >= prompt_lengths.unsqueeze(1) - - new_attention_mask = attention_mask.to(device=new_input_ids.device, dtype=new_input_ids.dtype) - - new_labels = torch.full_like(new_input_ids, -100) - new_labels[completion_mask & new_attention_mask.bool()] = new_input_ids[ - completion_mask & new_attention_mask.bool() - ] - - return new_attention_mask, new_labels - - def _extract_images_and_prompts(self, examples: list[dict]) -> tuple[list | None, list]: - """ - Extract per-example images and build prompts with multimodal messages, mirroring GRPOTrainer. - - Returns `(images, prompts)` where `images` is a per-example list (entries may be `None`), or `None` when the - batch carries no images, and `prompts` are the prepared multimodal messages. - """ - if "images" in examples[0]: - images = [example.get("images") for example in examples] - elif "image" in examples[0]: - images = [[example.get("image")] if example.get("image") is not None else None for example in examples] - else: - images = None - if images is not None and all(img_list is None or img_list == [] for img_list in images): - images = None - - prompts = [example["prompt"] for example in examples] - if images is not None: - prompts = [ - prepare_multimodal_messages(prompt, images=img_list) - for prompt, img_list in zip(prompts, images, strict=True) - ] - return images, prompts - - def _decode_completion_texts_from_labels(self, slice_inputs: dict[str, torch.Tensor | Any]) -> list[str] | None: - """Decode completion text from labels when raw text is absent.""" - labels = slice_inputs.get("labels") - if labels is None or not isinstance(labels, torch.Tensor): - return None - - labels_cpu = labels.detach().cpu() - decoded_completion_tokens: list[list[int]] = [] - for row in labels_cpu: - token_ids = row[row != -100].tolist() - decoded_completion_tokens.append(token_ids) - - return self.processing_class.batch_decode( - decoded_completion_tokens, - skip_special_tokens=False, - clean_up_tokenization_spaces=False, - ) - - def _ensure_original_text_fields( - self, slice_inputs: dict[str, torch.Tensor | Any] - ) -> dict[str, torch.Tensor | Any]: - """Populate original prompt/completion text fields when missing.""" - if "original_prompt_text" in slice_inputs and "original_completion_text" in slice_inputs: - return slice_inputs - - prompts = slice_inputs.get("prompts") - if prompts is None or not isinstance(prompts, torch.Tensor): - return slice_inputs - - prompt_texts = self.processing_class.batch_decode( - prompts, - skip_special_tokens=False, - clean_up_tokenization_spaces=False, - ) - completion_texts = self._decode_completion_texts_from_labels(slice_inputs) - if completion_texts is None: - return slice_inputs - - updated_slice = dict(slice_inputs) - updated_slice["original_prompt_text"] = prompt_texts - updated_slice["original_completion_text"] = completion_texts - return updated_slice - - @staticmethod - def _get_prompt_sequence_key(inputs: dict[str, torch.Tensor | Any], key: str) -> torch.Tensor: - """Align a sequence-length-dependent key with the left-padded prompt tensor.""" - values = inputs[key] - prompts = inputs["prompts"] - prompt_attention_mask = inputs.get("prompt_attention_mask") - - if prompt_attention_mask is None: - return values[:, : prompts.shape[1]] - - prompt_values = values.new_zeros(prompts.shape) - for i, mask in enumerate(prompt_attention_mask.bool()): - prompt_length = int(mask.sum().item()) - if prompt_length: - prompt_values[i, mask] = values[i, :prompt_length] - return prompt_values - - _SEQUENCE_KEYS = ("token_type_ids", "mm_token_type_ids") - _MODEL_INPUT_RESERVED_KEYS = frozenset( - ( - "input_ids", - "attention_mask", - "labels", - "prompts", - "prompt_attention_mask", - "completion_mask", - "assistant_masks", - "original_prompt_text", - "original_completion_text", - "byte_offsets", - ) - ) - - def _get_model_forward_kwargs( - self, inputs: dict[str, torch.Tensor | Any], exclude: tuple[str, ...] = () - ) -> dict[str, torch.Tensor]: - reserved_keys = self._MODEL_INPUT_RESERVED_KEYS | set(exclude) - return { - k: v - for k, v in inputs.items() - if k not in reserved_keys and not k.startswith("_") and isinstance(v, torch.Tensor) - } - - def _maybe_add_completion_byte_offsets(self, updated_slice: dict[str, torch.Tensor | Any]) -> None: - """Attach completion-relative byte offsets to on-policy ULD batches. - - Derived from the sampled ids via ``piece_byte_len`` (no decode→re-encode round-trip). - """ - if not ( - self.use_uld_loss - and self.teacher_tokenizer is not None - and self.uld_loss_fn is not None - and self.uld_loss_fn.use_extended_uld - ): - return - - new_input_ids = updated_slice["input_ids"] - new_labels = updated_slice["labels"] - seq_len = new_input_ids.shape[1] - - # convert_ids_to_tokens is a tokenizer method; VLM processors expose it via `.tokenizer`. - tokenizer = self._tokenizer - - rows: list[list[tuple[int, int]]] = [] - for row_ids, row_labels in zip(new_input_ids.cpu().tolist(), new_labels.cpu().tolist(), strict=True): - offs: list[tuple[int, int]] = [(0, 0)] * seq_len - cumulative = 0 - for pos, (tid, label) in enumerate(zip(row_ids, row_labels, strict=True)): - if label == -100: - continue - nb = piece_byte_len(tokenizer.convert_ids_to_tokens([tid])[0]) - offs[pos] = (cumulative, cumulative + nb) - cumulative += nb - rows.append(offs) - updated_slice["byte_offsets"] = torch.tensor(rows, dtype=torch.long, device=new_input_ids.device) - - @profiling_decorator - def _fill_buffer( - self, - generation_batch: dict[str, torch.Tensor | Any] | list[dict], - buffer_steps: int, - ): - # The batch is split into `buffer_steps` equal chunks via floor division (here and in - # `split_tensor_dict` below). A batch size that is not divisible by `buffer_steps` would - # silently drop examples or produce empty chunks, so fail fast instead. - batch_len = ( - len(generation_batch) - if isinstance(generation_batch, list) - else next(t for t in generation_batch.values() if t is not None).shape[0] - ) - if batch_len % buffer_steps != 0: - raise ValueError( - "The generation batch size must be divisible by gradient_accumulation_steps. Set " - "dataloader_drop_last=True or adjust per_device_train_batch_size / gradient_accumulation_steps." - ) - - if self._vlm_collator is not None: - # Identity collator path: generation_batch is list[dict] with raw PIL images. - # Split into chunks via list slicing, then collate on-the-fly per slice. - chunk_size = len(generation_batch) // buffer_steps - raw_slices = [generation_batch[i * chunk_size : (i + 1) * chunk_size] for i in range(buffer_steps)] - slices = None # not used in this path - else: - raw_slices = None # not used in this path - slices = split_tensor_dict(generation_batch, buffer_steps) - - if self.accelerator.is_main_process: - on_policy_flags = [random.random() <= self.lmbda for _ in range(buffer_steps)] - else: - on_policy_flags = [False] * buffer_steps - - on_policy_flags = broadcast_object_list(on_policy_flags, from_process=0) - on_policy_indices = [i for i, flag in enumerate(on_policy_flags) if flag] - - self._buffered_inputs = [None] * buffer_steps - self._buffered_on_policy = on_policy_flags - self._buffered_text_logs = [None] * buffer_steps - - for i, flag in enumerate(on_policy_flags): - if not flag: - if self._vlm_collator is not None: - # Extract raw images and prompts BEFORE collation, since the collator - # mutates examples in place (pops "image", overwrites "prompt"). - slice_inputs = {"_gold_vlm_lazy_examples": raw_slices[i]} - if self._teacher_processor is not None: - raw_images, raw_prompts = self._extract_images_and_prompts(raw_slices[i]) - slice_inputs["_gold_vlm_raw_images"] = raw_images - slice_inputs["_gold_vlm_raw_prompts"] = raw_prompts - self._buffered_inputs[i] = slice_inputs - continue - - slice_inputs = slices[i] - - if ( - self.use_uld_loss - and self.teacher_tokenizer is not None - and ("original_prompt_text" not in slice_inputs or "original_completion_text" not in slice_inputs) - ): - raise ValueError( - "Off-policy batch missing 'original_prompt_text' or 'original_completion_text' fields. " - "Use the default DataCollatorForChatML (or a collator that emits these fields) so the " - "teacher tokenizer has source text to align against." - ) - if ( - self.use_uld_loss - and self.teacher_tokenizer is not None - and self.uld_loss_fn.use_extended_uld - and "byte_offsets" not in slice_inputs - ): - raise ValueError( - "Off-policy batch missing `byte_offsets`. Use the default DataCollatorForChatML or set " - "`use_extended_uld=False`." - ) - - self._buffered_inputs[i] = slice_inputs - - if on_policy_indices: - if self._vlm_collator is not None: - self._generate_on_policy_vlm_raw(raw_slices, on_policy_indices) - else: - self._generate_on_policy_for_slices(slices, on_policy_indices) - - @profiling_decorator - def _generate_on_policy_for_slices( - self, slices: list[dict[str, torch.Tensor | Any]], on_policy_indices: list[int] - ): - prompt_ids_list = [] - local_slice_indices = [] - # Budget the prompt to max_length - max_new_tokens before generation, keeping the END of the - # prompt (the generation marker), so the student generates from and trains on the same context - # and prompt + completion fit in max_length. - max_completion_length = self.generation_config.max_new_tokens - prompt_max_length = max(1, self.args.max_length - max_completion_length) if self.args.max_length else None - for slice_idx in on_policy_indices: - slice_inputs = slices[slice_idx] - prompt_attention_mask = slice_inputs.get("prompt_attention_mask") - for prompt_idx, prompt in enumerate(slice_inputs["prompts"]): - if prompt_attention_mask is not None: - prompt = prompt[prompt_attention_mask[prompt_idx].bool()] - prompt_ids = prompt.tolist() - if prompt_max_length is not None and len(prompt_ids) > prompt_max_length: - prompt_ids = prompt_ids[-prompt_max_length:] - prompt_ids_list.append(prompt_ids) - local_slice_indices.append(slice_idx) - - prompts_text = self.processing_class.batch_decode( - prompt_ids_list, - skip_special_tokens=False, - ) - - if not self.use_vllm: - self._generate_non_vllm_for_slices(slices, on_policy_indices) - return - - if ( - self.state.global_step != self._last_vllm_sync_step - and self.state.global_step >= self._last_vllm_sync_step + self.vllm_sync_frequency - ): - self.vllm_generation.sync_weights() - self._last_vllm_sync_step = self.state.global_step - - _, completion_ids, _, _ = self.vllm_generation.generate( - prompts=prompt_ids_list, - images=None, - num_generations=self.num_generations, - ) - - self._process_completions_to_buffer( - slices, - on_policy_indices, - local_slice_indices, - completion_ids, - prompt_ids_list, - prompts_text, - self.generation_config.max_new_tokens, - ) - - def _generate_non_vllm_for_slices(self, slices: list[dict[str, torch.Tensor | Any]], on_policy_indices: list[int]): - """Fallback generation without vLLM (uses model.generate per slice).""" - with unwrap_model_for_generation( - self.model, - self.accelerator, - generation_kwargs=self.generation_kwargs, - ) as unwrapped_model: - max_completion_length = self.generation_config.max_new_tokens - prompt_max_length = max(1, self.args.max_length - max_completion_length) if self.args.max_length else None - pad_token_id = self.processing_class.pad_token_id if self.processing_class.pad_token_id is not None else 0 - for slice_idx in on_policy_indices: - slice_inputs = slices[slice_idx] - if prompt_max_length is not None: - # Budget per real token (keep-end), mirroring the vLLM path so both keep the same tokens. - # Trigger on the real-token count, not the padded width, and rebuild only when a row - # overflows, so no pad columns are carried into the budgeted prompt. - prompts = slice_inputs["prompts"] - prompt_attention_mask = slice_inputs.get("prompt_attention_mask") - budgeted_rows = [] - needs_budget = False - for row_idx, prompt in enumerate(prompts): - real_ids = ( - prompt[prompt_attention_mask[row_idx].bool()] - if prompt_attention_mask is not None - else prompt - ) - if real_ids.shape[0] > prompt_max_length: - real_ids = real_ids[-prompt_max_length:] - needs_budget = True - budgeted_rows.append(real_ids) - if needs_budget: - slice_inputs = dict(slice_inputs) - slice_inputs["prompts"] = pad(budgeted_rows, padding_side="left", padding_value=pad_token_id) - slice_inputs["prompt_attention_mask"] = pad( - [torch.ones(row.shape[0], dtype=torch.long, device=row.device) for row in budgeted_rows], - padding_side="left", - padding_value=0, - ) - result = self.generate_on_policy_outputs( - unwrapped_model, - slice_inputs, - self.generation_config, - ) - ( - new_input_ids, - new_attention_mask, - new_labels, - prompt_texts, - completion_texts, - ) = result - - updated_slice = dict(slice_inputs) - updated_slice["input_ids"] = new_input_ids - updated_slice["attention_mask"] = new_attention_mask - updated_slice["labels"] = new_labels - updated_slice["original_prompt_text"] = prompt_texts - updated_slice["original_completion_text"] = completion_texts - self._maybe_add_completion_byte_offsets(updated_slice) - - self._buffered_inputs[slice_idx] = updated_slice - self._buffered_text_logs[slice_idx] = (prompt_texts, completion_texts) - - def _generate_on_policy_vlm_raw(self, raw_slices: list[list[dict]], on_policy_indices: list[int]): - """On-policy generation from raw VLM examples, preserving PIL images for vLLM.""" - all_prompt_ids = [] - all_images = [] - all_prompts = [] - all_raw_examples = [] - local_slice_indices = [] - slice_raw_data = {} - max_completion_length = self.generation_config.max_new_tokens - prompt_max_length = max(1, self.args.max_length - max_completion_length) if self.args.max_length else None - - for slice_idx in on_policy_indices: - raw_examples = raw_slices[slice_idx] - - images, prompts = self._extract_images_and_prompts(raw_examples) - - prompts = [ - [ - ( - {**msg, "content": [{"type": "text", "text": msg["content"]}]} - if isinstance(msg.get("content"), str) - else msg - ) - for msg in prompt - ] - for prompt in prompts - ] - - if not self.use_vllm: - slice_raw_data[slice_idx] = (raw_examples, images, prompts, None) - continue - - # TODO: add self.tools support - # Workaround for a bug in transformers 5.3.0 where some processors (e.g. Qwen2.5-VL) crash on - # batched unpadded input (transformers#44514). - # Fixed in transformers 5.4.0 (transformers#44563). - needs_padding_workaround = Version("5.3.0") <= Version(transformers.__version__) < Version("5.4.0") - tokenized = self.processing_class.apply_chat_template( - conversation=prompts, - add_generation_prompt=True, - tokenize=True, - return_dict=True, - **({"padding": True} if needs_padding_workaround else {}), - ) - if needs_padding_workaround: - # Unpad input_ids: remove padding tokens using attention_mask to get per-sequence lists - prompt_ids_list = [ - [tok for tok, m in zip(ids, mask, strict=True) if m] - for ids, mask in zip(tokenized["input_ids"], tokenized["attention_mask"], strict=True) - ] - else: - prompt_ids_list = tokenized["input_ids"] - if prompt_max_length is not None: - # Do not truncate VLM prompts: a token-level slice can cut through the expanded image-token block, - # desyncing input_ids from the image features (garbage generations), and the buffered training - # examples are rebuilt from the untruncated rows, so the completion would be sampled under a - # different prompt than the one used for the loss. Fail fast instead. - for prompt_ids in prompt_ids_list: - if len(prompt_ids) > prompt_max_length: - raise ValueError( - f"A tokenized VLM prompt has {len(prompt_ids)} tokens (including expanded image " - f"tokens), exceeding the prompt budget of max_length - max_completion_length = " - f"{self.args.max_length} - {max_completion_length} = {prompt_max_length}. Increase `max_length` or set it to `None`." - ) - - slice_raw_data[slice_idx] = (raw_examples, images, prompts, prompt_ids_list) - - for i, example in enumerate(raw_examples): - all_prompt_ids.append(prompt_ids_list[i]) - all_images.append(images[i] if images is not None else None) - all_prompts.append(prompts[i]) - all_raw_examples.append(example) - local_slice_indices.append(slice_idx) - - if not self.use_vllm: - for slice_idx in on_policy_indices: - raw_examples, images, prompts, _ = slice_raw_data[slice_idx] - has_images = images is not None and any(img is not None for img in images) - pending_slice = {"_gold_vlm_on_policy_raw_examples": raw_examples} - if self._teacher_processor is not None: - pending_slice["_gold_vlm_raw_images"] = images if has_images else None - pending_slice["_gold_vlm_raw_prompts"] = prompts - self._buffered_inputs[slice_idx] = pending_slice - return - - all_prompts_text = self.processing_class.batch_decode(all_prompt_ids, skip_special_tokens=True) - if ( - self.state.global_step != self._last_vllm_sync_step - and self.state.global_step >= self._last_vllm_sync_step + self.vllm_sync_frequency - ): - self.vllm_generation.sync_weights() - self._last_vllm_sync_step = self.state.global_step - - if any(img is not None for img in all_images): - generate_images = all_images - else: - generate_images = None - _, completion_ids, _, _ = self.vllm_generation.generate( - prompts=all_prompt_ids, - images=generate_images, - num_generations=self.num_generations, - ) - - all_completion_texts = [] - for comp_ids in completion_ids: - if len(comp_ids) > max_completion_length: - comp_ids = comp_ids[:max_completion_length] - all_completion_texts.append( - self.processing_class.decode( - comp_ids, - skip_special_tokens=False, - clean_up_tokenization_spaces=False, - ) - ) - - # Redistribute completions to slices. The RepeatSampler has already duplicated examples - # `num_generations` times, so completions align 1:1 with the sampled input entries. - slice_completions = {idx: [] for idx in on_policy_indices} - slice_raw = {idx: [] for idx in on_policy_indices} - slice_images = {idx: [] for idx in on_policy_indices} - slice_prompts = {idx: [] for idx in on_policy_indices} - slice_prompts_text = {idx: [] for idx in on_policy_indices} - - for i, slice_idx in enumerate(local_slice_indices): - slice_completions[slice_idx].append(all_completion_texts[i]) - slice_raw[slice_idx].append(all_raw_examples[i]) - slice_images[slice_idx].append(all_images[i]) - slice_prompts[slice_idx].append(all_prompts[i]) - slice_prompts_text[slice_idx].append(all_prompts_text[i]) - - for slice_idx in on_policy_indices: - completion_texts = slice_completions[slice_idx] - raw_for_slice = slice_raw[slice_idx] - images_for_slice = slice_images[slice_idx] - prompts_for_slice = slice_prompts[slice_idx] - - # Build synthetic examples: original prompt + generated completion - synthetic_examples = [] - for i, example in enumerate(raw_for_slice): - synthetic = dict(example) - # Wrap as content blocks so VLM chat templates (e.g. SmolVLM) that index - # `message.content[0]` can render the synthetic assistant turn. - synthetic["completion"] = [ - { - "role": "assistant", - "content": [{"type": "text", "text": completion_texts[i]}], - } - ] - synthetic_examples.append(synthetic) - - has_images = any(img is not None for img in images_for_slice) - pending_slice = { - "_gold_vlm_lazy_examples": synthetic_examples, - } - if self._teacher_processor is not None: - pending_slice["_gold_vlm_raw_images"] = images_for_slice if has_images else None - pending_slice["_gold_vlm_raw_prompts"] = prompts_for_slice - self._buffered_inputs[slice_idx] = pending_slice - self._buffered_text_logs[slice_idx] = ( - slice_prompts_text[slice_idx], - completion_texts, - ) - - def _process_completions_to_buffer( - self, - slices: list[dict[str, torch.Tensor | Any]], - on_policy_indices: list[int], - local_slice_indices: list[int], - completion_ids: list, - prompt_ids_list: list[list[int]], - prompts_text: list[str], - max_completion_length: int, - ): - """ - Process vLLM completions and update buffered inputs for on-policy slices. - """ - device = self.accelerator.device - pad_token_id = self.processing_class.pad_token_id if self.processing_class.pad_token_id is not None else 0 - - slice_completions = {idx: [] for idx in on_policy_indices} - slice_prompt_ids = {idx: [] for idx in on_policy_indices} - slice_prompts = {idx: [] for idx in on_policy_indices} - - for i, slice_idx in enumerate(local_slice_indices): - slice_completions[slice_idx].append(completion_ids[i]) - slice_prompt_ids[slice_idx].append(prompt_ids_list[i]) - slice_prompts[slice_idx].append(prompts_text[i]) - - for slice_idx in on_policy_indices: - slice_inputs = slices[slice_idx] - completion_ids_for_slice = slice_completions[slice_idx] - prompt_ids_for_slice = slice_prompt_ids[slice_idx] - prompt_txts = slice_prompts[slice_idx] - - # Prompts are already budgeted to max_length - max_new_tokens before generation by the - # callers, so no truncation happens here; just build tensors and left-pad. - prompt_tensors = [] - prompt_attention_masks = [] - for prompt_ids in prompt_ids_for_slice: - prompt_tensor = torch.tensor(prompt_ids, device=device, dtype=torch.long) - prompt_tensors.append(prompt_tensor) - prompt_attention_masks.append(torch.ones(len(prompt_ids), device=device, dtype=torch.long)) - - prompt_ids = pad(prompt_tensors, padding_side="left", padding_value=pad_token_id) - prompt_attention_mask = pad(prompt_attention_masks, padding_side="left", padding_value=0) - - # Decode the prompt (already budgeted by the callers) so the teacher conditions on the same - # context the student saw. `clean_up_tokenization_spaces=False` matches the completion decode - # below so byte counts stay aligned. - prompt_txts_with_special = self.processing_class.batch_decode( - [ids.tolist() for ids in prompt_tensors], - skip_special_tokens=False, - clean_up_tokenization_spaces=False, - ) - - completion_ids_tensors = [torch.tensor(ids, device=device) for ids in completion_ids_for_slice] - completion_ids_for_text: list[list[int]] = [] - padded_completion_ids_list = [] - completion_attention_masks = [] - for completion_tensor in completion_ids_tensors: - if len(completion_tensor) > max_completion_length: - truncated_completion_tensor = completion_tensor[:max_completion_length] - padded_completion_ids_list.append(truncated_completion_tensor) - completion_ids_for_text.append(truncated_completion_tensor.tolist()) - completion_attention_masks.append( - torch.ones( - len(truncated_completion_tensor), - device=device, - dtype=torch.long, - ) - ) - elif len(completion_tensor) < max_completion_length: - padding_needed = max_completion_length - len(completion_tensor) - padded_tensor = torch.cat( - [ - completion_tensor, - torch.full( - (padding_needed,), - pad_token_id, - device=device, - dtype=completion_tensor.dtype, - ), - ] - ) - padded_completion_ids_list.append(padded_tensor) - completion_ids_for_text.append(completion_tensor.tolist()) - completion_attention_masks.append( - torch.cat( - [ - torch.ones( - len(completion_tensor), - device=device, - dtype=torch.long, - ), - torch.zeros(padding_needed, device=device, dtype=torch.long), - ] - ) - ) - else: - padded_completion_ids_list.append(completion_tensor) - completion_ids_for_text.append(completion_tensor.tolist()) - completion_attention_masks.append( - torch.ones(len(completion_tensor), device=device, dtype=torch.long) - ) - - completion_ids_padded = torch.stack(padded_completion_ids_list) - completion_attention_mask = torch.stack(completion_attention_masks) - - new_input_ids = torch.cat([prompt_ids, completion_ids_padded], dim=1) - new_attention_mask = torch.cat([prompt_attention_mask, completion_attention_mask], dim=1) - prompt_lengths = torch.full((prompt_ids.shape[0],), prompt_ids.shape[1], device=device) - new_attention_mask, new_labels = self._build_sequence_batch( - new_input_ids, - prompt_lengths, - attention_mask=new_attention_mask, - ) - - completion_texts = self.processing_class.batch_decode( - completion_ids_for_text, - skip_special_tokens=False, - clean_up_tokenization_spaces=False, - ) - - updated_slice = dict(slice_inputs) - updated_slice["input_ids"] = new_input_ids - updated_slice["attention_mask"] = new_attention_mask - updated_slice["labels"] = new_labels - # Update prompts to match the new padding width so prompt_length stays consistent - updated_slice["prompts"] = prompt_ids - updated_slice["prompt_attention_mask"] = prompt_attention_mask - updated_slice["original_prompt_text"] = prompt_txts_with_special - updated_slice["original_completion_text"] = completion_texts - self._maybe_add_completion_byte_offsets(updated_slice) - - self._buffered_inputs[slice_idx] = updated_slice - self._buffered_text_logs[slice_idx] = (prompt_txts, completion_texts) - - def _prepare_dataset( - self, - dataset: Dataset | IterableDataset, - processing_class: (PreTrainedTokenizerBase | BaseImageProcessor | FeatureExtractionMixin | ProcessorMixin), - args, - packing: bool, - formatting_func: Callable[[dict], str] | None, - dataset_name: str, - ) -> Dataset | IterableDataset: - """Preserve original text fields for ULD when needed.""" - # For VLM datasets, skip dataset preparation entirely — the VLM collator handles tokenization - # and image processing on the fly, similar to how SFTTrainer skips prep for vision datasets. - if self._is_vision_dataset: - return dataset - - column_names = list(next(iter(dataset)).keys()) - is_processed = "input_ids" in column_names - - if packing and self.use_uld_loss and self.teacher_tokenizer is not None: - raise ValueError( - "Packing is not supported with cross-tokenizer ULD because byte-offset alignment is defined per " - "prompt/completion example." - ) - - if not is_processed or (self.use_uld_loss and self.teacher_tokenizer is not None): - return self._prepare_dataset_with_original_text( - dataset, processing_class, args, packing, formatting_func, dataset_name - ) - - return super()._prepare_dataset(dataset, processing_class, args, packing, formatting_func, dataset_name) - - def _prepare_dataset_with_original_text( - self, - dataset: Dataset | IterableDataset, - processing_class: (PreTrainedTokenizerBase | BaseImageProcessor | FeatureExtractionMixin | ProcessorMixin), - args, - packing: bool, - formatting_func: Callable[[dict], str] | None, - dataset_name: str, - ) -> Dataset | IterableDataset: - """ - Prepare dataset while preserving original text for cross-tokenizer distillation. - """ - # Build the kwargs for the `map` function - map_kwargs = {} - if isinstance(dataset, Dataset): # IterableDataset does not support num_proc - map_kwargs["num_proc"] = args.dataset_num_proc - - with global_then_local_main_first(): - # Apply the formatting function if any - if formatting_func is not None: - if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` - map_kwargs["desc"] = f"Applying formatting function to {dataset_name} dataset" - - def _func(example): - return {"text": formatting_func(example)} - - dataset = dataset.map(_func, batched=False, **map_kwargs) - - # Convert the dataset to ChatML if needed - if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` - map_kwargs["desc"] = f"Converting {dataset_name} dataset to ChatML" - column_names = next(iter(dataset)).keys() - dataset = dataset.map( - maybe_convert_to_chatml, - remove_columns=("conversations" if "conversations" in column_names else None), - **map_kwargs, - ) - - # Tokenize the dataset while preserving original text - if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` - map_kwargs["desc"] = f"Tokenizing {dataset_name} dataset (preserving original text)" - - def tokenize_with_original_text( - example, processing_class, dataset_text_field, max_length, use_extended_uld - ): - """Emit input_ids, attention_mask, byte_offsets, completion_mask, and the original prompt/completion - text. - """ - result = {} - - if "prompt" in example: # prompt-completion case - if is_conversational(example): - prompt_text = processing_class.apply_chat_template( - example["prompt"], - add_generation_prompt=True, - tokenize=False, - **example.get("chat_template_kwargs", {}), - ) - full_text = processing_class.apply_chat_template( - example["prompt"] + example["completion"], - tokenize=False, - **example.get("chat_template_kwargs", {}), - ) - prompt_text = "".join( - x - for x, _ in takewhile( - lambda x: x[0] == x[1], - zip(prompt_text, full_text, strict=False), - ) - ) - completion_text = full_text[len(prompt_text) :] - else: - prompt_text = example["prompt"] - completion_text = example["completion"] - full_text = prompt_text + completion_text - result["original_prompt_text"] = prompt_text - result["original_completion_text"] = completion_text - elif is_conversational(example): - messages = example["messages"] - assistant_indices = [idx for idx, msg in enumerate(messages) if msg["role"] == "assistant"] - if assistant_indices: - completion_idx = assistant_indices[-1] - prompt_messages = messages[:completion_idx] - full_messages = messages[: completion_idx + 1] - if prompt_messages: - prompt_text = processing_class.apply_chat_template( - prompt_messages, - add_generation_prompt=True, - tokenize=False, - **example.get("chat_template_kwargs", {}), - ) - else: - prompt_text = "" - full_text = processing_class.apply_chat_template( - full_messages, - add_generation_prompt=False, - tokenize=False, - **example.get("chat_template_kwargs", {}), - ) - prompt_text = "".join( - x - for x, _ in takewhile( - lambda x: x[0] == x[1], - zip(prompt_text, full_text, strict=False), - ) - ) - completion_text = full_text[len(prompt_text) :] - result["original_prompt_text"] = prompt_text - result["original_completion_text"] = completion_text - else: - full_text = processing_class.apply_chat_template( - messages, - tokenize=False, - **example.get("chat_template_kwargs", {}), - ) - prompt_text = "" - result["original_prompt_text"] = "" - result["original_completion_text"] = full_text - else: - text = example.get(dataset_text_field, example.get("text", "")) - prompt_text = "" - completion_text = text - full_text = text - result["original_prompt_text"] = "" - result["original_completion_text"] = text - - if use_extended_uld: - # Single backend call: ids and char-derived byte offsets from the same encoding, - # so input_ids[i] is described by full_offs[i] without any boundary slop. - backend = processing_class.backend_tokenizer - [(input_ids, full_offs)] = encode_with_byte_offsets(backend, [full_text], add_special_tokens=False) - prompt_byte_len = len(prompt_text.encode("utf-8")) - # A token straddling the boundary is pulled into the completion here: extended ULD - # re-splits it against the teacher's own byte offsets via `_align_by_byte_offsets`. - completion_start = next( - (idx for idx, (_, e) in enumerate(full_offs) if e > prompt_byte_len), - len(input_ids), - ) - # Completion-relative: prompt positions zeroed, completion offsets shifted to - # the assistant content's first byte (matches build_teacher_inputs_from_texts). - byte_offsets = [(0, 0)] * completion_start + [ - (max(0, s - prompt_byte_len), e - prompt_byte_len) for s, e in full_offs[completion_start:] - ] - else: - # Works with slow tokenizers too (e.g. SentencePiece): no `backend_tokenizer` needed. - encoding = processing_class(full_text, add_special_tokens=False, return_offsets_mapping=True) - input_ids = encoding["input_ids"] - # Same boundary rule as extended ULD: a leading space at the seam is normally attached to - # the following (completion) token by the tokenizer, so matching on `start` instead would - # drop that token's content from the completion entirely. - completion_start = next( - (idx for idx, (_, end) in enumerate(encoding["offset_mapping"]) if end > len(prompt_text)), - len(input_ids), - ) - byte_offsets = [(0, 0)] * len(input_ids) - - append_eos = ( - not is_conversational(example) - and processing_class.eos_token_id is not None - and (not input_ids or input_ids[-1] != processing_class.eos_token_id) - ) - if append_eos: - input_ids.append(processing_class.eos_token_id) - if use_extended_uld: - completion_byte_len = len(completion_text.encode("utf-8")) - byte_offsets.append((completion_byte_len, completion_byte_len)) - else: - byte_offsets.append((0, 0)) - - # Keep the last `max_length` tokens (the completion end). `completion_mask` tracks the - # boundary so it survives truncation without re-tokenizing the prompt. - if max_length is not None and len(input_ids) > max_length: - drop = len(input_ids) - max_length - input_ids = input_ids[drop:] - byte_offsets = byte_offsets[drop:] - completion_start = max(0, completion_start - drop) - # If truncation ate into the completion, rebase the kept completion offsets so they're - # relative to the new (truncated) `original_completion_text` the teacher will re-encode. - if completion_start < len(byte_offsets): - base = byte_offsets[completion_start][0] - if base > 0: - byte_offsets = byte_offsets[:completion_start] + [ - (s - base, e - base) for s, e in byte_offsets[completion_start:] - ] - # Resync the strings the teacher will re-encode with the ids the student kept. - decode = partial( - processing_class.decode, - skip_special_tokens=False, - clean_up_tokenization_spaces=False, - ) - result["original_prompt_text"] = decode(input_ids[:completion_start]) - completion_ids = input_ids[completion_start:] - if append_eos: - completion_ids = completion_ids[:-1] - result["original_completion_text"] = decode(completion_ids) - - result["input_ids"] = input_ids - result["attention_mask"] = [1] * len(input_ids) - result["byte_offsets"] = byte_offsets - result["completion_mask"] = [0] * completion_start + [1] * (len(input_ids) - completion_start) - return result - - dataset = dataset.map( - tokenize_with_original_text, - fn_kwargs={ - "processing_class": processing_class, - "dataset_text_field": args.dataset_text_field, - "max_length": args.max_length, - "use_extended_uld": args.use_extended_uld, - }, - **map_kwargs, - ) - - # Pack if requested. Truncation already happened in `tokenize_with_original_text`, keeping - # the completion end — so the generic front-truncating `truncate_dataset` is intentionally - # not applied here (it would drop the completion). - if packing: - if args.max_length is None: - raise ValueError("When packing is enabled, `max_length` can't be `None`.") - if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` - map_kwargs["desc"] = f"Packing {dataset_name} dataset" - - columns_to_keep = [ - "input_ids", - "original_prompt_text", - "original_completion_text", - ] - existing_columns = set(dataset.column_names) - columns_to_select = [col for col in columns_to_keep if col in existing_columns] - - dataset = dataset.select_columns(columns_to_select) - dataset = pack_dataset(dataset, args.max_length, args.packing_strategy, map_kwargs) - - return dataset - - @staticmethod - def generalized_jsd_loss( - student_logits, - teacher_logits, - labels=None, - beta=0.5, - temperature=1.0, - reduction="batchmean", - logits_are_probs=False, - num_items_in_batch=None, - ): - """ - Compute the generalized Jensen-Shannon Divergence loss for knowledge distillation using F.kl_div. See Eq. (1) - of https://huggingface.co/papers/2306.13649 for the definition. - - Args: - student_logits: - Tensor of shape (batch_size, sequence_length, vocab_size) - teacher_logits: - Tensor of shape (batch_size, sequence_length, vocab_size) - labels: - Tensor of shape (batch_size, sequence_length) with -100 for padding tokens to ignore when computing - loss - beta: - Interpolation coefficient between 0 and 1 (default: 0.5) - temperature: - Softmax temperature (default: 1.0) - reduction: - Specifies the reduction to apply to the output (default: 'batchmean') - - Returns: - loss: Scalar tensor with the generalized JSD loss - """ - - if logits_are_probs: - student_log_probs = torch.log(student_logits.clamp_min(1e-8)) - teacher_log_probs = torch.log(teacher_logits.clamp_min(1e-8)) - else: - # Apply temperature scaling to logits before computing probabilities - student_logits = student_logits / temperature - teacher_logits = teacher_logits / temperature - # Compute log probabilities for student and probabilities for teacher - student_log_probs = F.log_softmax(student_logits, dim=-1) - teacher_log_probs = F.log_softmax(teacher_logits, dim=-1) - - if beta == 0: - jsd = F.kl_div(student_log_probs, teacher_log_probs, reduction="none", log_target=True) - elif beta == 1: - jsd = F.kl_div(teacher_log_probs, student_log_probs, reduction="none", log_target=True) - else: - # Compute the log of the mixture distribution - # log(a + b) = log(exp(log(a)) + exp(log(b))) -> for mixture - beta = torch.tensor(beta, dtype=student_log_probs.dtype, device=student_log_probs.device) - mixture_log_probs = torch.logsumexp( - torch.stack( - [ - student_log_probs + torch.log1p(-beta), - teacher_log_probs + torch.log(beta), - ] - ), - dim=0, - ) - - # Compute KL divergences using F.kl_div - # PyTorch differs from the standard mathematical definition, so the order of the probability distributions is swapped compared to that defined in the paper. - kl_teacher = F.kl_div(mixture_log_probs, teacher_log_probs, reduction="none", log_target=True) - kl_student = F.kl_div(mixture_log_probs, student_log_probs, reduction="none", log_target=True) - - # Compute the Generalized Jensen-Shannon Divergence - jsd = beta * kl_teacher + (1 - beta) * kl_student - - # Masking - if labels is not None: - mask = labels != -100 - jsd = jsd[mask] - - # Apply reduction - if num_items_in_batch is not None: - # Normalize by the global number of valid tokens for gradient-accumulation-correct loss (see issue #4719). - jsd_sum = jsd.sum() - if isinstance(num_items_in_batch, torch.Tensor): - num_items_in_batch = num_items_in_batch.to(jsd_sum.device) - return jsd_sum / num_items_in_batch - if reduction == "batchmean": - # clamp_min(1) avoids 0/0 -> nan when a sample has no unmasked positions - # (e.g. completion fully truncated). jsd[mask] is empty -> jsd.sum() == 0, - # so 0/1 == 0 with a valid grad path. - denom = mask.sum().clamp_min(1) if labels is not None else max(jsd.size(0), 1) - return jsd.sum() / denom - elif reduction == "sum": - return jsd.sum() - elif reduction == "mean": - return jsd.mean() - else: - return jsd - - def _build_teacher_vlm_inputs( - self, completion_texts: list[str], raw_images: list, raw_prompts: list - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]: - """Build image-aware teacher inputs for cross-architecture VLM ULD distillation. - - The teacher prompt is rendered through the teacher's own processor so image placeholders and pixel tensors - match the teacher model, while completion tokens carry byte offsets relative to the original - ``completion_texts`` — the same coordinate system the student uses — so cross-tokenizer byte-offset alignment - stays valid. Mirrors ``build_teacher_inputs_from_texts`` but injects the teacher's multimodal prompt and - returns the teacher's forward kwargs (``pixel_values``, ...). - - Returns ``(input_ids, labels, attention_mask, byte_offsets, forward_kwargs)``. - """ - pad_token_id = self.teacher_tokenizer.pad_token_id - eos_token_id = self.teacher_tokenizer.eos_token_id - - teacher_prompt_texts = self._teacher_processor.apply_chat_template( - raw_prompts, tokenize=False, add_generation_prompt=True - ) - teacher_prompt_processed = self._teacher_processor( - images=raw_images, - text=teacher_prompt_texts, - padding=True, - return_tensors="pt", - ) - if self.uld_loss_fn.use_extended_uld: - # Only the extended path needs the fast tokenizer's byte offsets; positional ULD works with slow - # (e.g. SentencePiece) tokenizers that have no `backend_tokenizer`. - backend = self.teacher_tokenizer.backend_tokenizer - completion_encs = encode_with_byte_offsets(backend, completion_texts, add_special_tokens=False) - else: - completion_ids = self.teacher_tokenizer(completion_texts, add_special_tokens=False)["input_ids"] - completion_encs = [(ids, [(0, 0)] * len(ids)) for ids in completion_ids] - - sequences: list[torch.Tensor] = [] - attention_masks: list[torch.Tensor] = [] - labels_list: list[torch.Tensor] = [] - offsets_list: list[list[tuple[int, int]]] = [] - sequence_kwargs = defaultdict(list) - - for row, ((enc_ids, enc_offs), completion_text) in enumerate( - zip(completion_encs, completion_texts, strict=True) - ): - prompt_mask = teacher_prompt_processed["attention_mask"][row].bool() - prompt_ids = teacher_prompt_processed["input_ids"][row][prompt_mask].tolist() - # Remove trailing EOS from prompt so completions can extend cleanly - if eos_token_id is not None and prompt_ids and prompt_ids[-1] == eos_token_id: - prompt_ids = prompt_ids[:-1] - - completion_ids = list(enc_ids) - completion_offs = list(enc_offs) - content_len = len(completion_text.encode("utf-8")) - - sequence = list(prompt_ids) + completion_ids - offsets = [(0, 0)] * len(prompt_ids) + completion_offs - if eos_token_id is not None: - sequence.append(eos_token_id) - offsets.append((content_len, content_len) if self.uld_loss_fn.use_extended_uld else (0, 0)) - - seq_tensor = torch.tensor(sequence, dtype=torch.long) - sequences.append(seq_tensor) - attention_masks.append(torch.ones_like(seq_tensor)) - offsets_list.append(offsets) - - labels = seq_tensor.clone() - labels[: len(prompt_ids)] = -100 - labels_list.append(labels) - - # Sequence-aligned multimodal keys (e.g. token_type_ids for Gemma): keep the prompt values - # and zero-fill the completion span. - for key in self._SEQUENCE_KEYS: - if key in teacher_prompt_processed: - prompt_values = teacher_prompt_processed[key][row][prompt_mask][: len(prompt_ids)] - completion_values = torch.zeros(len(sequence) - len(prompt_ids), dtype=prompt_values.dtype) - sequence_kwargs[key].append(torch.cat((prompt_values, completion_values))) - - teacher_input_ids = pad( - sequences, - padding_side="right", - padding_value=pad_token_id if pad_token_id is not None else 0, - ) - teacher_attention_mask = pad(attention_masks, padding_side="right", padding_value=0).bool() - teacher_labels = pad(labels_list, padding_side="right", padding_value=-100) - - target_len = teacher_input_ids.size(1) - teacher_byte_offsets = torch.stack( - [pad_byte_offsets(offs, target_len, padding_side="right") for offs in offsets_list], - dim=0, - ) - - # Multimodal forward kwargs from the teacher processor (pixel_values, image_grid_thw, ...). - forward_kwargs = { - k: v.to(self.accelerator.device) - for k, v in self._get_model_forward_kwargs(teacher_prompt_processed, exclude=self._SEQUENCE_KEYS).items() - } - for key, values in sequence_kwargs.items(): - forward_kwargs[key] = pad(values, padding_side="right", padding_value=0).to(self.accelerator.device) - - return ( - teacher_input_ids, - teacher_labels, - teacher_attention_mask, - teacher_byte_offsets, - forward_kwargs, - ) - - def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): - # Extract multimodal fields (pixel_values, image_grid_thw, ...) for student forward passes. - # Standard JSD reuses these for the teacher (same-family VLM); cross-tokenizer ULD rebuilds - # teacher inputs separately below. - student_forward_kwargs = self._get_model_forward_kwargs(inputs) - if self.use_uld_loss and self.teacher_tokenizer is not None: - # Both DataCollatorForChatML and the on-policy generation path attach these - # fields, so cross-tokenizer ULD never has to round-trip through batch_decode. - prompt_texts = inputs["original_prompt_text"] - completion_texts = inputs["original_completion_text"] - - if self._teacher_processor is not None: - # VLM teacher: render the prompt through the teacher's own processor so image - # placeholders and pixel tensors match the teacher model. - if "_raw_images" not in inputs or "_raw_prompts" not in inputs: - raise ValueError( - "VLM ULD distillation requires `_raw_images` and `_raw_prompts` in the batch so teacher " - "inputs can be rendered with the teacher processor. Use the default GOLD VLM data collator " - "or ensure your custom collator preserves these fields." - ) - ( - teacher_input_ids, - teacher_labels, - teacher_attention_mask, - teacher_completion_byte_offsets, - teacher_forward_kwargs, - ) = self._build_teacher_vlm_inputs(completion_texts, inputs["_raw_images"], inputs["_raw_prompts"]) - else: - teacher_forward_kwargs = {} - ( - teacher_input_ids, - teacher_labels, - teacher_attention_mask, - teacher_completion_byte_offsets, - ) = build_teacher_inputs_from_texts( - self.teacher_tokenizer, - prompt_texts, - completion_texts, - use_extended_uld=self.uld_loss_fn.use_extended_uld, - ) - - teacher_input_ids = teacher_input_ids.to(self.accelerator.device) - teacher_labels = teacher_labels.to(self.accelerator.device) - teacher_attention_mask = teacher_attention_mask.to(self.accelerator.device) - - outputs_student = model( - input_ids=inputs["input_ids"], - attention_mask=inputs["attention_mask"], - use_cache=False, - **student_forward_kwargs, - ) - - self.teacher_model.eval() - with torch.no_grad(): - outputs_teacher = self.teacher_model( - input_ids=teacher_input_ids, - attention_mask=teacher_attention_mask, - **teacher_forward_kwargs, - ) - else: - outputs_student = model( - input_ids=inputs["input_ids"], - attention_mask=inputs["attention_mask"], - **student_forward_kwargs, - ) - - self.teacher_model.eval() - with torch.no_grad(): - outputs_teacher = self.teacher_model( - input_ids=inputs["input_ids"], - attention_mask=inputs["attention_mask"], - **student_forward_kwargs, - ) - - # Standard causal shift: logits at position i predict the token at i + 1. The `labels != -100` mask - # inside `generalized_jsd_loss` already excludes prompt (and padding) positions, so we do not slice by - # prompt length. Slicing by `inputs["prompts"].shape[1]` (the batch-max prompt width) would drop real - # completion tokens for samples whose prompt is shorter than the batch maximum, since `labels` is - # padded to the full-sequence width independently of `prompts`. - shifted_student_logits = outputs_student.logits[:, :-1, :] - shifted_teacher_logits = outputs_teacher.logits[:, :-1, :] - shifted_labels = inputs["labels"][:, 1:] - loss = self.generalized_jsd_loss( - student_logits=shifted_student_logits, - teacher_logits=shifted_teacher_logits, - labels=shifted_labels, - beta=self.beta, - temperature=self.temperature, - num_items_in_batch=num_items_in_batch, - ) - - if self.use_uld_loss and self.teacher_tokenizer is not None: - student_labels = inputs["labels"] - - student_byte_offsets = inputs.get("byte_offsets") - if self.uld_loss_fn.use_extended_uld and student_byte_offsets is None: - raise ValueError("Input batches must include `byte_offsets` when `use_extended_uld=True`.") - - loss = self.uld_loss_fn( - student_logits=outputs_student.logits, - teacher_logits=outputs_teacher.logits, - student_labels=student_labels, - teacher_labels=teacher_labels, - student_input_ids=inputs["input_ids"], - teacher_input_ids=teacher_input_ids, - student_byte_offsets=student_byte_offsets, - teacher_byte_offsets=teacher_completion_byte_offsets, - ) - - if hasattr(self.uld_loss_fn, "last_matched_loss") and hasattr(self.uld_loss_fn, "last_unmatched_loss"): - ga = max(1, int(self.args.gradient_accumulation_steps)) - step_eq = 1.0 / ga - matched_val = ( - self.uld_loss_fn.last_matched_loss.item() - if self.uld_loss_fn.last_matched_loss is not None - else 0.0 - ) - unmatched_val = ( - self.uld_loss_fn.last_unmatched_loss.item() - if self.uld_loss_fn.last_unmatched_loss is not None - else 0.0 - ) - - self._matched_sum += matched_val - self._unmatched_sum += unmatched_val - self._matched_step_eq += step_eq - self._unmatched_step_eq += step_eq - - empty_cache() - - return (loss, outputs_student) if return_outputs else loss - - def generate_on_policy_outputs(self, model, inputs, generation_config): - # Generate output with respect to the prompt only - # Drop sequence-aligned multimodal keys (token_type_ids, mm_token_type_ids): generate recomputes them from - # `input_ids` as it extends the sequence, and some models (e.g. Qwen3-VL) reject them as unknown - # `generate` kwargs. - generate_kwargs = self._get_model_forward_kwargs(inputs, exclude=self._SEQUENCE_KEYS) - generated_outputs = model.generate( - input_ids=inputs["prompts"], - attention_mask=inputs["prompt_attention_mask"], - generation_config=generation_config, - return_dict_in_generate=True, - **generate_kwargs, - ) - # Get the generated token IDs - generated_tokens = generated_outputs.sequences - - batch_size = generated_tokens.size(0) - device = generated_tokens.device - - prompt_mask = inputs["prompt_attention_mask"] - - # model.generate() returns full sequences (prompt + completion), so completions start - # after the full padded prompt width. - prompt_length = inputs["prompts"].shape[1] - prompt_lengths = torch.full((batch_size,), prompt_length, dtype=torch.long, device=device) - completion_ids = generated_tokens[:, prompt_length:] - - # Mask everything after the first EOS token. eos_token_id can be a list of stop tokens (Llama 3), - # so match with torch.isin; with no eos id nothing stops early, so keep the whole completion. - if generation_config.eos_token_id is None: - completion_mask = torch.ones_like(completion_ids) - else: - is_eos = torch.isin(completion_ids, torch.tensor(generation_config.eos_token_id, device=device)) - eos_idx = torch.full((is_eos.size(0),), is_eos.size(1), dtype=torch.long, device=device) - eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)] - sequence_indices = torch.arange(is_eos.size(1), device=device).expand(is_eos.size(0), -1) - completion_mask = (sequence_indices <= eos_idx.unsqueeze(1)).int() - - # Pad ids can appear inside real prompt text, so use the prompt mask instead of matching ids - new_attention_mask = torch.ones_like(generated_tokens) - new_attention_mask[:, prompt_length:] = completion_mask - new_attention_mask[:, :prompt_length] = prompt_mask - - new_input_ids = generated_tokens - new_attention_mask, new_labels = self._build_sequence_batch( - new_input_ids, prompt_lengths, attention_mask=new_attention_mask - ) - - prompt_texts = [] - completion_texts = [] - for idx in range(batch_size): - length = int(prompt_lengths[idx].item()) - prompt_tokens = inputs["prompts"][idx] - prompt_tokens = prompt_tokens[prompt_mask[idx].bool()] - prompt_texts.append( - self.processing_class.decode( - prompt_tokens.tolist(), - skip_special_tokens=False, - clean_up_tokenization_spaces=False, - ) - ) - completion_tokens = new_input_ids[idx, length:] - completion_tokens = completion_tokens[new_labels[idx, length:] != -100] - completion_texts.append( - self.processing_class.decode( - completion_tokens.tolist(), - skip_special_tokens=False, - clean_up_tokenization_spaces=False, - ) - ) - - return ( - new_input_ids, - new_attention_mask, - new_labels, - prompt_texts, - completion_texts, - ) - - # During eval, Trainer calls prediction_step. The inherited SFT prediction_step indexes the raw inputs before - # collation, which breaks the VLM identity-collator path (inputs is a list of raw examples). We override it to - # collate via _prepare_inputs and force compute_loss, evaluating the off-policy distillation loss over the - # ground-truth completion (no generation). - def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys: list[str] | None = None): - inputs = self._prepare_inputs(inputs) - with torch.no_grad(): - with self.compute_loss_context_manager(): - loss = self.compute_loss(model, inputs) - loss = loss.mean().detach() - return loss, None, None - - @profiling_decorator - def training_step( - self, - model: nn.Module, - inputs: dict[str, torch.Tensor | Any], - num_items_in_batch: int | None = None, - ) -> torch.Tensor: - """ - Perform a training step for the General Online Logit Distillation (GOLD) model. - - This method implements the on-policy learning approach described in the GOLD blog post. With probability - `self.lmbda`, it generates new responses using the student model, which are then used for training instead of - the offline original inputs. - """ - buffer_steps = self.args.gradient_accumulation_steps - - loss = super().training_step(model, inputs, num_items_in_batch) - - slice_idx = (self._step - 1) % buffer_steps - - on_policy = False - if self._buffered_on_policy is not None and slice_idx < len(self._buffered_on_policy): - on_policy = self._buffered_on_policy[slice_idx] - - if on_policy and self._buffered_text_logs is not None and self._buffered_text_logs[slice_idx] is not None: - prompt_texts, completion_texts = self._buffered_text_logs[slice_idx] - self._textual_logs["prompt"].extend(gather_object(prompt_texts)) - self._textual_logs["completion"].extend(gather_object(completion_texts)) - - loss_scalar = float(loss.detach()) - step_equiv = 1.0 / self.args.gradient_accumulation_steps - - if on_policy: - self._on_policy_loss_total += loss_scalar - self._on_policy_step_equiv += step_equiv - else: - self._off_policy_loss_total += loss_scalar - self._off_policy_step_equiv += step_equiv - return loss - - def log(self, logs: dict[str, float], start_time: float | None = None) -> None: - mode = "train" if self.model.training else "eval" - metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()} # average the metrics - - if mode == "train": - device = self.accelerator.device if hasattr(self.accelerator, "device") else torch.device("cpu") - vec = torch.tensor( - [ - self._on_policy_loss_total, - self._off_policy_loss_total, - self._on_policy_step_equiv, - self._off_policy_step_equiv, - self._matched_sum, - self._unmatched_sum, - self._matched_step_eq, - self._unmatched_step_eq, - ], - dtype=torch.float32, - device=device, - ) - - if ( - getattr(self.accelerator, "distributed_type", DistributedType.NO) != DistributedType.NO - and dist.is_available() - and dist.is_initialized() - ): - dist.all_reduce(vec, op=dist.ReduceOp.SUM) - - ( - on_sum, - off_sum, - on_eq, - off_eq, - matched_sum, - unmatched_sum, - matched_eq, - unmatched_eq, - ) = vec.tolist() - - if on_eq > 0: - logs["on_policy_loss"] = round(on_sum / on_eq, 4) - if off_eq > 0: - logs["off_policy_loss"] = round(off_sum / off_eq, 4) - - if matched_eq > 0: - logs["matched_loss"] = round(matched_sum / matched_eq, 4) - if unmatched_eq > 0: - logs["unmatched_loss"] = round(unmatched_sum / unmatched_eq, 4) - - self._on_policy_loss_total = self._off_policy_loss_total = 0.0 - self._on_policy_step_equiv = self._off_policy_step_equiv = 0.0 - self._matched_sum = self._unmatched_sum = 0.0 - self._matched_step_eq = self._unmatched_step_eq = 0.0 - - # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` - # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. - if mode == "eval": - metrics = {f"eval_{key}": val for key, val in metrics.items()} - - logs.update(metrics) - super().log(logs, start_time) - self._metrics[mode].clear() - - if ( - self.accelerator.is_main_process - and self.log_completions - and ((self.state.global_step % self.log_completion_steps) == 0) - ): - if is_rich_available(): - print_prompt_completions_sample_uld( - self._textual_logs["prompt"], - self._textual_logs["completion"], - self.state.global_step, - self.num_completions_to_print, - ) - - if self.args.report_to and "wandb" in self.args.report_to and wandb.run is not None: - import pandas as pd - - table = { - "step": [self.state.global_step] * len(self._textual_logs["prompt"]), - "prompt": self._textual_logs["prompt"], - "completion": self._textual_logs["completion"], - } - df = pd.DataFrame(table) - if self.wandb_log_unique_prompts: - df = df.drop_duplicates(subset=["prompt"]) - if self.num_completions_to_print and len(df) > 0: - df = df.sample(n=self.num_completions_to_print, random_state=42) - wandb.log({"completions": wandb.Table(dataframe=df)}) -class UnslothGOLDTrainer(_UnslothGOLDTrainer): - """ - - """ - def __init__( - self, - model = None, - teacher_model = None, - args = None, - data_collator = None, - train_dataset = None, - eval_dataset = None, - processing_class = None, - compute_metrics = None, - callbacks = None, - preprocess_logits_for_metrics = None, - peft_config = None, - **kwargs - ): - if args is None: args = UnslothGOLDConfig() - _unsloth_precision_requested = tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) - use_bf16 = getattr(args, 'bf16', False) - if type(use_bf16) is not bool: use_bf16 = False - use_fp16 = getattr(args, 'fp16', False) - if type(use_fp16) is not bool: use_fp16 = False - force_float32 = False - try: - from unsloth_zoo.device_type import device_is_bf16_supported as _bf16_supported - except Exception: - _bf16_supported = torch.cuda.is_bf16_supported - try: - from unsloth_zoo.device_type import DEVICE_TYPE as _unsloth_device_type - except Exception: - _unsloth_device_type = None - full_finetuning = getattr(model, '_unsloth_full_finetuning', None) - if full_finetuning is None: full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1' - model_forced_float32 = getattr(model, '_unsloth_forced_float32', None) - if model_forced_float32 is None: model_forced_float32 = os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1' - if model_forced_float32 and not (full_finetuning and _bf16_supported()): - print('Unsloth: Switching to float32 training since model cannot work with float16') - force_float32 = True - mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') - dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None) - if dtype is None: dtype = model.get_input_embeddings().weight.dtype - from unsloth_zoo.utils import _get_dtype - dtype = _get_dtype(dtype) - float16 = dtype == torch.float16 - bfloat16 = dtype == torch.bfloat16 - float32 = dtype == torch.float32 - user_float32 = bool(getattr(model, '_unsloth_user_float32', False)) - if full_finetuning: - if bfloat16 and use_fp16: use_fp16 = False - if float16 and use_bf16: use_bf16 = False - if not force_float32 and bfloat16 and use_fp16 and _unsloth_device_type == 'xpu' and not _bf16_supported(): - print('Unsloth: The model was loaded in bfloat16 on purpose, so training uses bfloat16 mixed precision instead of float16.') - use_fp16 = False - use_bf16 = True - args.fp16 = False - args.bf16 = True - if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`') - if not force_float32 and (bfloat16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`') - if force_float32: - # Forced float32 training - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32' and float32 and user_float32 and not _bf16_supported(): - print('Unsloth: Model is in float32 and this GPU has no bfloat16 support, so training stays in float32. Pass fp16 = True to force float16 mixed precision instead.') - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32': - # Mixed precision training. bf16 only if the GPU supports it; V100/T4 use fp16. - use_bf16_amp = (not float16) and (_bf16_supported() or (bfloat16 and _unsloth_device_type == 'xpu')) - args.fp16 = not use_bf16_amp - args.bf16 = use_bf16_amp - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16_amp else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16_amp else 'fp16' - # args.mixed_precision is a new argument which needs to be set now - elif mixed_precision_dtype == 'bfloat16': - # Both False since bfloat16 full finetuning doesn't do any autocasting. - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif use_bf16 or use_fp16: - # transformers <5 exported this itself from fp16/bf16; 5.x dropped the write, so an - # explicit flag left it unset and GRPO readers defaulted to 'fp16', wrapping a - # bfloat16 model in a float16 autocast. See unslothai/unsloth#4891. - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16 else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16 else 'fp16' - - if getattr(args, 'eval_dataset', None) is not None and getattr(args, 'eval_strategy', 'no') == 'no': - args.eval_strategy = 'steps' - if getattr(args, 'eval_steps', None) is None: args.eval_steps = 0.1 - ga_steps = getattr(args, 'gradient_accumulation_steps', None) - if ga_steps is not None and ga_steps > 1: - from transformers import __version__ as transformers_version - if Version(transformers_version) <= Version('4.45.2'): - print('**** Unsloth: Please use our fixed gradient_accumulation_steps by updating transformers, TRL and Unsloth!\n' - '`pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth transformers trl unsloth_zoo`') - if getattr(args, 'eval_strategy', 'no') != 'no': - eval_bsz = getattr(args, 'per_device_eval_batch_size', 8) - eval_generations = getattr(args, 'num_generations_eval', None) or getattr(args, 'num_generations', None) or 1 - if eval_bsz == 8 and args.per_device_train_batch_size < eval_bsz and (args.per_device_train_batch_size * args.world_size) % eval_generations == 0: args.per_device_eval_batch_size = args.per_device_train_batch_size - if getattr(args, 'eval_accumulation_steps', None) is None and ga_steps is not None: args.eval_accumulation_steps = ga_steps - fp16_full_eval = getattr(args, 'fp16_full_eval', False) - if type(fp16_full_eval) is not bool: fp16_full_eval = False - bf16_full_eval = getattr(args, 'bf16_full_eval', False) - if type(bf16_full_eval) is not bool: bf16_full_eval = False - if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True - if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False - if force_float32: - args.bf16_full_eval = False - args.fp16_full_eval = False - elif os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16': - args.bf16_full_eval = True - args.fp16_full_eval = False - elif not bf16_full_eval and not fp16_full_eval: - args.bf16_full_eval = args.bf16 - args.fp16_full_eval = args.fp16 - if tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) != _unsloth_precision_requested: - from unsloth.models.loader_utils import sync_deepspeed_mixed_precision - sync_deepspeed_mixed_precision(args) - if model is not None: - from unsloth.models.loader_utils import apply_deepspeed_zero3_model_setup - apply_deepspeed_zero3_model_setup(model) - if model is not None: - _warnings_issued = getattr(model, 'warnings_issued', None) - if _warnings_issued is None: - model.warnings_issued = {} - elif not isinstance(_warnings_issued, dict): - try: - model.warnings_issued = dict(_warnings_issued) - except Exception: - model.warnings_issued = {} - if 'max_seq_length' not in locals() and not hasattr(args, 'max_seq_length'): - pass - else: - model_max_seq_length = getattr(model, 'max_seq_length', None) - args_max_seq_length = getattr(args, 'max_seq_length', None) - if args_max_seq_length is None and model_max_seq_length is not None: - max_seq_length = model.max_seq_length - if hasattr(args, 'max_seq_length'): args.max_seq_length = max_seq_length - elif args_max_seq_length is not None and model_max_seq_length is not None: - if args_max_seq_length > model_max_seq_length: - print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but ' - 'the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.') - args.max_seq_length = model_max_seq_length - if model is not None and hasattr(model, 'for_training'): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right' - if 'processing_class' in locals(): - if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right' - if hasattr(processing_class, 'tokenizer') and hasattr(processing_class.tokenizer, 'padding_side'): processing_class.tokenizer.padding_side = 'right' - __tokenizer = processing_class if 'processing_class' in locals() else tokenizer - from unsloth_zoo.vision_utils import UnslothVisionDataCollator - if not isinstance(data_collator, UnslothVisionDataCollator): - if isinstance(data_collator, DataCollatorForSeq2Seq) and 'labels' not in _unsloth_dataset_column_names(train_dataset): - data_collator = TransformersDataCollatorForLanguageModeling( - __tokenizer, - mlm = False, - mlm_probability = 0.0, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling) and 'labels' in _unsloth_dataset_column_names(train_dataset): - data_collator = DataCollatorForSeq2Seq( - __tokenizer, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - else: - if hasattr(args, 'remove_unused_columns'): args.remove_unused_columns = False - if hasattr(args, 'dataset_text_field'): args.dataset_text_field = '' - if hasattr(args, 'dataset_kwargs'): args.dataset_kwargs = {'skip_prepare_dataset': True} - if not isinstance(data_collator, UnslothVisionDataCollator): - if not hasattr(__tokenizer, 'pad') and hasattr(__tokenizer, 'tokenizer'): - if isinstance(data_collator, DataCollatorForSeq2Seq): - data_collator = DataCollatorForSeq2Seq( - __tokenizer.tokenizer, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling): - data_collator = TransformersDataCollatorForLanguageModeling( - __tokenizer.tokenizer, - mlm = False, - mlm_probability = 0.0, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - other_metrics = [] - - from unsloth_zoo.logging_utils import PatchRLStatistics - PatchRLStatistics('gold_trainer', other_metrics) - - # [TODO] Fix up DataParallel multiplying batch sizes - # [TODO] DDP works, but DP seems to not work? [TODO] - if getattr(args, "parallel_mode", None) == ParallelMode.NOT_DISTRIBUTED and args.n_gpu > 1: - if getattr(args, "_n_gpu", 1) != 1: - args._n_gpu = 1 - if "model" in locals(): - try: - from unsloth.models._gc_config_default import honour_gradient_checkpointing_off - honour_gradient_checkpointing_off(args, model) - except Exception: - pass - if "model" in locals() and hasattr(model, "for_training"): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if "model" in locals(): - try: - from unsloth.models._ddp_defaults import unsloth_default_ddp_find_unused_parameters - unsloth_default_ddp_find_unused_parameters(args, model) - except Exception: - pass - super().__init__( - model = model, - teacher_model = teacher_model, - args = args, - data_collator = data_collator, - train_dataset = train_dataset, - eval_dataset = eval_dataset, - processing_class = processing_class, - compute_metrics = compute_metrics, - callbacks = callbacks, - preprocess_logits_for_metrics = preprocess_logits_for_metrics, - peft_config = peft_config,**kwargs) - # A tensor parallel model (tp_native.tp_trainer_setup) stays on its training path: for_inference would - # route the decoder layers onto decode paths outside the tensor parallel regions. - if "model" in locals() and hasattr(model, "for_inference") and not getattr(self, "_unsloth_tensor_parallel", False): - model.for_inference() - if hasattr(self, 'neftune_hook_handle'): - self.neftune_hook_handle.remove() - if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle - if getattr(args, 'neftune_noise_alpha', None) is not None: - model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha - pass - if hasattr(self, 'accelerator'): - scaler = self.accelerator.scaler - current_model = model - while hasattr(current_model, 'model'): - current_model.accelerator_scaler = scaler - current_model = current_model.model - current_model.accelerator_scaler = scaler - pass - if hasattr(self, 'train'): - self.train = MethodType(prepare_for_training_mode(self.__class__.train), self) - pass - if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'): - _vllm_tok = self.llm.get_tokenizer() - _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None) - if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None: - _vllm_tok.chat_template = _pc.chat_template - pass - if getattr(self, 'aux_loss_enabled', False) and hasattr(getattr(self, 'model', None), 'config'): - _text_config = self.model.config - if hasattr(_text_config, 'get_text_config'): _text_config = _text_config.get_text_config() - _n_experts = [getattr(_text_config, _k) for _k in ('num_local_experts', 'num_experts', 'n_routed_experts', 'moe_num_experts') if isinstance(getattr(_text_config, _k, None), int)] - if _n_experts and all(_n <= 1 for _n in _n_experts): - self.aux_loss_enabled = False - _text_config.output_router_logits = False - pass - -pass diff --git a/unsloth_compiled_cache/UnslothGRPOTrainer.py b/unsloth_compiled_cache/UnslothGRPOTrainer.py deleted file mode 100644 index dfd6a059858f293f9ce47c2d7420d04770f29abd..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/UnslothGRPOTrainer.py +++ /dev/null @@ -1,9224 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from unsloth_zoo.temporary_patches.utils import torch_compile_with_fallback -from unsloth_zoo.temporary_patches.common import torch_compile -from unsloth_zoo.temporary_patches.common import _maybe_compile -import functools -import transformers -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from trl.trainer.grpo_trainer import (Any, AutoModelForSequenceClassification, AutoProcessor, AutoTokenizer, BaseTunerLayer, BitsAndBytesConfig, Callable, CommitScheduler, DataLoader, Dataset, DatasetCard, DatasetCardData, DatasetDict, DistributedBackend, EnvironmentFactory, GRPOConfig, GRPOTrainer, GenerationConfig, IterableDataset, IterableDatasetDict, LoraConfig, PREFIX_CHECKPOINT_DIR, Path, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, RepeatSampler, RewardFunc, RolloutFunc, Sampler, SyncRefModelCallback, TrainerCallback, VLLMGeneration, Version, _BaseTrainer, _SUPPORTS_RESPONSE_TEMPLATE, add_fused_lm_head, add_response_schema, apply_chat_template, asyncio, atexit, copy, create_model_from_path, create_repo, defaultdict, deque, disable_dropout_in_model, disable_gradient_checkpointing, gather, gather_object, get_callable_name, get_config_model_id, get_peft_model, get_training_chat_template, identity, inspect, is_async_callable, is_chat_template_prefix_preserving, is_conversational, is_jmespath_available, is_peft_available, is_peft_model, is_rich_available, json, logger, math, nanmax, nanmin, nanstd, nn, np, os, pad, parse_response, pd, peft, pkg_resources, prepare_deepspeed, prepare_fsdp, prepare_multimodal_messages, print_prompt_completions_sample, profiling_context, profiling_decorator, repeat_iterable_dataset, set_seed, shuffle_sequence_dict, shutdown_event_loop_in_daemon, split_pixel_values_by_grid, split_tensor_dict, start_event_loop_in_daemon, strip_images_from_messages, supports_tool_calling, sys, textwrap, time, torch, transformers, unsplit_pixel_values_by_grid, unwrap_model_for_generation, use_adapter, warnings, AutoModelForSequenceClassification, AutoProcessor, AutoTokenizer, BaseTunerLayer, BitsAndBytesConfig, Callable, CommitScheduler, Dataset, DatasetCard, DatasetCardData, DatasetDict, DistributedBackend, EnvironmentFactory, GRPOConfig, GRPOTrainer, GenerationConfig, IterableDataset, IterableDatasetDict, LoraConfig, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, RepeatSampler, RewardFunc, RolloutFunc, Sampler, SyncRefModelCallback, TrainerCallback, VLLMGeneration, Version, add_fused_lm_head, add_response_schema, atexit, copy, create_model_from_path, create_repo, defaultdict, deque, disable_dropout_in_model, gather, get_callable_name, get_config_model_id, get_peft_model, get_training_chat_template, identity, inspect, is_async_callable, is_chat_template_prefix_preserving, is_jmespath_available, is_peft_available, is_peft_model, logger, nn, np, os, pad, pd, peft, pkg_resources, prepare_deepspeed, prepare_fsdp, repeat_iterable_dataset, set_seed, shutdown_event_loop_in_daemon, start_event_loop_in_daemon, supports_tool_calling, sys, time, torch, transformers, warnings, Version, copy, gather, is_conversational, np, os, pad, parse_response, profiling_context, time, torch, transformers, Any, apply_chat_template, copy, disable_gradient_checkpointing, gather, gather_object, inspect, is_conversational, math, nanmax, nanmin, nanstd, nn, np, os, pad, pd, peft, prepare_multimodal_messages, torch, use_adapter, gather, np, os, pad, profiling_context, torch, transformers, unwrap_model_for_generation, Any, copy, gather, math, nn, np, os, pad, time, torch, transformers, warnings, Any, np, profiling_decorator, shuffle_sequence_dict, split_pixel_values_by_grid, split_tensor_dict, torch, unsplit_pixel_values_by_grid, PeftModel, PreTrainedModel, is_peft_available, logger, os, peft, torch, GRPOTrainer, gather, inspect, nanmax, nanmin, np, os, pad, time, torch) - - -import os -import math -import logging -from typing import * -from dataclasses import dataclass, field -from packaging.version import Version -import torch -import numpy as np -from contextlib import nullcontext -from torch.nn import functional as F -import inspect -from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling -from transformers.training_args import ParallelMode -from unsloth_zoo.device_type import DEVICE_TYPE, DEVICE_TYPE_TORCH, device_synchronize - -# Wrap trainer with padding to right and enable training mode -import functools -from types import MethodType -try: - from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers -except: - def reset_unsloth_gradient_checkpointing_buffers(): pass -# Canonical reset lives in unsloth.models._utils so the SFT auto-packing wrapper and the plain -# Trainer loop can import the same helper; fall back to a no-op only if it can't be imported. -try: - from unsloth.models._utils import _unsloth_reset_stray_compile_cache -except Exception: - def _unsloth_reset_stray_compile_cache(self): pass -try: - from unsloth.models._utils import _unsloth_dataset_column_names -except Exception: - def _unsloth_dataset_column_names(dataset): return dataset.column_names -# Drops/renames config arguments the installed TRL no longer accepts, so a -# script pinned to an older TRL keeps working after an upgrade. Falls back to -# the historical raw passthrough so this can never break trainer construction. -try: - from unsloth.models.rl_config_compat import filter_config_init_kwargs as _unsloth_filter_config_init_kwargs - # A cache file generated here can be imported by an older Unsloth whose filter - # predates `mirrored_from`, so drop the argument rather than raise TypeError. - if "mirrored_from" not in inspect.signature(_unsloth_filter_config_init_kwargs).parameters: - _unsloth_filter_config_init_kwargs_old = _unsloth_filter_config_init_kwargs - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): - return _unsloth_filter_config_init_kwargs_old(config_class, kwargs) -except Exception: - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): return kwargs -def prepare_for_training_mode(f): - @functools.wraps(f) - def wrapper(self, *args, **kwargs): - # Drop any torch.compile graph cache poisoned by a stray pre-train forward. - try: - _unsloth_reset_stray_compile_cache(self) - except Exception: - pass - # Finish the previous W&B run if this is a subsequent train() call. - # We do this at the START of train() (not the end) so that - # evaluate() / log() still work after train() completes. - # HF's WandbCallback.setup() will call wandb.init() for the new run. - # See: https://github.com/unslothai/unsloth/issues/3954 - if getattr(self, '_unsloth_training_completed', False): - try: - import wandb - if wandb.run is not None: - wandb.finish() - # Reset HF's WandbCallback so it calls wandb.init() for the new run - for cb in self.callback_handler.callbacks: - if type(cb).__name__ == 'WandbCallback': - cb._initialized = False - break - except: - pass - # Enable training mode - _was_training = None - # Restore the GC mode the model was configured with at setup; fall back to - # the training args only when it wasn't recorded (issue #4735). Use hasattr, - # not a None sentinel, so a deliberately-recorded None is restored verbatim. - _model = getattr(self, 'model', None) - if hasattr(_model, '_unsloth_gradient_checkpointing'): - use_gc = _model._unsloth_gradient_checkpointing - else: - use_gc = getattr(self.args, 'gradient_checkpointing', True) - if hasattr(self, 'model') and hasattr(self.model, "training"): - _was_training = self.model.training - if hasattr(self, 'model') and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - output = f(self, *args, **kwargs) - # Restore previous mode when possible - if hasattr(self, 'model') and hasattr(self.model, "for_inference"): - if _was_training is False: - self.model.for_inference() - elif _was_training is True and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - # Reset gradient checkpointing buffers to free memory while staying ready for next run - try: - reset_unsloth_gradient_checkpointing_buffers() - except: - pass - # Mark that training completed so the next train() call can - # finish this W&B run before starting a new one - self._unsloth_training_completed = True - return output - return wrapper -pass - -torch_compile_options = { - "epilogue_fusion" : True, - "max_autotune" : False, - "shape_padding" : True, - "trace.enabled" : False, - "triton.enable_persistent_tma_matmul": False, - "cuda.cutlass_epilogue_fusion_enabled": (lambda caps: bool(caps) and len(caps) == torch.cuda.device_count() and min(c[0] for c in caps) >= 9)(__import__("unsloth.device_type", fromlist = ["cuda_visible_capabilities"]).cuda_visible_capabilities()), - "cuda.cutlass_tma_only": (lambda caps: bool(caps) and len(caps) == torch.cuda.device_count() and min(c[0] for c in caps) >= 9)(__import__("unsloth.device_type", fromlist = ["cuda_visible_capabilities"]).cuda_visible_capabilities()), - "cuda.compile_opt_level" : "-O2", - "cuda.enable_cuda_lto" : True, - } - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_hidden_states_selective_log_softmax( - hidden_states: torch.Tensor, - lm_head: torch.Tensor, - index: torch.Tensor, - chunks: int = 4, - logit_scale_multiply: float = 0.0, - logit_scale_divide: float = 0.0, - logit_softcapping: float = 0.0, - temperature: float = 1.0, - # Rows per chunk cap. Read HERE, in the default, not in the body: the body - # is traced with fullgraph = True, and `os.environ` is an unsupported op - # there on torch 2.4 -- Dynamo raises - # torch._dynamo.exc.Unsupported: const method call bytes.decode - # from os._Environ.__getitem__, which the eager fallback does not catch - # (it only catches recompile-limit and disabled-hook breaks), so the very - # first call would die even with the variable unset. A default is evaluated - # once when the def runs, which is import time, outside any traced region. - # It is also a plain int argument, so Dynamo guards on it instead of - # constant-folding an unguarded read (2.7+ never notice a later change). - # A non-numeric value is ignored rather than raised on, so a typo cannot - # break the import. 0 keeps the previous chunk boundaries exactly. - max_rows_per_chunk: int = ( - int(os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip()) - if os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip().isdigit() - else 0 - ), - lm_head_bias: Optional[torch.Tensor] = None, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, - # Read in the default for the same Dynamo reason as max_rows_per_chunk. Pairs with vLLM's - # head_dtype=float32 that load_vllm requests under the same flag. - fp32_lm_head: bool = ( - os.environ.get("UNSLOTH_VLLM_FP32_LM_HEAD", "0").strip().lower() in ("1", "true", "on", "yes") - ), - # Older Unsloth scores old / ref log-probs here and the policy with the fused kernel: one scorer keeps ratio at 1. - fused_linear: bool = ( - os.environ.get("UNSLOTH_FUSED_LINEAR", "1") != "0" - and os.environ.get("UNSLOTH_GRPO_CHUNKED_LOGPROBS", "1") != "0" - ), -) -> torch.Tensor: - # All Unsloth Zoo code licensed under AGPL3 - from unsloth_zoo.fused_losses import chunked_logprob as _chunked_logprob - if (fused_linear and _chunked_logprob._FUSED_LINEAR and not fp32_lm_head - and support_top_k == 0 and support_top_p == 1.0 and support_min_p == 0.0 - and support_excluded_ids is None): - return _chunked_logprob.unsloth_chunked_logprobs( - hidden_states, lm_head, index, - temperature = temperature, - logit_scale_multiply = logit_scale_multiply, - logit_scale_divide = logit_scale_divide, - logit_softcapping = logit_softcapping, - lm_head_bias = lm_head_bias, - max_rows_per_chunk = max_rows_per_chunk if max_rows_per_chunk > 0 else None, - ).reshape((hidden_states.shape[0], hidden_states.shape[1])) - # Reshape on this tensor's own last dim: a no-op, so a wrong-width caller - # cannot have its row count silently rewritten and instead fails at the - # matmul below, which prints both operands. Do not swap in a bare - # torch._check: it reports only "Expected cond to be True", naming neither - # operand, and Dynamo rejects a message-carrying one. Callers dispatch on - # the width first -- see `compute_logprobs_chunk`, the packed path and - # `_pg_grad_forward`. - flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) - flat_index = index.reshape(-1) - - # Each chunk materialises rows x vocab logits and then a float32 copy of - # them, all on the device holding the output head. With a large vocabulary - # and a fixed chunk count that grows with the batch, so the peak scales with - # the batch rather than staying bounded. max_rows_per_chunk caps the rows - # per chunk instead, which is pure loop splitting: more, smaller chunks, - # same concatenated result. 0 (the default) keeps the previous chunk - # boundaries exactly. - if max_rows_per_chunk > 0: - n_rows = flat_hidden_states.shape[0] - chunks = max(chunks, -(-n_rows // max_rows_per_chunk)) - chunks = min(chunks, max(n_rows, 1)) - - # A one-row chunk's head gradient is a K=1 matmul Inductor mis-lowers under dynamic shapes (illegal - # memory access or NaN head gradients, torch 2.13), so keep chunks at 2+ rows when the head trains. - # Exactly `chunks` row ranges, so the unrolled loop length never depends on T: torch.chunk - # returns ceil(T / ceil(T / chunks)) pieces and recompiled per new count; tensor_split guards - # on T % chunks. At most rows // 2 chunks when the head trains keeps every piece at 2+ rows. - lm_head_grad = torch.is_grad_enabled() and lm_head.requires_grad - if lm_head_grad and flat_hidden_states.shape[0] < 2 * chunks: - chunks = max(flat_hidden_states.shape[0] // 2, 1) - - n_rows = flat_hidden_states.shape[0] - chunks = max(1, min(chunks, n_rows)) - bounds = [(i * n_rows) // chunks for i in range(chunks + 1)] - chunked_hidden_states = [flat_hidden_states[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - - # Cast once, outside the loop, so every chunk shares one fp32 head copy (also for autograd). - if fp32_lm_head: - lm_head = lm_head.to(torch.float32) - if lm_head_bias is not None: - lm_head_bias = lm_head_bias.to(torch.float32) - - all_per_token_logps = [] - - for chunk_hidden_states, chunk_index in zip(chunked_hidden_states, chunked_index): - # When the model is dispatched over several devices, the output head can - # sit on a different one from the hidden states, because accelerate - # places the tail of the model on the last device it fills. Co-locate on - # the head's device before the matmul, otherwise this raises - # Unhandled FakeTensor Device Propagation for aten.mm.default, - # found two different devices cuda:0, cuda:1 - # On a single device every .to() here is a no-op and the result is - # bit-identical to before. - chunk_hidden_states = chunk_hidden_states.to(device = lm_head.device, dtype = lm_head.dtype) - chunk_index = chunk_index.to(lm_head.device) - if fp32_lm_head: - # The GRPO loss runs this under bf16 / fp16 autocast, which would downcast the matmul. - with torch.autocast(device_type = chunk_hidden_states.device.type, enabled = False): - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - elif lm_head_bias is None: - chunk_logits = chunk_hidden_states @ lm_head.t() - else: - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - - if logit_scale_multiply != 0.0: - chunk_logits = chunk_logits * logit_scale_multiply - if logit_scale_divide != 0.0: - chunk_logits = chunk_logits / logit_scale_divide - if logit_softcapping != 0.0: - chunk_logits = logit_softcapping * torch.tanh(chunk_logits / logit_softcapping) - - chunk_logits = chunk_logits.to(torch.float32) - - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - per_token_logps = _truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - ) - else: - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - # Return to the caller's device so the concatenation below and every - # downstream consumer see the device they started on. - all_per_token_logps.append(per_token_logps.to(hidden_states.device)) - - all_per_token_logps = torch.concat(all_per_token_logps) - - all_per_token_logps = all_per_token_logps.reshape((hidden_states.shape[0], hidden_states.shape[1])) - return all_per_token_logps - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_selective_log_softmax( - logits, - index, - temperature: float = 1.0, - chunks: int = 4, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, -): - # Fixed `chunks` row ranges; see chunked_hidden_states_selective_log_softmax. - flat_logits = logits.reshape(-1, logits.shape[-1]) - flat_index = index.reshape(-1) - # Fewer rows than chunks: one row per range, as torch.chunk did, never empty ranges. - chunks = max(1, min(chunks, flat_logits.shape[0])) - bounds = [(i * flat_logits.shape[0]) // chunks for i in range(chunks + 1)] - chunked_logits = [flat_logits[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - all_per_token_logps = [] - # Per-chunk selective_log_softmax. - for chunk_logits, chunk_index in zip(chunked_logits, chunked_index): - chunk_logits = chunk_logits.to(torch.float32) - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - all_per_token_logps.append(_truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - )) - continue - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - all_per_token_logps.append(per_token_logps) - pass - all_per_token_logps = torch.concat(all_per_token_logps) - all_per_token_logps = all_per_token_logps.reshape((logits.shape[0], logits.shape[1])) - return all_per_token_logps - -def calculate_pad_tokens_in_prompt( - input_ids: torch.Tensor, - logits_to_keep: int, - pad_token_id: int -) -> torch.Tensor: - """Count left-padded tokens per sequence, e.g. [pad, pad, pad, cat] -> 3.""" - if logits_to_keep >= input_ids.shape[1]: - raise ValueError("logits_to_keep must be smaller than the sequence length.") - - prompt_section = input_ids[:, :-logits_to_keep] - - padding_mask = (prompt_section == pad_token_id) - - pad_token_counts = padding_mask.sum(dim=1) - - return pad_token_counts - -def create_completion_attention_mask( - completion_input_ids: torch.Tensor, - left_pad_tokens_per_prompt: torch.Tensor, - max_left_pad: int, - pad_token_id: int -) -> torch.Tensor: - """Build a completion mask that zeros leading prompt and trailing pad tokens. - - For [p,p,p,c,c,c,pad,pad,pad] (p=sliced prompt, c=completion, pad=padding) - this returns [0,0,0,1,1,1,0,0,0]. - """ - batch_size, completion_len = completion_input_ids.shape - device = completion_input_ids.device - - num_tokens_to_mask = max_left_pad - left_pad_tokens_per_prompt - - indices = torch.arange(completion_len, device=device).unsqueeze(0) - shift_mask = indices >= num_tokens_to_mask.unsqueeze(1) - - non_padding_mask = (completion_input_ids != pad_token_id) - - final_mask = shift_mask & non_padding_mask - - return final_mask - -def left_pack_padding(tensor: torch.Tensor, pad_id: int) -> torch.Tensor: - """Move all padding tokens in each sequence to the right.""" - mask = (tensor != pad_id) - # stable=True since the binary mask is unordered. - sorted_indices = torch.argsort(mask, dim=1, descending=True, stable=True) - packed_tensor = torch.gather(tensor, 1, sorted_indices) - return packed_tensor - -def align_logprobs_with_mask( - logprob_tensor: torch.Tensor, - attention_mask: torch.Tensor, - pad_value: float = 0.0 -) -> torch.Tensor: - """Align a log probability tensor with a given attention mask.""" - - device = logprob_tensor.device - batch_size, logprob_seq_len = logprob_tensor.shape - mask_seq_len = attention_mask.shape[1] - - padded_logprobs = torch.full( - attention_mask.shape, - fill_value=pad_value, - dtype=logprob_tensor.dtype, - device=device - ) - - left_pad_counts = torch.argmax(attention_mask, dim=1) - - cols = torch.arange(logprob_seq_len, device=device) - dest_indices = left_pad_counts.unsqueeze(1) + cols - - # Destination row indices, shape [batch_size, logprob_seq_len]. - row_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand_as(dest_indices) - - # Keep only in-bounds destinations, then scatter via advanced indexing. - valid_mask = dest_indices < mask_seq_len - valid_rows = row_indices[valid_mask] - valid_cols = dest_indices[valid_mask] - valid_vals = logprob_tensor[valid_mask] - padded_logprobs[valid_rows, valid_cols] = valid_vals - - return padded_logprobs - -def align_completion_tool_mask( - tool_mask: torch.Tensor, - completion_mask: torch.Tensor, -) -> torch.Tensor: - """Align a raw completion-length tool/env mask with Unsloth's repacked loss mask.""" - if tool_mask is None: - return completion_mask - if tool_mask.shape[0] != completion_mask.shape[0]: - raise ValueError("tool_mask batch size must match completion_mask batch size.") - - tool_mask = tool_mask.to(device=completion_mask.device) - if tool_mask.shape == completion_mask.shape: - aligned_tool_mask = tool_mask - else: - aligned_tool_mask = align_logprobs_with_mask( - tool_mask, - completion_mask, - pad_value=0, - ) - return completion_mask * aligned_tool_mask.to(dtype=completion_mask.dtype) - -def autotune_batch_and_chunks( - total_input_rows, - seq_len, - hidden_size, - vocab_size, - dtype_bytes=16, - multiplier=None -): - if multiplier is None: - final_m = max(4, seq_len // 4096) - else: - final_m = multiplier - - if torch.cuda.is_available(): - free_bytes, _ = torch.cuda.mem_get_info() - limit_gb = (free_bytes / (1024**3))*.80 - elif hasattr(torch, "xpu") and torch.xpu.is_available(): - # XPU: estimate free memory as total - reserved. - total_mem = torch.xpu.get_device_properties(0).total_memory - reserved_mem = torch.xpu.memory_reserved() - free_bytes = total_mem - reserved_mem - limit_gb = (free_bytes / (1024**3)) * 0.80 - else: - # Fallback: assume 8GB available. - limit_gb = 8.0 - - bytes_to_gb = 1024**3 - - b_vals = torch.arange(total_input_rows, 0, -1, device='cpu', dtype=torch.float32) - - hidden_gb = (b_vals * seq_len * hidden_size * dtype_bytes) / bytes_to_gb - - base_logits = ((b_vals/total_input_rows) * b_vals * seq_len * vocab_size * dtype_bytes) / bytes_to_gb - logits_gb = base_logits / final_m - - total_mem_gb = hidden_gb + logits_gb - - valid_mask = total_mem_gb <= limit_gb - valid_indices = torch.nonzero(valid_mask, as_tuple=False) - - if valid_indices.shape[0] == 0: - #This means your GPU will OOM - # Capped at the row count: unsloth's no-grad pass divides rows by this without max(1, ...). - return max(1, min(4, total_input_rows)), final_m - - best_idx = valid_indices[0].item() - final_b = int(b_vals[best_idx].item()) - - return final_b, final_m - -def sanitize_logprob(logprob): - """Local port of trl.scripts.vllm_serve.sanitize_logprob. - Filters NaN logprobs from vLLM outputs.""" - value = logprob.logprob - if math.isnan(value): - logging.getLogger(__name__).warning( - f"Generated NaN logprob, token logprob '{logprob}' will be ignored" - ) - return None - return value -def _unsloth_grpo_autocast_forced_float32(self): - # Stamped by from_pretrained: UNSLOTH_FORCE_FLOAT32 is process wide, so a model loaded after this trainer was built would answer for it here. - forced = getattr(getattr(self, "model", None), "_unsloth_forced_float32", None) - if forced is None: - forced = os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1" - return forced - -def _unsloth_grpo_autocast(self): - """Decide the GRPO autocast once and latch it on the trainer. ACCELERATE_MIXED_PRECISION is process wide, so a trainer built later but run first would hand this trainer its precision; args belongs to this trainer.""" - if not hasattr(self, "_autocast_enabled"): - args = getattr(self, "args", None) - precision = getattr(args, "mixed_precision", None) - use_bf16 = getattr(args, "bf16", None) - use_fp16 = getattr(args, "fp16", None) - if not isinstance(precision, str): - # transformers < 5 has no args.mixed_precision, but rl.py sets the fp16 / bf16 flags on this same args for every branch it takes. - if isinstance(use_bf16, bool) and isinstance(use_fp16, bool): - precision = "bf16" if use_bf16 else ("fp16" if use_fp16 else "no") - else: - precision = os.environ.get("ACCELERATE_MIXED_PRECISION", "fp16") - self._autocast_dtype = torch.float16 if precision == "fp16" else torch.bfloat16 - # "no" is a real value: full finetuning and an explicit float32 load both set it, and reading it as bfloat16 raises on a T4 or V100. - self._autocast_enabled = precision != "no" - self._autocast_force_float32 = False - forced = _unsloth_grpo_autocast_forced_float32(self) - if forced and precision != "bf16": - # Gemma3 / gpt-oss set "no" but still want float16 autocast; a trainer already on bf16 keeps it, since float16 is what the forced list avoids. - self._autocast_dtype = torch.float16 - self._autocast_enabled = True - self._autocast_force_float32 = True - - return self._autocast_enabled, self._autocast_dtype - -def _unsloth_grpo_autocast_unscored(self): - """Latch for a step that ran no scoring forward. Where the ordinary latch would force float16 autocast on a - forced float32 model (Gemma3 on a T4) without vLLM, record no loss autocast instead: nothing latched there - before the TRL old-logprob gate came back, so zoo trained it without autocast. A forced-family trainer on - bf16 (full finetuning) keeps its bf16 latch. Recorded on the trainer rather than left to zoo's - process-wide fallback, which a later model or trainer may have changed. Under vLLM the old-policy - forward always ran and latched, so it latches as before.""" - fresh = not hasattr(self, "_autocast_enabled") - latched = _unsloth_grpo_autocast(self) - if ( - fresh - and not getattr(self, "use_vllm", False) - and getattr(self, "_autocast_force_float32", False) - ): - self._autocast_dtype = torch.float16 - self._autocast_enabled = False - latched = (self._autocast_enabled, self._autocast_dtype) - return latched - -def _unsloth_grpo_autocast_kwargs(self, device_type = DEVICE_TYPE_TORCH): - """torch.amp.autocast kwargs for GRPO generation.""" - enabled, dtype = _unsloth_grpo_autocast(self) - if not getattr(self, "_autocast_force_float32", False) and torch.is_autocast_enabled( - device_type - ): - # Already inside an autocast: inherit its dtype by omitting the key, since autocast passes whatever it gets to set_autocast_dtype. - return {"enabled": enabled} - return {"enabled": enabled, "dtype": dtype} - -def _unsloth_grpo_support_mismatch(self): - """Why kept-set scoring (UNSLOTH_GRPO_TRUNCATED_SUPPORT) cannot score this trainer, or [] when it can. - - Each scoring call rebuilds the top_k / top_p / min_p kept set from its own logits and always keeps the - scored token. That is the sampler's set only while the scored weights are the rollout's. A reference - model (beta != 0), or old logps TRL keeps because generation is not aligned to the optimizer step - (gradient_accumulation_steps % (steps_per_generation * num_iterations) != 0), keeps another set: a - token outside it gets a denominator that depends on the scored token, so the implied probabilities do - not sum to 1 and the PPO ratio / KL inputs are wrong. A shared set would need the sampler's kept ids. - """ - reasons = [] - if getattr(self, "_unsloth_vllm_sync_skip", False): - reasons.append("vLLM weight sync is skipped, so rollouts use the initial weights") - beta = getattr(self, "beta", 0.0) - if beta is not None and float(beta) != 0.0: - reasons.append(f"beta={beta} scores a reference model") - args = getattr(self, "args", None) - accumulation = int(getattr(args, "gradient_accumulation_steps", 1) or 1) - steps_per_generation = int(getattr(args, "steps_per_generation", None) or accumulation) - num_iterations = getattr(self, "num_iterations", None) - if num_iterations is None: - num_iterations = getattr(args, "num_iterations", 1) - num_iterations = int(num_iterations or 1) - if accumulation % (steps_per_generation * num_iterations) != 0: - reasons.append( - f"num_iterations={num_iterations} / steps_per_generation={steps_per_generation} " - f"reuse rollouts across optimizer steps" - ) - return reasons - -def _unsloth_grpo_refuse_unscored_support(self): - """With compute_efficient no old / ref logps are scored here, so the gate below in - _get_per_token_logps_and_entropies never runs, and the zoo loss scores the current policy over - its own kept set against vLLM's logprobs. When that set cannot be the sampler's (an unsynced - vLLM sampler, a reference model, reused rollouts), mark the trainer refused, which the zoo loss - honours by scoring the full vocabulary.""" - if getattr(self, "_unsloth_truncated_support_refused", False): - return - try: - from unsloth_zoo.rl_replacements import grpo_truncated_support_kwargs - if not grpo_truncated_support_kwargs(self): - return - except Exception: - return - mismatch = _unsloth_grpo_support_mismatch(self) - if mismatch: - import warnings - warnings.warn( - "Unsloth: not scoring over the rollouts' top_k / top_p / min_p kept set " - "(UNSLOTH_GRPO_TRUNCATED_SUPPORT=1, or automatic for truncated vLLM rollouts " - "feeding the IS correction or off-policy mask) because " - f"{', '.join(mismatch)}: the scored policy is not the sampler, so its top_k / top_p / " - "min_p kept set differs from vLLM's and the PPO ratio would be unnormalised. " - "Scoring over the full vocabulary instead.", - stacklevel = 2, - ) - self._unsloth_truncated_support_refused = True - -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_text_configs(config): - """``config`` and its text sub-config, the two places a transform can live: Gemma-4-style configs keep it on ``config.text_config``, T5Gemma-style ones only reach it via ``config.get_text_config()``.""" - if config is None: - return [] - holders = [config] - 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: - holders.append(text_cfg) - return holders - -def _unsloth_get_final_logit_softcapping(model): - """The soft cap the loss applies, under any of its three spellings: ``final_logit_softcapping`` (Gemma), ``logits_soft_cap`` (RecurrentGemma), ``output_logit_soft_cap`` (xLSTM). Returns 0 if unset.""" - config = _unsloth_get_model_config(model) - if config is None: - return 0 - for holder in _unsloth_text_configs(config): - for name in ("final_logit_softcapping", "logits_soft_cap", "output_logit_soft_cap"): - softcap = getattr(holder, name, None) - if softcap: - return softcap - return 0 - -def _unsloth_resolve_logit_scales(model_config): - """``(multiply, divide)`` for the logits, read with the same field table and per-family overrides ``detect_logit_transforms`` uses. - - Only reached when the installed unsloth_zoo predates that helper. Must stay in step with ``resolve_logit_transforms`` in unsloth/models/llama.py: if the forward applies a transform GRPO does not, the policy log-probabilities come from different logits than the ones generated, which silently shifts every importance ratio. - """ - # Granite divides, HyperCLOVA X multiplies, MiniCPM3 scales hidden states (not a logit transform). - overrides = { - ("logits_scaling", "hyperclovax"): "multiply", - ("logits_scaling", "minicpm3"): None, - } - found = {"multiply": 0, "divide": 0} - for holder in _unsloth_text_configs(model_config): - model_type = getattr(holder, "model_type", "") or "" - for bucket, names in ( - ("multiply", ("logit_scale", "lm_head_multiplier", "output_multiplier")), - ("divide", ("logits_scaling",)), - ): - for name in names: - target = overrides.get((name, model_type), bucket) - if target is None or found[target]: - continue - value = getattr(holder, name, 0) or 0 - if value: - found[target] = value - break - return found["multiply"], found["divide"] - -def _unsloth_grpo_returns_hidden_states(model, tensor, lm_head): - """Does ``tensor`` (a forward's ``.logits``) carry hidden states or real logits? - - ``_get_per_token_logps_and_entropies`` sets ``UNSLOTH_RETURN_HIDDEN_STATES=1``, but only a forward that honours the name hands hidden states back as ``.logits``; any other forward returns a real ``[.., vocab]`` tensor that must not reach the ``lm_head`` matmul. - - The primary test is an explicit signal that the forward honours the flag, both set outside this file: ``__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__`` on the generated class, written by ``unsloth_zoo.compiler.create_standalone_class`` exactly when ``apply_fused_lm_head`` gave that forward its own branch; and ``_unsloth_grpo_hidden_states_forward_wrapped``, set by ``_install_grpo_hidden_states_forward_wrapper`` in ``unsloth/models/rl.py`` for models the compiler did not rewrite. That wrapper degrades to real logits when the model cannot produce hidden states and records whether it did so in ``_unsloth_grpo_hidden_states_degraded`` before returning, so reading the pair after a forward describes the call that just finished; degradation is per call, not per model. - - The width comparison stays as the fallback, for an ``unsloth_zoo`` old enough that it never writes the marker. It is decisive whenever ``vocab_size != hidden_size``, and the signal may only overrule it when it is not, which is the one case the shape cannot answer. - """ - if tensor.shape[-1] != lm_head.shape[1]: - return False # vocab-wide: real logits, whatever any signal claims - if lm_head.shape[0] != lm_head.shape[1]: - return True # hidden-wide and vocab_size != hidden_size: hidden states - return _unsloth_grpo_hidden_states_signal(model) is not False - -def _unsloth_grpo_hidden_states_signal(model): - """``True``/``False`` if the forward honours ``UNSLOTH_RETURN_HIDDEN_STATES``, ``None`` when neither marker is present. See ``_unsloth_grpo_returns_hidden_states`` for where each marker is set. Walks the wrapper chain because the markers are set on whichever object the trainer saw, which may be the DDP module or the PEFT base model rather than the object handed to the logprob loop.""" - candidates = [] - pending = [model] - while pending and len(candidates) < 8: - candidate = pending.pop(0) - if candidate is None or any(candidate is seen for seen in candidates): - continue - candidates.append(candidate) - get_base_model = getattr(candidate, "get_base_model", None) - if callable(get_base_model): - try: - pending.append(get_base_model()) - except Exception: - pass - for _attr in ("module", "base_model", "model"): - child = getattr(candidate, _attr, None) - if child is not None and hasattr(child, "forward"): - pending.append(child) - for candidate in candidates: - if getattr(candidate, "__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__", False): - return True - if getattr(type(candidate), "__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__", False): - return True - if any( - getattr(candidate, "_unsloth_grpo_hidden_states_forward_wrapped", False) - for candidate in candidates - ): - # The wrapper honours the flag unless it recorded that this call could not. - if any( - hasattr(candidate, "_unsloth_grpo_hidden_states_degraded") for candidate in candidates - ): - return not any( - getattr(candidate, "_unsloth_grpo_hidden_states_degraded", False) - for candidate in candidates - ) - # An unsloth/models/rl.py predating the per-call attribute set only the warn-once flag; it is the best signal such a wrapper offers. - return not any( - getattr(candidate, "_unsloth_grpo_hidden_states_warning_issued", False) - for candidate in candidates - ) - return None - -def _unsloth_get_mm_token_id(processing_class, attr_name, token): - tokenizer = getattr(processing_class, "tokenizer", processing_class) - token_id = getattr(processing_class, attr_name, None) - if token_id is None: - token_id = getattr(tokenizer, attr_name, None) - - convert_tokens_to_ids = getattr(tokenizer, "convert_tokens_to_ids", None) - if token_id is None and convert_tokens_to_ids is not None: - token_id = convert_tokens_to_ids(token) - - if type(token_id) is int and token_id >= 0: - if token_id != getattr(tokenizer, "unk_token_id", None): - return token_id - return None - -def _unsloth_fix_mm_token_type_ids( - processing_class, input_ids, mm_token_type_ids = None, completion_ids = None -): - image_token_id = _unsloth_get_mm_token_id( - processing_class, "image_token_id", "<|image_pad|>" - ) - video_token_id = _unsloth_get_mm_token_id( - processing_class, "video_token_id", "<|video_pad|>" - ) - - if image_token_id is not None or video_token_id is not None: - rebuilt = input_ids.new_zeros(input_ids.shape) - if image_token_id is not None: - rebuilt = rebuilt.masked_fill(input_ids == image_token_id, 1) - if video_token_id is not None: - rebuilt = rebuilt.masked_fill(input_ids == video_token_id, 2) - return rebuilt - - if ( - mm_token_type_ids is not None - and completion_ids is not None - and mm_token_type_ids.shape[0] == input_ids.shape[0] - and mm_token_type_ids.shape[1] + completion_ids.shape[1] == input_ids.shape[1] - ): - return torch.cat( - [mm_token_type_ids, mm_token_type_ids.new_zeros(completion_ids.shape)], - dim = 1, - ) - return mm_token_type_ids - -def _unsloth_grpo_accumulation_steps(trainer): - """Gradient accumulation divisor for the GRPO loss. 1 whenever the model is not training. - - TRL divides by `current_gradient_accumulation_steps` in train mode and by 1.0 in eval, since - an eval pass accumulates nothing. Trainer sets that attribute inside the training loop and - never clears it, so an in-training evaluation still sees the training window's size and - eval_loss comes back that many times too small. Read off `model.training` like TRL does - rather than a trainer flag, because that is what the eval loop actually switches. - The attribute is also missing when evaluate() runs standalone (#2464), hence the default. - """ - model = getattr(trainer, "model", None) - if model is not None and not getattr(model, "training", True): - return 1 - return getattr(trainer, "current_gradient_accumulation_steps", 1) - -def _unsloth_grpo_is_metric_values(delta, flat_is_ratio, mask, sequence_level): - # Inputs are zero-filled outside the mask; TRL reduces over masked tokens only, else min ratio reads 0. - if mask is None or delta.shape != mask.shape or flat_is_ratio.shape != mask.shape: - return delta.reshape(-1), flat_is_ratio.reshape(-1) - keep = mask.to(torch.bool) - if not sequence_level: - return delta[keep], flat_is_ratio[keep] - # A row with no kept tokens is exp(0) = 1, as in TRL. - counts = keep.sum(dim = -1) - per_row = torch.where( - counts > 0, - (flat_is_ratio * keep).sum(dim = -1) / counts.clamp(min = 1), - torch.ones_like(counts, dtype = flat_is_ratio.dtype), - ) - return delta[keep], per_row - -def _unsloth_grpo_is_sequence_level(args_mode): - # One IS ratio per sequence: TRL's sequence_* and unsloth_zoo's geometric_* (verl seq_mean_k1), - # with the UNSLOTH_GRPO_VLLM_IS_MODE override unsloth_zoo's grpo_accumulated_loss applies. - import os as _os - mode = _os.environ.get("UNSLOTH_GRPO_VLLM_IS_MODE", "") or args_mode or "token_truncate" - return str(mode).startswith(("sequence_", "geometric_")) - -def _unsloth_grpo_sample_weight_mode(): - import os - - exclude = os.environ.get("UNSLOTH_GRPO_ZERO_VARIANCE_EXCLUDE", "0") == "1" - aggregation = os.environ.get("UNSLOTH_GRPO_LOSS_AGGREGATION", "").strip().lower() - if aggregation not in ("", "prompt"): - raise ValueError( - f"Unsloth: UNSLOTH_GRPO_LOSS_AGGREGATION must be 'prompt' or unset, got {aggregation}." - ) - if not exclude and aggregation != "prompt": - return None - return exclude, aggregation == "prompt" - -def _unsloth_grpo_rollout_logps_as_old( - trainer, - sampling_per_token_logps, - completion_mask, - tool_mask = None, -): - """UNSLOTH_GRPO_ROLLOUT_LOGPROBS_AS_OLD=1: use the rollout engine's logprobs (vLLM, or a - rollout_func) as pi_old and skip the old-policy forward, as verl rollout_correction - bypass_mode (loss_type ppo_clip) and vime --use-rollout-logprobs do. The PPO ratio is then - pi_theta / pi_rollout and its clip (plus the dual clip `delta`) is the correction, so no - vLLM importance weight is applied on top. True when this generation batch is marked.""" - import os - - if os.environ.get("UNSLOTH_GRPO_ROLLOUT_LOGPROBS_AS_OLD", "0") != "1": - return False - - def _skip(reason): - trainer._unsloth_rollout_logps_as_old_fallbacks = ( - getattr(trainer, "_unsloth_rollout_logps_as_old_fallbacks", 0) + 1 - ) - warned = getattr(trainer, "_unsloth_rollout_logps_as_old_warned", None) - if not isinstance(warned, set): - warned = set() - trainer._unsloth_rollout_logps_as_old_warned = warned - if reason not in warned: - warned.add(reason) - import warnings - warnings.warn( - f"Unsloth: UNSLOTH_GRPO_ROLLOUT_LOGPROBS_AS_OLD=1 is not applied to this batch: " - f"{reason}. The trainer scores pi_old itself instead.", - stacklevel = 2, - ) - return False - - reason = None - if sampling_per_token_logps is None: - reason = "there are no rollout logprobs (needs use_vllm=True or a rollout_func)" - elif getattr(trainer.args, "loss_type", None) == "vespo": - reason = "loss_type='vespo' weights by the vLLM importance ratio itself" - elif not hasattr(trainer.args, "loss_type"): - reason = "this TRL version predates loss_type" - else: - try: - from unsloth_zoo.rl_replacements import grpo_rollout_logps_as_old # noqa: F401 - except Exception: - reason = "it needs a newer unsloth_zoo: pip install -U unsloth_zoo" - if reason is None: - loss_mask = completion_mask.to(torch.bool) - if tool_mask is not None: - loss_mask = loss_mask & tool_mask.to(device = loss_mask.device, dtype = torch.bool) - if not bool(torch.isfinite(sampling_per_token_logps[loss_mask]).all()): - reason = "some rollout logprobs on loss tokens are not finite" - accelerator = getattr(trainer, "accelerator", None) - if getattr(accelerator, "num_processes", 1) > 1: - # One reduction on every rank, before any of them returns: under FSDP / ZeRO-3 the - # old-policy forward a fallback runs is collective, so all ranks take the same branch. - eligible = torch.tensor( - [int(reason is None)], dtype = torch.int32, device = completion_mask.device - ) - if not bool(accelerator.gather(eligible).min()) and reason is None: - reason = "another process cannot use its rollout logprobs for this batch" - if reason is not None: - return _skip(reason) - trainer._unsloth_rollout_logps_as_old_batches = ( - getattr(trainer, "_unsloth_rollout_logps_as_old_batches", 0) + 1 - ) - return True - -def _unsloth_trl_masks_unscorable(): - try: - import trl - from packaging.version import Version - return Version(trl.__version__) >= Version("1.6.0") - except Exception: - return True - -def _unsloth_grpo_sample_weight( - trainer, mode, rewards_per_func, completion_mask, tool_mask, process_slice -): - """Per-completion loss weights normalised over the whole generation batch (all processes), - or None when off. ScaleRL (arXiv 2510.13786): zero-variance filtering drops groups whose - rewards are all equal from the loss and its token count; prompt-level aggregation averages - tokens within each group, then over groups. Training batches only.""" - settings = _unsloth_grpo_sample_weight_mode() - if settings is None or mode != "train": - return None - exclude, prompt_level = settings - loss_type = getattr(trainer.args, "loss_type", None) - if (prompt_level and loss_type == "luspo") or ( - not prompt_level and loss_type not in ("bnpo", "dapo", "cispo", "vespo") - ): - if not getattr(trainer, "_unsloth_sample_weight_warned", False): - import warnings - warnings.warn( - f"Unsloth: zero-variance exclusion alone needs a token-normalised loss_type " - f"(bnpo, dapo, cispo, vespo) and prompt-level aggregation does not support " - f"luspo; ignored for loss_type={loss_type}.", - stacklevel = 2, - ) - trainer._unsloth_sample_weight_warned = True - return None - num_generations = trainer.num_generations - loss_mask = completion_mask if tool_mask is None else completion_mask * tool_mask - lengths = trainer.accelerator.gather(loss_mask.sum(dim = 1).to(torch.float32)) - - def nan_variance(x, dim): - # TRL's nanstd squared: Bessel corrected, nan (never "zero") with fewer than two values. - centred = x - torch.nanmean(x, dim = dim, keepdim = True) - count = (~torch.isnan(x)).sum(dim = dim, keepdim = True) - variance = torch.nansum(centred * centred, dim = dim, keepdim = True) / (count - 1).clamp(min = 1) - return torch.where(count > 1, variance, torch.full_like(variance, torch.nan)) - - weights = trainer.reward_weights.to(rewards_per_func.device).unsqueeze(0) - if getattr(trainer, "multi_objective_aggregation", None) == "normalize_then_sum": - # TRL normalises each reward function within its group before the weighted sum. - per_func = rewards_per_func.view(-1, num_generations, rewards_per_func.shape[1]) - per_func = (per_func - torch.nanmean(per_func, dim = 1, keepdim = True)) / ( - nan_variance(per_func, 1).sqrt() + 1e-4 - ) - rewards = (per_func.view_as(rewards_per_func) * weights).nansum(dim = 1) - else: - rewards = (rewards_per_func * weights).nansum(dim = 1) - if _unsloth_trl_masks_unscorable(): - # TRL >= 1.6 leaves a completion no reward function scored out of its group's statistics; - # older TRL scores it 0 (nansum), so its group has variance and real advantages there. - rewards[torch.isnan(rewards_per_func).all(dim = 1)] = torch.nan - variance = nan_variance(rewards.view(-1, num_generations), 1).squeeze(1) - group_keep = torch.ones_like(variance, dtype = torch.bool) - if exclude: - # TRL's frac_reward_zero_std test, per group. - group_keep = ~torch.isclose(variance.sqrt(), torch.zeros_like(variance)) - keep = group_keep.repeat_interleave(num_generations).to(torch.float32).to(lengths.device) - if prompt_level: - group_tokens = (lengths * keep).view(-1, num_generations).sum(dim = 1) - groups = group_keep.sum().to(torch.float32).clamp(min = 1.0) - weight = keep / (group_tokens.clamp(min = 1.0).repeat_interleave(num_generations) * groups) - else: - weight = keep / (lengths * keep).sum().clamp(min = 1.0) - trainer._metrics[mode]["frac_reward_zero_std_excluded"].append( - float(1.0 - group_keep.float().mean()) if exclude else 0.0 - ) - return weight[process_slice].to(completion_mask.device) - -def _unsloth_grpo_vision_inputs(source): - """unsloth_zoo owns this key tuple; the copy below is the fallback for a zoo predating - GRPO_VISION_KEYS, and test_grpo_vision_kwargs_forwarded.py fails if the two diverge.""" - try: - from unsloth_zoo.rl_replacements import grpo_get_vision_inputs - return grpo_get_vision_inputs(source) - except Exception: - pass - if source is None: - return {} - get = getattr(source, "get", None) - if get is None: - return {} - return { - key: get(key, None) - for key in ( - "pixel_values", - "image_grid_thw", - "pixel_attention_mask", - "image_sizes", - "spatial_shapes", - "num_tiles", - # both Gemma 4 spellings: image_position_ids is TRL >= 1.1.0 - "image_position_ids", - "pixel_position_ids", - "num_images", - "token_type_ids", - "mm_token_type_ids", - "pixel_values_videos", - "video_grid_thw", - "second_per_grid_ts", - "num_videos", - ) - } - -def _unsloth_grpo_split_vision_by_sample(batch): - """TRL only splits the tile and image indexed vision tensors per sample from its own - split_pixel_values_by_grid, and before TRL 1.1.0 that function knows only the - image_grid_thw layout. Every other layout is left flat, and _prepare_inputs then - shuffles and slices the batch by sample index, so an LFM2-VL or Gemma batch has its - tiles reordered away from the samples they belong to. This mirrors the current - trl.trainer.utils.split_pixel_values_by_grid so an older TRL keeps them together.""" - - def _counts(value): - if value is None: - return None - if hasattr(value, "tolist"): - value = value.tolist() - try: - return [int(count) for count in value] - except TypeError: - return None - - batch = _unsloth_grpo_split_videos_by_sample(batch) - pixel_values = batch.get("pixel_values", None) - if pixel_values is None: - return batch - image_grid_thw = batch.get("image_grid_thw", None) - num_images = _counts(batch.get("num_images", None)) - - if isinstance(pixel_values, list): - # TRL 1.1.0+ splits by sample; 0.22.x-0.23.x splits by grid row and leaves image_grid_thw flat, so - # a two-image row is mis-shuffled. Regroup it. Not len(pixel_values) == len(num_images): [0, 2] - # makes them agree while the axes differ. - _prompt_ids = batch.get("prompt_ids", None) - if ( - image_grid_thw is None - or isinstance(image_grid_thw, list) - or not num_images - or not pixel_values - or len(pixel_values) != sum(num_images) - or all(_count == 1 for _count in num_images) - or (_prompt_ids is not None and len(num_images) != _prompt_ids.shape[0]) - ): - return batch - split = dict(batch) - _empty = pixel_values[0][:0] - _grouped = [] - _offset = 0 - for _count in num_images: - _group = pixel_values[_offset : _offset + _count] - _offset += _count - _grouped.append(torch.cat(_group, dim = 0) if _group else _empty) - split["pixel_values"] = _grouped - split["image_grid_thw"] = list(torch.split(image_grid_thw, num_images, dim = 0)) - for _image_key in ("pixel_attention_mask", "image_sizes"): - _per_image = batch.get(_image_key, None) - if ( - _per_image is not None - and not isinstance(_per_image, list) - and _per_image.shape[0] == sum(num_images) - and _per_image.shape[0] != len(num_images) - ): - split[_image_key] = list(torch.split(_per_image, num_images, dim = 0)) - return split - - if image_grid_thw is not None: - # An unsplit grid batch: TRL owns this layout in every version that persists it. - return batch - if not num_images: - return batch - rows = pixel_values.shape[0] - num_tiles = _counts(batch.get("num_tiles", None)) - split = dict(batch) - for _position_key in ("image_position_ids", "pixel_position_ids"): - _position_ids = batch.get(_position_key, None) - if _position_ids is None or isinstance(_position_ids, list): - continue - if rows != sum(num_images) or _position_ids.shape[0] != sum(num_images): - continue - split["pixel_values"] = list(torch.split(pixel_values, num_images, dim = 0)) - split[_position_key] = list(torch.split(_position_ids, num_images, dim = 0)) - return split - if num_tiles and rows == sum(num_tiles): - split["pixel_values"] = list(torch.split(pixel_values, num_tiles, dim = 0)) - for _tile_key in ("pixel_attention_mask", "spatial_shapes"): - _tiled = batch.get(_tile_key, None) - if ( - _tiled is not None - and not isinstance(_tiled, list) - and _tiled.shape[0] == sum(num_tiles) - ): - split[_tile_key] = list(torch.split(_tiled, num_tiles, dim = 0)) - return split - if rows != sum(num_images): - # One padded row per sample already (Idefics, SmolVLM): TRL leaves this alone. - return batch - if rows == len(num_images) and any(_count != 1 for _count in num_images): - # Padded-sample and flat-image axes are ambiguous: keep TRL layout, else [2, 0] misassigns images - return batch - split["pixel_values"] = list(torch.split(pixel_values, num_images, dim = 0)) - _image_sizes = batch.get("image_sizes", None) - if ( - _image_sizes is not None - and not isinstance(_image_sizes, list) - and _image_sizes.shape[0] == sum(num_images) - ): - split["image_sizes"] = list(torch.split(_image_sizes, num_images, dim = 0)) - return split - -def _unsloth_grpo_unsplit_vision(batch): - """Undo _unsloth_grpo_split_vision_by_sample once this step's slice has been taken, so - the forward sees the layout the processor produced. TRL's own unsplit merges only - pixel_values before 1.1.0, and pixel_values plus the grid and position ids from 1.1.0.""" - merged = None - for key in ( - "pixel_values", - "image_grid_thw", - "pixel_attention_mask", - "spatial_shapes", - "image_sizes", - "image_position_ids", - "pixel_position_ids", - "pixel_values_videos", - "video_grid_thw", - "second_per_grid_ts", - ): - value = batch.get(key, None) - if not isinstance(value, list) or len(value) == 0: - continue - if not hasattr(value[0], "shape"): - # num_images and num_tiles are plain counts, not tensors to merge. - continue - if merged is None: - merged = dict(batch) - merged[key] = torch.cat(value, dim = 0) - return batch if merged is None else merged - -def _unsloth_grpo_split_videos_by_sample(batch): - """Group the flat video rows, grid and second_per_grid_ts by sample (num_videos) so the - _prepare_inputs shuffle keeps them with their sample; _unsloth_grpo_unsplit_vision merges back.""" - pixel_values_videos = batch.get("pixel_values_videos", None) - video_grid_thw = batch.get("video_grid_thw", None) - num_videos = batch.get("num_videos", None) - if ( - pixel_values_videos is None - or video_grid_thw is None - or num_videos is None - or isinstance(pixel_values_videos, list) - ): - return batch - num_videos = [int(n) for n in num_videos] - if sum(num_videos) != video_grid_thw.shape[0]: - return batch - rows_per_sample = [ - int(group.prod(dim = -1).sum()) for group in torch.split(video_grid_thw, num_videos) - ] - split = dict(batch) - split["pixel_values_videos"] = list(torch.split(pixel_values_videos, rows_per_sample)) - split["video_grid_thw"] = list(torch.split(video_grid_thw, num_videos)) - second_per_grid_ts = batch.get("second_per_grid_ts", None) - if second_per_grid_ts is not None: - second_per_grid_ts = torch.as_tensor(second_per_grid_ts) - split["second_per_grid_ts"] = list(torch.split(second_per_grid_ts, num_videos)) - return split - -def _unsloth_grpo_prompt_videos(prompt, kind = "video"): - """Number of `kind` parts in a conversational prompt.""" - if not isinstance(prompt, list): - return 0 - count = 0 - for message in prompt: - content = message.get("content", None) if isinstance(message, dict) else None - if isinstance(content, list): - count += sum( - 1 for part in content if isinstance(part, dict) and part.get("type") == kind - ) - return count - -def _unsloth_grpo_has_video_key(prompt): - if not isinstance(prompt, list): - return False - for message in prompt: - content = message.get("content", None) if isinstance(message, dict) else None - if isinstance(content, list) and any( - isinstance(part, dict) and "video" in part for part in content - ): - return True - return False - -def _unsloth_grpo_clean_video_prompts(prompts, trainer = None): - """Arrow adds `video: None` to every content part of a column holding videos; Qwen2-VL - templates test `'video' in content` and render a placeholder per part (StopIteration).""" - if not any(_unsloth_grpo_has_video_key(prompt) for prompt in prompts): - return prompts - if getattr(trainer, "use_vllm", False) and any( - _unsloth_grpo_prompt_videos(prompt) for prompt in prompts - ): - # Before generation, which vLLM would otherwise run without the videos. - raise NotImplementedError( - "Unsloth: video GRPO runs on the transformers generation path only. " - "Load the model with fast_inference = False (use_vllm = False)." - ) - cleaned = [] - for prompt in prompts: - if not _unsloth_grpo_has_video_key(prompt): - cleaned.append(prompt) - continue - messages = [] - for message in prompt: - content = message.get("content", None) - if isinstance(content, list): - message = dict(message) - message["content"] = [ - {k: v for k, v in part.items() if v is not None} - if isinstance(part, dict) - else part - for part in content - ] - messages.append(message) - cleaned.append(messages) - return cleaned - -def _unsloth_grpo_prompt_key(value): - """Dedup key: equal for the repeated copies of a prompt, never shared by different media. - repr() truncates arrays, so in-memory frames are keyed by identity, not by content.""" - if isinstance(value, dict): - return tuple((k, _unsloth_grpo_prompt_key(v)) for k, v in value.items()) - if isinstance(value, (list, tuple)): - return tuple(_unsloth_grpo_prompt_key(v) for v in value) - if value is None or isinstance(value, (str, int, float, bool)): - return value - return ("id", id(value)) - -def _unsloth_grpo_video_inputs( - trainer, - prompts, - images = None, -): - """Video tensors for the logprob forwards, which TRL builds from images only. Rendered with - the same chat template call generation makes, once per distinct prompt. None without video.""" - num_videos = [_unsloth_grpo_prompt_videos(prompt) for prompt in prompts] - if not any(num_videos): - return None - if getattr(trainer, "use_vllm", False): - raise NotImplementedError( - "Unsloth: video GRPO runs on the transformers generation path only. " - "Load the model with fast_inference = False (use_vllm = False)." - ) - if (images is not None and any(images)) or any( - _unsloth_grpo_prompt_videos(prompt, "image") for prompt in prompts - ): - raise NotImplementedError( - "Unsloth: GRPO does not support image and video inputs in the same batch yet." - ) - try: - from unsloth_zoo.rl_replacements import GRPO_VISION_KEYS - except Exception: - GRPO_VISION_KEYS = () - if "pixel_values_videos" not in GRPO_VISION_KEYS: - raise RuntimeError( - "Unsloth: video GRPO needs a newer unsloth_zoo that forwards video inputs to " - "both GRPO logprob passes. Please upgrade: pip install -U unsloth_zoo" - ) - unique, order = {}, [] - for prompt in prompts: - order.append(unique.setdefault(_unsloth_grpo_prompt_key(prompt), len(unique))) - first = {} - for i, u in enumerate(order): - first.setdefault(u, i) - unique_prompts = [prompts[first[u]] for u in range(len(unique))] - template_kwargs = dict(getattr(trainer, "chat_template_kwargs", None) or {}) - if getattr(trainer, "chat_template", None) is not None: - # GRPOConfig.chat_template, which generation renders with too - template_kwargs["chat_template"] = trainer.chat_template - processed = trainer.processing_class.apply_chat_template( - conversation = unique_prompts, - add_generation_prompt = True, - tokenize = True, - return_dict = True, - padding = True, - return_tensors = "pt", - **template_kwargs, - ) - pixel_values_videos = processed.get("pixel_values_videos", None) - video_grid_thw = processed.get("video_grid_thw", None) - if pixel_values_videos is None or video_grid_thw is None: - raise NotImplementedError( - f"Unsloth: video GRPO supports processors that return pixel_values_videos and " - f"video_grid_thw (Qwen2-VL, Qwen2.5-VL, Qwen3-VL, Qwen3.5); " - f"{type(trainer.processing_class).__name__} does not." - ) - second_per_grid_ts = processed.get("second_per_grid_ts", None) - if second_per_grid_ts is not None: - second_per_grid_ts = torch.as_tensor(second_per_grid_ts) - unique_counts = [num_videos[first[u]] for u in range(len(unique))] - video_index = torch.split(torch.arange(video_grid_thw.shape[0]), unique_counts) - row_ends = video_grid_thw.prod(dim = -1).cumsum(0) - row_starts = row_ends - video_grid_thw.prod(dim = -1) - videos, rows = [], [] - for u in order: - for v in video_index[u].tolist(): - videos.append(v) - rows.append(pixel_values_videos[int(row_starts[v]) : int(row_ends[v])]) - device = trainer.accelerator.device - out = { - "pixel_values_videos": torch.cat(rows, dim = 0).to(device), - "video_grid_thw": video_grid_thw[videos].to(device), - "num_videos": num_videos, - } - if second_per_grid_ts is not None: - out["second_per_grid_ts"] = second_per_grid_ts[videos].to(device) - return out - -def _unsloth_grpo_image_cell(value): - if value is None: - return None - if isinstance(value, (list, tuple)): - return list(value) - return [value] - -def _unsloth_reject_grpo_image_list(inputs, trainer = None): - """Refuse it where the rewrite above missed TRL's spelling: the processor's own error - names neither the column nor the fix. - - `trainer` narrows the refusal to the one runtime mode a legacy TRL cannot carry a multi - image row through. Its vLLM server path sends the raw cells to `VLLMClient.generate`, - which does `[pil_to_base64(img) for img in images]` over the top level entries, so a cell - holding two images reaches `list.save(...)` and dies with an AttributeError naming neither. - Colocate mode and the no vLLM path both go through the processor, which this change fixed, - so they keep working and must not be refused.""" - if trainer is not None: - if not getattr(trainer, "use_vllm", False): - return - if getattr(trainer, "vllm_mode", None) != "server": - return - # Every row: one list cell anywhere desyncs the images and the placeholders. - try: - rows = list(inputs) - except Exception: - return - for _row_index, _row in enumerate(rows): - value = _row.get("image", None) if isinstance(_row, dict) else None - if isinstance(value, (list, tuple)) and len(value) > 1: - raise ValueError( - f"Unsloth: GRPO received a singular `image` column holding {len(value)} " - f"images in row {_row_index}, and this TRL version cannot carry more than " - "one image per row through to the model. " - "Rename the column to `images`, which TRL reads as the per example list of " - "images, or keep one image per row. " - "See https://github.com/unslothai/unsloth/issues/3605" - ) - -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 - -def grpo_compute_loss( - ref, - new, - old, - sampling_per_token_logps, - input_ids, - mask, - beta, - advantages, - **kwargs -): - # All Unsloth Zoo code licensed under AGPL3 - # Optional argument defaults. - loss_type = kwargs.get("loss_type", "grpo") - epsilon_low = kwargs.get("epsilon_low", 0.2) - epsilon_high = kwargs.get("epsilon_high", 0.2) - max_completion_length = kwargs.get("max_completion_length", 8192) - delta = kwargs.get("delta", None) - importance_sampling_level = kwargs.get("importance_sampling_level", "token") - num_items_in_batch = kwargs.get("num_items_in_batch", None) - current_gradient_accumulation_steps = kwargs.get("current_gradient_accumulation_steps", 1) - steps_per_generation = kwargs.get("steps_per_generation", None) - num_processes = kwargs.get("num_processes", 1) - use_vllm = kwargs.get("use_vllm", False) - # The off-policy mask uses vLLM sampling logprobs whenever the batch supplies them (matching TRL); - # the vLLM importance-sampling ratio is applied to the loss only when this flag is on. - vllm_importance_sampling_correction = kwargs.get("vllm_importance_sampling_correction", False) - vllm_importance_sampling_mode = kwargs.get("vllm_importance_sampling_mode", "sequence_mask") - vllm_importance_sampling_cap = kwargs.get("vllm_importance_sampling_cap", 2.0) - vllm_importance_sampling_clip_min = kwargs.get("vllm_importance_sampling_clip_min", None) - vllm_importance_sampling_clip_max = kwargs.get("vllm_importance_sampling_clip_max", 3.0) - get_sapo_token_loss = kwargs.get("get_sapo_token_loss", None) - sapo_temperature_pos = kwargs.get("sapo_temperature_pos", 1.0) - sapo_temperature_neg = kwargs.get("sapo_temperature_neg", 1.05) - get_gamma_weights = kwargs.get("get_gamma_weights", None) - vespo_k_pos = kwargs.get("vespo_k_pos", 2.0) - vespo_lambda_pos = kwargs.get("vespo_lambda_pos", 3.0) - vespo_k_neg = kwargs.get("vespo_k_neg", 3.0) - vespo_lambda_neg = kwargs.get("vespo_lambda_neg", 2.0) - get_off_policy_mask = kwargs.get("get_off_policy_mask", None) - off_policy_mask_threshold = kwargs.get("off_policy_mask_threshold", None) - # TRL top_entropy_quantile: the policy term (not the KL) on high-entropy tokens only. - entropy_mask = kwargs.get("entropy_mask", None) - # Parsed (option, lower_log, upper) tuples from grpo_parse_rollout_rs, or None. - rollout_rs = kwargs.get("rollout_rs", None) - # Per-completion weights from unsloth's generation step (ScaleRL zero-variance exclusion / - # prompt-level aggregation), normalised over the whole generation batch; replaces the loss_type normaliser. - sample_weight = kwargs.get("sample_weight", None) - # Only direct callers see this fallback; the trainer always forwards an explicit value. - use_bias_correction_kl = kwargs.get("use_bias_correction_kl", False) - input_ids = input_ids.unsqueeze(-1) - - importance_sampling_ratio = None - # Bool (B, T) keep mask from rollout_rs; None when rejection is off. - rs_keep = None - - # exp(new - old) and exp(ref - new) below are taken before `mask` is applied. A sequence-packed - # logp path leaves the masked (prompt/pad) columns at 0 while a padded one fills them with a real - # logp, so when new and old/ref disagree there those ratios can overflow to inf and inf * 0 (the - # masked-out loss) becomes nan. Force new/old/ref to share 0 on the masked columns so both ratios - # are exp(0) = 1 there; every loss term below multiplies by `mask`, so this changes nothing. - if mask is not None: - _keep = mask.to(torch.bool) - new = torch.where(_keep, new, 0.0) - if old is not None: old = torch.where(_keep, old, 0.0) - if ref is not None: ref = torch.where(_keep, ref, 0.0) - - if advantages.dim() == 1: - advantages = advantages.unsqueeze(1) - - if off_policy_mask_threshold is not None: - # DeepSeek-V3.2 off-policy mask. The mismatch logprobs are sampling_per_token_logps (vLLM - # sampling logprobs) if present, else old, else new.detach() when both are absent - # (num_iterations == 1 with no vLLM). This mirrors TRL, which defaults old_per_token_logps to - # per_token_logps.detach() so get_off_policy_mask never receives None (it computes - # mismatch - per_token_logps.detach(), so new.detach() yields a zero-KL keep-all mask). The - # callable is a signature-stable adapter installed in grpo_accumulated_loss, so this stays - # fixed across TRL versions with no signature introspection inside this compiled function. - off_policy_mask = get_off_policy_mask( - advantages=advantages, - per_token_logps=new, - sampling_per_token_logps=sampling_per_token_logps if sampling_per_token_logps is not None else (old if old is not None else new.detach()), - mask=mask, - off_policy_threshold=off_policy_mask_threshold, - ) - - with torch.no_grad(): - if use_vllm and sampling_per_token_logps is not None and vllm_importance_sampling_correction: - # Filter out extra leading prompt tokens after left-padding input_ids. - # Match TRL: aggregate log-ratios then exp (product), not sum of exp ratios. - importance_sampling_ratio = (old - sampling_per_token_logps) * mask - # Unscored vLLM tokens arrive as nan and nan * 0 survives the mask: ratio 1, as TRL does. - importance_sampling_ratio = torch.nan_to_num(importance_sampling_ratio, nan = 0.0) - - if vllm_importance_sampling_mode in ["sequence_mask", "sequence_truncate"]: - importance_sampling_ratio = importance_sampling_ratio.sum(dim=-1, keepdim=True) - elif vllm_importance_sampling_mode in ["geometric_mask", "geometric_truncate"]: - # Geometric mean of the token ratios (verl seq_mean_k1): the product blows up on long rows. - importance_sampling_ratio = importance_sampling_ratio.sum(dim=-1, keepdim=True) / mask.sum(dim=-1, keepdim=True).clamp(min=1.0) - - importance_sampling_ratio = torch.exp(importance_sampling_ratio) - - if vllm_importance_sampling_mode in ["token_truncate", "sequence_truncate", "geometric_truncate"]: - importance_sampling_ratio = torch.clamp( - importance_sampling_ratio, - min=vllm_importance_sampling_clip_min, - max=vllm_importance_sampling_clip_max - ) - elif vllm_importance_sampling_mode in ["token_mask", "sequence_mask", "geometric_mask"]: - min_val = ( - vllm_importance_sampling_clip_min - if vllm_importance_sampling_clip_min is not None - else -math.inf - ) - - max_val = ( - vllm_importance_sampling_clip_max - if vllm_importance_sampling_clip_max is not None - else math.inf - ) - - invalid_mis_mask = (importance_sampling_ratio < min_val) | ( - importance_sampling_ratio > max_val - ) - - importance_sampling_ratio = importance_sampling_ratio.masked_fill( - invalid_mis_mask, value=0.0 - ) - else: - raise ValueError( - f"Unknown vLLM importance sampling mode: {vllm_importance_sampling_mode}. Possible values are 'token_truncate', 'token_mask', 'sequence_truncate', 'sequence_mask', 'geometric_truncate' and 'geometric_mask'." - ) - - # verl compute_rollout_rejection_mask: every option must pass. verl then replaces - # response_mask with response_mask * keep, so a rejected token leaves the policy term - # and the KL loss (both selected out below). It stays in the loss normaliser like - # TRL's *_mask modes, since num_items_in_batch is counted before the loss. - # Log ratio is pi_old / pi_rollout; with rollout logps as pi_old (verl bypass mode) it is - # pi_theta / pi_rollout, verl compute_policy_loss_bypass_mode passing old_log_prob=log_prob. - rs_log_ratio = None - if rollout_rs is not None: - if use_vllm and sampling_per_token_logps is not None and vllm_importance_sampling_correction: - rs_log_ratio = old - sampling_per_token_logps - elif kwargs.get("rollout_logps_as_old", False) and old is not None: - rs_log_ratio = new.detach() - old - if rs_log_ratio is not None: - rs_log_ratio = torch.clamp(torch.nan_to_num(rs_log_ratio * mask, nan = 0.0), min = -20.0, max = 20.0) - rs_mask = mask.to(rs_log_ratio.dtype) - rs_count = rs_mask.sum(dim=-1, keepdim=True).clamp(min=1.0) - rs_keep = torch.ones_like(rs_log_ratio) - for rs_option, rs_lower, rs_upper in rollout_rs: - if rs_option.endswith("k1"): - rs_stat = -rs_log_ratio - elif rs_option.endswith("k2"): - rs_stat = 0.5 * rs_log_ratio * rs_log_ratio - else: - rs_stat = torch.expm1(rs_log_ratio) - rs_log_ratio - rs_stat = rs_stat * rs_mask - if rs_option.startswith("seq_sum"): - rs_stat = rs_stat.sum(dim=-1, keepdim=True) - elif rs_option.startswith("seq_mean"): - rs_stat = rs_stat.sum(dim=-1, keepdim=True) / rs_count - elif rs_option.startswith("seq_max"): - # k2, k3 >= 0 and masked entries are 0, so this is the masked max (0 on an empty row). - rs_stat = rs_stat.amax(dim=-1, keepdim=True) - if rs_lower is not None: - rs_pass = (rs_stat >= rs_lower) & (rs_stat <= rs_upper) - else: - rs_pass = rs_stat <= rs_upper - rs_keep = rs_keep * rs_pass.to(rs_keep.dtype) - rs_keep = rs_keep > 0 - if importance_sampling_ratio is not None: - # Select, not multiply: an unclipped ratio can overflow to inf and inf * 0 is nan. - importance_sampling_ratio = torch.where(rs_keep, importance_sampling_ratio, 0.0) - pass - - # Must detach when old is None: exp(new - new.detach()) == 1 but keeps grads correct. - if old is not None: - log_ratio = new - old - else: - log_ratio = new - new.detach() - if kwargs.get("rollout_logps_as_old", False): - # Rollout gap can pass 88 nats (float32 exp overflow, 0 * inf backward); verl - # compute_policy_loss_vanilla clamps the log ratio to [-20, 20] too. - log_ratio = torch.clamp(log_ratio, min = -20.0, max = 20.0) - - if importance_sampling_level == "token": - log_importance_weights = log_ratio - elif importance_sampling_level == "sequence": - log_importance_weights = (log_ratio * mask).sum(-1) / mask.sum(-1).clamp(min=1.0) - log_importance_weights = log_importance_weights.unsqueeze(-1) - else: - raise ValueError( - f"Unknown importance sampling level: {importance_sampling_level}. Possible values are 'token' " - "and 'sequence'." - ) - - coef_1 = torch.exp(log_importance_weights) - - # Reverse KL: low-variance low-bias estimator as used in the GRPO paper. - if beta != 0.0: - # Keep expm1 finite, then cap the final KL term (TRL does not clamp). - kl_log_ratio = torch.clamp(ref - new, min = -20.0, max = 20.0) - kl_i = torch.expm1(kl_log_ratio) - kl_log_ratio - # TRL order: pre-clamp non-detached coef_1, before the loss_type dispatch. - if use_bias_correction_kl: - kl_i = kl_i * coef_1 - # Clamp after correction so a saturated KL also stops the ratio's gradient. - kl_i = torch.clamp(kl_i, min = -10.0, max = 10.0) - else: - # Zeros with the correct shape. - if importance_sampling_level == "sequence": - kl_i = new.new_zeros(new.size(0), 1) - else: - kl_i = torch.zeros_like(new) - - if loss_type == "cispo": - clamped_ratios = torch.clamp(coef_1, max=epsilon_high).detach() - loss_i = -clamped_ratios * advantages * new - elif loss_type in ["grpo", "bnpo", "dr_grpo", "dapo", "luspo"]: - coef_2 = torch.clamp(coef_1, 1 - epsilon_low, 1 + epsilon_high) - - if delta is not None: - loss_1 = torch.clamp(coef_1, max=delta) * advantages - else: - loss_1 = coef_1 * advantages - pass - loss_2 = coef_2 * advantages - loss_i = -torch.min(loss_1, loss_2) - elif loss_type == "sapo": - temperatures = torch.where(advantages > 0, sapo_temperature_pos, sapo_temperature_neg) - soft_coef_1 = torch.sigmoid(temperatures * (coef_1 - 1)) * 4 / temperatures - loss_i = -soft_coef_1 * advantages - elif loss_type == "vespo": - if get_gamma_weights is None: - raise Exception("vespo is only available in TRL 0.26.0+") - phi_seq = get_gamma_weights( - advantages=advantages, - log_ratio_per_token=log_ratio, - mask=mask, - importance_sampling_ratio=importance_sampling_ratio, - k_pos=vespo_k_pos, - lambda_pos=vespo_lambda_pos, - k_neg=vespo_k_neg, - lambda_neg=vespo_lambda_neg, - ) - loss_i = -phi_seq * advantages * new - else: - raise ValueError(f"Unknown loss type: {loss_type}") - - if off_policy_mask_threshold is not None: - loss_i = loss_i * off_policy_mask - - if entropy_mask is not None: - loss_i = loss_i * entropy_mask - - if use_vllm and sampling_per_token_logps is not None and vllm_importance_sampling_correction: - # vespo applies the IS ratio inside get_gamma_weights, so skip it here. - if loss_type != "vespo": - loss_i = loss_i * importance_sampling_ratio - # delta for the metric. - with torch.no_grad(): - delta = torch.abs(old - sampling_per_token_logps) - delta = delta * mask - # Zeroed rather than filtered like TRL, to keep the returned shape. - delta = torch.nan_to_num(delta, nan = 0.0, posinf = math.inf) - flat_is_ratio = importance_sampling_ratio * mask - else: - delta = torch.tensor([]).detach() - flat_is_ratio = torch.tensor([]).detach() - if beta != 0.0: - loss_i = loss_i + beta * kl_i - if rs_keep is not None: - # Rejected tokens drop out of every loss term, KL included (verl rollout_rs). A select, so - # an inf or nan there neither reaches the loss nor turns into a nan gradient. - loss_i = torch.where(rs_keep, loss_i, 0.0) - - mask = mask.to(torch.float32) - n_mask_per_reward = mask.sum(1) - - # https://github.com/huggingface/trl/blob/e8b8499f1f8d76838155b515e414ee98f757d6d5/trl/trainer/grpo_trainer.py#L1624 - if sample_weight is not None: - # Same process / accumulation rescale as the dapo branch, whose weights are 1 / num_items_in_batch. - normalizer = 1.0 / num_processes - if steps_per_generation: - normalizer = normalizer * current_gradient_accumulation_steps / steps_per_generation - loss = (loss_i * mask * sample_weight.to(loss_i.dtype).unsqueeze(-1)).sum() / normalizer - elif loss_type in ["grpo", "sapo"]: - loss = ((loss_i * mask).sum(-1) / mask.sum(-1).clamp(min=1.0)).mean() - loss = loss / current_gradient_accumulation_steps - elif loss_type == "bnpo": - # BNPO uses the local micro-batch token count even when TRL supplies a global count. - loss = (loss_i * mask).sum() / mask.sum().clamp(min=1.0) - loss = loss / current_gradient_accumulation_steps - elif loss_type == "dr_grpo": - loss = (loss_i * mask).sum() / (loss_i.size(0) * max_completion_length) - loss = loss / current_gradient_accumulation_steps - elif loss_type in ["cispo", "dapo", "vespo"]: - # Floor at 1 like TRL: a fully masked batch (mask_truncated_completions) is 0/0 = nan otherwise. - if torch.is_tensor(num_items_in_batch): - normalizer = num_items_in_batch.clamp(min = 1.0) / num_processes - else: - normalizer = max(float(num_items_in_batch), 1.0) / num_processes - # num_items_in_batch spans the whole generation batch; rescale to one accumulation window like TRL. - if steps_per_generation: - normalizer = normalizer * current_gradient_accumulation_steps / steps_per_generation - loss = (loss_i * mask).sum() / normalizer - elif loss_type == "luspo": - # loss_i is (B, T) unless sequence level with beta 0, so mask elementwise (TRL >= 1.10). - loss = (loss_i * mask).sum(-1).mean() - normalizer = current_gradient_accumulation_steps - loss = loss / normalizer - else: - raise ValueError(f"Unknown loss type: {loss_type}") - - # Folded metrics. - def masked_batch_mean(x): - with torch.inference_mode(): - completion_length = n_mask_per_reward.mean() - # Sequence level only if x is narrower than mask: one-token completions are (B, 1) too. - if x.shape[1] == 1 and mask.shape[1] != 1: - return completion_length, x.mean() - else: - # Rows with no tokens (mask_truncated_completions) are left out, as in TRL. - mean_kl_per_reward = (x * mask).sum(1) / n_mask_per_reward.clamp(min = 1.0) - kept_rows = (n_mask_per_reward > 0).sum() - mean_kl = torch.where( - kept_rows == n_mask_per_reward.numel(), - mean_kl_per_reward.mean(), - mean_kl_per_reward.sum() / kept_rows.clamp(min = 1), - ) - return completion_length, mean_kl - # Rejected tokens add nothing to the reported KL either; the denominator keeps every completion - # token, like the loss normaliser. - kl_metric = kl_i if rs_keep is None or beta == 0.0 else torch.where(rs_keep, kl_i, 0.0) - completion_length, mean_kl = masked_batch_mean(kl_metric) - return loss, completion_length, mean_kl, delta, flat_is_ratio, coef_1, mask - -class UnslothEfficientGRPO(torch.autograd.Function): - # All Unsloth Zoo code licensed under AGPL3 - @staticmethod - def forward(ctx, _new_logps, _old_logps, _ref_logps, _sampling_per_token_logps, lm_head, _input_ids, _mask, _advantages, beta, scaler = None, n_chunks = 1, extra_kwargs=None, upstream_scale = None): - if extra_kwargs is None: - extra_kwargs = {} - def compute_loss(new_logps, old_logps, ref_logps, sampling_per_token_logps, input_ids, mask, advantages, scaling): - loss, completion_length, mean_kl, delta, flat_is_ratio, coef_1, _mask = grpo_compute_loss( - ref_logps, - new_logps, - old_logps, - sampling_per_token_logps, - input_ids, - mask, - beta, - advantages, - **extra_kwargs, - ) - - # Scale for mixed precision; return loss.detach() or autograd uses 2x VRAM. - scaled_loss = loss * scaling - return scaled_loss, (loss.detach(), completion_length, mean_kl, delta, flat_is_ratio, coef_1) - pass - - device =_new_logps.device - grad_inputs = torch.empty_like(_new_logps) - accumulated_loss = torch.zeros(1, device = device)[0] - accumulated_completion_length = torch.zeros(1, device = device)[0] - accumulated_mean_kl = torch.zeros(1, device = device)[0] - accumulated_delta = [] - accumulated_flat_is_ratio = [] - accumulated_coef_1 = [] - - def accumulate_chunk( - new_logps_j, - old_logps_j, - ref_logps_j, - sampling_per_token_logps_j, - input_ids_j, - mask_j, - advantages_j, - scaling, - grad_inputs_j, - ): - (chunk_grad_input,), (chunk_loss, (unscaled_loss, chunk_completion_length, chunk_mean_kl, chunk_delta, chunk_flat_is_ratio, chunk_coef_1)) = torch.func.grad_and_value( - compute_loss, - argnums = (0,), - has_aux = True, - )(new_logps_j, old_logps_j, ref_logps_j, sampling_per_token_logps_j, input_ids_j, mask_j, advantages_j, scaling) - accumulated_loss .add_(unscaled_loss) - accumulated_completion_length.add_(chunk_completion_length) - accumulated_mean_kl .add_(chunk_mean_kl) - accumulated_delta .append(chunk_delta) - accumulated_flat_is_ratio .append(chunk_flat_is_ratio) - accumulated_coef_1 .append(chunk_coef_1) - grad_inputs_j[:] = chunk_grad_input - pass - - from unsloth_zoo.temporary_patches.utils import torch_compile_with_fallback - accumulate_chunk = torch_compile_with_fallback( - fullgraph = True, - # [TODO] Dynamic marking causes torch.compile errors if sequence length is long - dynamic = True, - options = torch_compile_options, - )(accumulate_chunk) - - grad_inputs_chunks = torch.chunk(grad_inputs, chunks = n_chunks, dim = 0) - new_logps = torch.chunk(_new_logps, chunks = n_chunks, dim = 0) - if _old_logps is not None: - old_logps = torch.chunk(_old_logps, chunks = n_chunks, dim = 0) - else: - old_logps = [None] * n_chunks - if _ref_logps is not None: - ref_logps = torch.chunk(_ref_logps, chunks = n_chunks, dim = 0) - else: - ref_logps = [None] * n_chunks - if _sampling_per_token_logps is not None: - sampling_per_token_logps = torch.chunk(_sampling_per_token_logps, chunks = n_chunks, dim = 0) - else: - sampling_per_token_logps = [None] * n_chunks - input_ids = torch.chunk(_input_ids, chunks = n_chunks, dim = 0) - mask = torch.chunk(_mask, chunks = n_chunks, dim = 0) - advantages = torch.chunk(_advantages, chunks = n_chunks, dim = 0) - - # Mixed precision scaling if present. - scaling = scaler.get_scale() if scaler is not None else 1.0 - # A float scale is a dynamo guard (recompile per change); a float32 0-dim tensor is not. - if scaler is not None and not torch.is_tensor(scaling): - scaling = torch.full((), float(scaling), dtype = torch.float32, device = device) - - for (grad_inputs_j, new_logps_j, old_logps_j, ref_logps_j, sampling_per_token_logps_j, input_ids_j, mask_j, advantages_j, ) in\ - zip(grad_inputs_chunks, new_logps, old_logps, ref_logps, sampling_per_token_logps, input_ids, mask, advantages): - - # [TODO] Dynamic marking causes torch.compile errors if sequence length is long - - # mark_dynamic(new_hidden_states_j) - # mark_dynamic(ref_hidden_states_j) - # if old_hidden_states_j is not None: - # mark_dynamic(old_hidden_states_j) - # mark_dynamic(input_ids_j) - # mark_dynamic(mask_j) - accumulate_chunk( - new_logps_j, - old_logps_j, - ref_logps_j, - sampling_per_token_logps_j, - input_ids_j, - mask_j, - advantages_j, - scaling, - grad_inputs_j, - ) - pass - - grad_inputs .div_(n_chunks) - accumulated_loss .div_(n_chunks) - accumulated_completion_length.div_(n_chunks) - accumulated_mean_kl .div_(n_chunks) - - if _sampling_per_token_logps is not None: - accumulated_delta = torch.cat(accumulated_delta, dim=0) - accumulated_flat_is_ratio = torch.cat(accumulated_flat_is_ratio, dim=0) - else: - accumulated_delta = None - accumulated_flat_is_ratio = None - accumulated_coef_1 = torch.cat(accumulated_coef_1, dim=0) - ctx.save_for_backward(grad_inputs) - ctx.upstream_scale = upstream_scale - return ( - accumulated_loss, - accumulated_completion_length, - accumulated_mean_kl, - accumulated_delta, - accumulated_flat_is_ratio, - accumulated_coef_1 - ) - pass - - @staticmethod - def backward(ctx, grad_output, dcompletion_length, dmean_kl, ddelta, ddflat_is_ratio, dcoef_1): - (grad_input,) = ctx.saved_tensors - if ctx.upstream_scale is not None: - grad_input = grad_input * (grad_output * ctx.upstream_scale) - return (grad_input, None, None, None, None, None, None, None, None, None, None, None, None) - pass - -def grpo_accumulated_loss( - trainer, - input_ids, - attention_mask, - logits_to_keep, - completion_mask, - advantages, - old_logps, - ref_logps, - n_chunks = -1, - tool_mask = None, - **kwargs, -): - # All Unsloth Zoo code licensed under AGPL3 - # Body-local import so the copy inlined into the generated trainer cache resolves. - try: - from unsloth_zoo.rl_replacements import _warn_unsupported_grpo_options - _warn_unsupported_grpo_options(trainer) - except Exception: - pass - - # Body-local: this source is copied into the generated trainer without its imports. - from unsloth_zoo.rl_replacements import ( - grpo_shared_vision_inputs as _grpo_get_vision_inputs, - grpo_vision_chunks as _grpo_vision_chunks, - grpo_chunked_hidden_states_logprobs as chunked_hidden_states_selective_log_softmax, - chunked_selective_log_softmax, - grpo_rollout_logps_as_old, - grpo_rollout_as_old_ratio_metrics, - grpo_rollout_old_fill_unscored, - ) - if os.environ.get("UNSLOTH_GRPO_LOGPROB_RECOMPUTE", "0") == "1": - from unsloth_zoo.rl_replacements import ( - chunked_hidden_states_selective_log_softmax_recompute as chunked_hidden_states_selective_log_softmax, - ) - # Adaptive entropy control is not implemented, so it disables the bonus (and warns). - _entropy_quantile = getattr(trainer.args, "top_entropy_quantile", 1.0) - _entropy_mask_on = ( - _entropy_quantile is not None and _entropy_quantile < 1.0 - and hasattr(trainer, "get_high_entropy_mask") - ) - _entropy_coef = float(getattr(trainer, "entropy_coef", getattr(trainer.args, "entropy_coef", 0.0)) or 0.0) - _entropy_bonus_on = _entropy_coef != 0.0 and not getattr(trainer.args, "use_adaptive_entropy", False) - _need_entropy = _entropy_mask_on or _entropy_bonus_on - if _need_entropy: - from unsloth_zoo.rl_replacements import ( - chunked_hidden_states_selective_log_softmax_recompute as _entropy_log_softmax, - UnslothGRPOEntropyBonus as _UnslothGRPOEntropyBonus, - grpo_entropy_bonus_weights as _grpo_entropy_bonus_weights, - grpo_high_entropy_mask as _grpo_high_entropy_mask, - ) - vision_inputs = _grpo_get_vision_inputs(kwargs) - pixel_values = vision_inputs.get('pixel_values', None) - image_grid_thw = vision_inputs.get('image_grid_thw', None) - video_grid_thw = vision_inputs.get('video_grid_thw', None) - if pixel_values is None: - # Only a sentinel below (the chunks carry the tensors): video rows must take the vision - # path too, since left and sequence packing move M-RoPE video positions. - pixel_values = vision_inputs.get('pixel_values_videos', None) - # Released unsloth 2026.9.4 decides whether multi-image GRPO is supported by grepping - # inspect.getsource(grpo_accumulated_loss) for "num_images", so moving the handling into - # grpo_vision_chunks makes that probe answer no and raise "Please upgrade unsloth_zoo" at - # the user who just did. The chunker reads num_images out of vision_inputs itself; this - # binding is what the released probe looks for, and it keeps the name meaningful here. - num_images = vision_inputs.get('num_images', None) - # Transformers 5.x requires token_type_ids/mm_token_type_ids for some vision models - token_type_ids = vision_inputs.get('token_type_ids', None) - mm_token_type_ids = vision_inputs.get('mm_token_type_ids', None) - if mm_token_type_ids is not None or image_grid_thw is not None or video_grid_thw is not None: - mm_token_type_ids = _unsloth_fix_mm_token_type_ids( - trainer.processing_class, input_ids, mm_token_type_ids - ) - vision_inputs['mm_token_type_ids'] = mm_token_type_ids - # Thread vLLM sampling logprobs when something actually consumes them: the off-policy mask - # (off_policy_mask_threshold) or the IS ratio (vllm_importance_sampling_correction). The mask - # needs them regardless of IS correction (matching TRL, which feeds them to get_off_policy_mask - # either way); the IS ratio stays gated on the correction flag inside grpo_compute_loss. On the - # plain vLLM path (neither active) they are dropped so nothing pays for an unused aligned/compiled - # input and grpo_compute_loss returns None (not empty) delta/flat_is_ratio. - # Set by unsloth when UNSLOTH_GRPO_ROLLOUT_LOGPROBS_AS_OLD=1 skipped the old-policy forward. - _rollout_logps_as_old = bool(kwargs.pop("rollout_logps_as_old", False)) - _sampling_logps_used = ( - getattr(trainer, "vllm_importance_sampling_correction", False) - or getattr(trainer.args, "off_policy_mask_threshold", None) is not None - or _rollout_logps_as_old - ) - sampling_per_token_logps = kwargs.get("sampling_per_token_logps", None) if _sampling_logps_used else None - temperature = kwargs.get("temperature", 1.0) - logit_scale_multiply = kwargs.get("logit_scale_multiply", 0.0) - logit_scale_divide = kwargs.get("logit_scale_divide", 0.0) - logit_softcapping = kwargs.get("logit_softcapping", 0.0) - prev_max_left_pad = kwargs.get("max_left_pad", 0) # max_left_pad for LLM training, enabled by default. - - # Pop from kwargs to avoid downstream issues. - _ = kwargs.pop("sampling_per_token_logps", None) - kwargs["vllm_importance_sampling_cap"] = getattr(trainer.args, "vllm_importance_sampling_cap", None) - # Older TRL lacks this arg; fall back to token_truncate (legacy clamp(max=cap) behavior). - kwargs["vllm_importance_sampling_mode"] = getattr(trainer.args, "vllm_importance_sampling_mode", None) or "token_truncate" - # Opt-in for TRL versions without the field or the geometric modes, e.g. geometric_mask. - if os.environ.get("UNSLOTH_GRPO_VLLM_IS_MODE", ""): - kwargs["vllm_importance_sampling_mode"] = os.environ["UNSLOTH_GRPO_VLLM_IS_MODE"] - # Like TRL, only check the mode where it is used: with the correction off it stays a no-op. - if getattr(trainer, "use_vllm", False) and getattr(trainer, "vllm_importance_sampling_correction", False) and kwargs["vllm_importance_sampling_mode"] not in ( - "token_truncate", "token_mask", "sequence_truncate", "sequence_mask", "geometric_truncate", "geometric_mask", - ): - raise ValueError(f"Unsloth: unknown vLLM importance sampling mode {kwargs['vllm_importance_sampling_mode']}.") - # TRL validates this in __init__, before the env override: get_gamma_weights sums a per-token log ratio. - if ( - kwargs.get("loss_type", getattr(trainer, "loss_type", None)) == "vespo" - and getattr(trainer, "use_vllm", False) - and getattr(trainer, "vllm_importance_sampling_correction", False) - and kwargs["vllm_importance_sampling_mode"] not in ("token_truncate", "token_mask") - ): - raise ValueError( - "Unsloth: VESPO loss requires vllm_importance_sampling_mode to be 'token_truncate' or 'token_mask'. " - f"Got: {kwargs['vllm_importance_sampling_mode']}." - ) - kwargs["vllm_importance_sampling_clip_min"] = getattr(trainer.args, "vllm_importance_sampling_clip_min", None) - kwargs["vllm_importance_sampling_clip_max"] = getattr(trainer.args, "vllm_importance_sampling_clip_max", None) - kwargs["get_sapo_token_loss"] = trainer.get_sapo_token_loss if hasattr(trainer, "get_sapo_token_loss") else None - kwargs["sapo_temperature_pos"] = trainer.args.sapo_temperature_pos if hasattr(trainer.args, "sapo_temperature_pos") else None - kwargs["sapo_temperature_neg"] = trainer.args.sapo_temperature_neg if hasattr(trainer.args, "sapo_temperature_neg") else None - kwargs["get_gamma_weights"] = trainer.get_gamma_weights if hasattr(trainer, "get_gamma_weights") else None - kwargs["vespo_k_pos"] = trainer.args.vespo_k_pos if hasattr(trainer.args, "vespo_k_pos") else 2.0 - kwargs["vespo_k_neg"] = trainer.args.vespo_k_neg if hasattr(trainer.args, "vespo_k_neg") else 3.0 - kwargs["vespo_lambda_pos"] = trainer.args.vespo_lambda_pos if hasattr(trainer.args, "vespo_lambda_pos") else 3.0 - kwargs["vespo_lambda_neg"] = trainer.args.vespo_lambda_neg if hasattr(trainer.args, "vespo_lambda_neg") else 2.0 - off_policy_mask_threshold = trainer.args.off_policy_mask_threshold if hasattr(trainer.args, "off_policy_mask_threshold") else None - kwargs["off_policy_mask_threshold"] = off_policy_mask_threshold - # get_off_policy_mask exists on TRL >= 0.27.0; its 3rd parameter was `old_per_token_logps` in 0.27.0 - # and renamed to `sampling_per_token_logps` in 0.27.1 (huggingface/trl#4857), so a fixed keyword call - # crashes on one side of the rename. Wrap it in a signature-stable adapter here, outside the compiled - # loss, so grpo_compute_loss always calls it with one keyword. Detect the real name once via inspect - # and cache the adapter on the trainer; a fresh closure every step would re-trigger torch.compile. - _off_policy_mask_fn = trainer.get_off_policy_mask if hasattr(trainer, "get_off_policy_mask") else None - if _off_policy_mask_fn is None or off_policy_mask_threshold is None: - kwargs["get_off_policy_mask"] = None - else: - _adapter = getattr(trainer, "_unsloth_off_policy_mask_adapter", None) - # Compare by value, not identity: trainer.get_off_policy_mask returns a fresh bound-method - # object on every access, so `is not` would always miss and rebuild the adapter each step - # (re-triggering torch.compile). `!=` on bound methods compares __self__ and __func__, so the - # cached adapter is reused; the first call still rebuilds since None != the bound method. - if getattr(_adapter, "_unsloth_wrapped", None) != _off_policy_mask_fn: - import inspect as _inspect - try: - _params = _inspect.signature(_off_policy_mask_fn).parameters - except (TypeError, ValueError): - _params = {} - if "old_per_token_logps" in _params and "sampling_per_token_logps" not in _params: - # TRL 0.27.0 named the mismatch-logprobs parameter old_per_token_logps. - def _adapter(advantages, per_token_logps, sampling_per_token_logps, mask, off_policy_threshold): - return _off_policy_mask_fn( - advantages=advantages, - per_token_logps=per_token_logps, - old_per_token_logps=sampling_per_token_logps, - mask=mask, - off_policy_threshold=off_policy_threshold, - ) - else: - # TRL >= 0.27.1 / 1.7.x use sampling_per_token_logps (also the default going forward). - def _adapter(advantages, per_token_logps, sampling_per_token_logps, mask, off_policy_threshold): - return _off_policy_mask_fn( - advantages=advantages, - per_token_logps=per_token_logps, - sampling_per_token_logps=sampling_per_token_logps, - mask=mask, - off_policy_threshold=off_policy_threshold, - ) - _adapter._unsloth_wrapped = _off_policy_mask_fn - trainer._unsloth_off_policy_mask_adapter = _adapter - kwargs["get_off_policy_mask"] = _adapter - # Read inside the compiled loss to gate the IS ratio, which the off-policy mask must not gate. - kwargs["vllm_importance_sampling_correction"] = getattr(trainer, "vllm_importance_sampling_correction", False) - # Follows TRL's own value; older TRL has no such field and False is correct there. - kwargs["use_bias_correction_kl"] = getattr(trainer.args, "use_bias_correction_kl", False) - kwargs["use_vllm"] = trainer.use_vllm - # Eval batches are not split or accumulated, so keep the full count. - _training = getattr(getattr(trainer, "model", None), "training", True) - kwargs["steps_per_generation"] = getattr(trainer.args, "steps_per_generation", None) if _training else None - # Generated trainers still pass unsloth_num_chunks; nothing downstream reads it. - try: - from unsloth_zoo.rl_replacements import _warn_deprecated_n_chunks - _warn_deprecated_n_chunks(n_chunks) - except Exception: - pass - - if kwargs["vllm_importance_sampling_clip_max"] is None and kwargs["vllm_importance_sampling_cap"] is not None: - kwargs["vllm_importance_sampling_clip_min"] = 0 - kwargs["vllm_importance_sampling_clip_max"] = kwargs["vllm_importance_sampling_cap"] - # "lower_upper" (either side may be empty), e.g. "0.999_1.001" for the geometric modes. - _is_clip_env = os.environ.get("UNSLOTH_GRPO_VLLM_IS_CLIP", "") - if _is_clip_env: - from unsloth_zoo.rl_replacements import grpo_parse_is_clip as _grpo_parse_is_clip - kwargs["vllm_importance_sampling_clip_min"], kwargs["vllm_importance_sampling_clip_max"] = _grpo_parse_is_clip( - _is_clip_env, kwargs["vllm_importance_sampling_mode"], - ) - # verl rollout_rs rejection (token_k3, seq_mean_k3, seq_mean_k1, ...), applied with the vLLM IS ratio. - from unsloth_zoo.rl_replacements import grpo_parse_rollout_rs as _grpo_parse_rollout_rs - _rollout_rs = _grpo_parse_rollout_rs( - os.environ.get("UNSLOTH_GRPO_ROLLOUT_RS", ""), os.environ.get("UNSLOTH_GRPO_ROLLOUT_RS_THRESHOLD", ""), - ) - if _rollout_rs is not None: - kwargs["rollout_rs"] = _rollout_rs - _rollout_metrics_on = os.environ.get("UNSLOTH_GRPO_ROLLOUT_METRICS", "0") == "1" - - if not hasattr(trainer, '_autocast_dtype'): - # "no" is float32 training (a T4 / V100 without bfloat16): autocasting it to bfloat16 raises there. - _mixed_precision = os.environ.get('ACCELERATE_MIXED_PRECISION', 'fp16') - trainer._autocast_dtype = None if _mixed_precision == 'no' else (torch.float16 if _mixed_precision == 'fp16' else torch.bfloat16) - if os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1': trainer._autocast_dtype = None - pass - # Restored in `finally`: an OOM or interrupt below must not leave later forwards returning hidden states. - _unsloth_prior_hidden_states = os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES") - os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" - try: - output_embeddings = trainer.model.get_output_embeddings() - lm_head = output_embeddings.weight - lm_head_bias = getattr(output_embeddings, "bias", None) - # ZeRO-3 released the head after the forward: score with one gathered copy per call. - _zero3_head = None - from unsloth_zoo.fused_losses.cross_entropy_loss import unsloth_zero3_partitioned, unsloth_zero3_gathered - if unsloth_zero3_partitioned(lm_head) or unsloth_zero3_partitioned(lm_head_bias): - _zero3_head = (unsloth_zero3_gathered(lm_head), unsloth_zero3_gathered(lm_head_bias)) - lm_head, lm_head_bias = _zero3_head - - # UNSLOTH_VLLM_FP32_LM_HEAD: one fp32 head copy per call, shared by every mini-batch kept for the loss. - import inspect as _inspect - from unsloth_zoo.rl_replacements import chunked_hidden_states_selective_log_softmax as _projector - _fp32_lm_head = _inspect.signature(_projector).parameters["fp32_lm_head"].default - _fp32_cache = {} - - def _fp32_bind(): - nonlocal lm_head, lm_head_bias - if not _fp32_lm_head or lm_head.dtype == torch.float32: - return - # FSDP2 unshard / ZeRO-3 gather swap the objects; a grad-mode change needs a fresh cast too. - key = (lm_head, lm_head_bias, torch.is_grad_enabled()) - if _fp32_cache.get("key") is None or any(a is not b for a, b in zip(_fp32_cache["key"], key)): - bias = None if lm_head_bias is None else lm_head_bias.to(torch.float32) - _fp32_cache.update(key = key, value = (lm_head.to(torch.float32), bias)) - lm_head, lm_head_bias = _fp32_cache["value"] - - def _refresh_lm_head(): - # FSDP2's forward unshard swaps the Parameters: the bindings above are metadata-only until then. - nonlocal lm_head, lm_head_bias - if _zero3_head is not None: - lm_head, lm_head_bias = _zero3_head - return - from unsloth_zoo.fused_losses.cross_entropy_loss import unsloth_fsdp_unshard - unsloth_fsdp_unshard(output_embeddings) - lm_head = output_embeddings.weight - lm_head_bias = getattr(output_embeddings, "bias", None) - - from unsloth_zoo.rl_replacements import grpo_truncated_support_kwargs - _support_kwargs = grpo_truncated_support_kwargs(trainer) - # old / ref logps come from Unsloth's _get_per_token_logps, which marks the trainer when - # it scored them over the same set; mixing the two normalisations would bias every ratio. - if _support_kwargs and (old_logps is not None or ref_logps is not None)\ - and not getattr(trainer, "_unsloth_truncated_support_scored", False): - if not getattr(trainer, "_unsloth_truncated_support_refused", False): - from unsloth_zoo.log import logger as _zoo_logger - _zoo_logger.warning( - "Unsloth: UNSLOTH_GRPO_TRUNCATED_SUPPORT (scoring over the rollouts' kept set) needs an " - "unsloth build that scores the old / reference logprobs over the same kept set. " - "Scoring the full vocabulary instead: pip install -U unsloth" - ) - trainer._unsloth_truncated_support_refused = True - _support_kwargs = {} - if _support_kwargs and getattr(trainer, "_unsloth_truncated_support_refused", False): - # Unsloth found the scored policy is not the sampler (an unsynced vLLM sampler, a reference - # model, reused rollouts): the current policy's own kept set would not be vLLM's either. - _support_kwargs = {} - if _support_kwargs: - from functools import partial - from unsloth_zoo.rl_replacements import ( - _GRPO_TRUNCATED_SUPPORT_STATS as _support_stats, - grpo_truncated_support_setting as _support_setting, - ) - _support_stats["on" if _support_setting() == "on" else "auto"] += 1 - if os.environ.get("UNSLOTH_GRPO_LOGPROB_RECOMPUTE", "0") == "1": - # The recompute helper scores the full vocabulary only, so the kept set uses the plain helper - from unsloth_zoo.rl_replacements import chunked_hidden_states_selective_log_softmax - if not getattr(trainer, "_unsloth_truncated_support_recompute_warned", False): - from unsloth_zoo.log import logger as _zoo_logger - _zoo_logger.warning( - "Unsloth: UNSLOTH_GRPO_LOGPROB_RECOMPUTE=1 is not used with UNSLOTH_GRPO_TRUNCATED_SUPPORT=1; " - "the logprobs are scored with the standard helper over the kept set." - ) - trainer._unsloth_truncated_support_recompute_warned = True - chunked_hidden_states_selective_log_softmax = partial( - chunked_hidden_states_selective_log_softmax, **_support_kwargs, - ) - chunked_selective_log_softmax = partial(chunked_selective_log_softmax, **_support_kwargs) - # Unsloth keeps _autocast_dtype set when it turns autocast off (float32 training on a GPU without bfloat16). - _autocast_on = trainer._autocast_dtype is not None and getattr(trainer, "_autocast_enabled", True) - dtype_bytes = 16 if _autocast_on and trainer._autocast_dtype in [torch.float16, torch.bfloat16] else 32 - if _support_kwargs: - # Kept-set scoring saves one more float32 [rows, vocab] tensor per chunk for the backward - # (the masked logits), so the mini-batch autotune counts the logits twice. No effect - # unless free memory already limits the mini-batch. - dtype_bytes = dtype_bytes * 2 - - 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 trainer.args.unsloth_grpo_mini_batch is None: - # Size per call, as unsloth's copy does: caching in args froze the first step's plan. - B, multiplier = autotune_batch_and_chunks( - total_rows, seq_len, hidden_dim, vocab_dim, dtype_bytes, trainer.args.unsloth_logit_chunk_multiplier - ) - B = max(1, total_rows//B) - else: - if trainer.args.unsloth_grpo_mini_batch > total_rows: - B = total_rows - else: - B = trainer.args.unsloth_grpo_mini_batch - - if trainer.args.unsloth_logit_chunk_multiplier is None: - multiplier = max(4, seq_len // 4096) - else: - multiplier = trainer.args.unsloth_logit_chunk_multiplier - - # The text path rebuilds completion_mask from token ids, which would undo TRL's - # mask_truncated_completions row zeroing; keep TRL's rows and reapply them before the loss. - kept_completion_rows = None - if ( - pixel_values is None - and getattr(trainer, "mask_truncated_completions", False) - and torch.is_tensor(completion_mask) - and completion_mask.dim() == 2 - and completion_mask.shape[0] == input_ids.shape[0] - ): - kept_completion_rows = completion_mask.sum(dim = 1, keepdim = True) > 0 - - if pixel_values is None: - left_pad_tokens_per_prompt = calculate_pad_tokens_in_prompt(input_ids, logits_to_keep, trainer.processing_class.pad_token_id) - - # Determine max_left_pad from precomputed logprobs shape for consistency - if old_logps is not None: - max_left_pad = old_logps.shape[1] - logits_to_keep - elif ref_logps is not None: - max_left_pad = ref_logps.shape[1] - logits_to_keep - else: - max_left_pad = torch.max(left_pad_tokens_per_prompt).item() - - input_ids = left_pack_padding(input_ids, trainer.processing_class.pad_token_id) - - completion_input_ids = input_ids[:, -(logits_to_keep +max_left_pad):] - completion_mask = create_completion_attention_mask(completion_input_ids, left_pad_tokens_per_prompt, max_left_pad, trainer.processing_class.pad_token_id).to(attention_mask.dtype) - - if sampling_per_token_logps is not None: - sampling_per_token_logps = align_logprobs_with_mask(sampling_per_token_logps, completion_mask) - else: - sampling_per_token_logps = None - completion_mask = align_completion_tool_mask(tool_mask, completion_mask) - attention_mask = input_ids != trainer.processing_class.pad_token_id - attention_mask = attention_mask.to(attention_mask.dtype) - else: - completion_input_ids = input_ids[:, -logits_to_keep:] - completion_mask = align_completion_tool_mask(tool_mask, completion_mask) - - if _rollout_logps_as_old: - old_logps, sampling_per_token_logps = grpo_rollout_logps_as_old(sampling_per_token_logps, old_logps) - - unwrapped_model = trainer.accelerator.unwrap_model(trainer.model, keep_fp32_wrapper = False) - - for module in unwrapped_model.modules(): - if hasattr(module, "_hf_hook") and hasattr(module._hf_hook, "io_same_decice"): - module._hf_hook.io_same_decice = False - pass - - all_logprobs_list = [] - all_entropies_list = [] - - # No width dispatch in these helpers: unsloth rewrites every - # `if .shape[-1] == lm_head.shape[1]:` in this function and refuses any other spelling. - def _hidden_logps_and_entropies(hidden_states, ids, chunks): - logps, entropies = _entropy_log_softmax( - hidden_states, lm_head, ids, chunks, - logit_scale_multiply, logit_scale_divide, logit_softcapping, temperature, - lm_head_bias = lm_head_bias, return_entropy = True, - # Entropy reaches the loss only through the entropy bonus; a mask / metric needs no gradient. - entropy_requires_grad = _entropy_bonus_on, - ) - if _support_kwargs: - # Entropies stay full-vocabulary as TRL's; the logprobs take the kept set - logps = chunked_hidden_states_selective_log_softmax( - hidden_states, lm_head, ids, chunks, - logit_scale_multiply, logit_scale_divide, logit_softcapping, temperature, - lm_head_bias = lm_head_bias, - ) - return logps, entropies - - def _raw_logits_entropies(logits): - # The forward already scaled / softcapped these; same temperature as the logprobs. - if _entropy_bonus_on: - log_p = torch.log_softmax(logits.to(torch.float32) / temperature, dim = -1) - return -(log_p.exp() * log_p).sum(-1) - # Mask only (top_entropy_quantile): TRL scores entropies under no_grad, so no graph is kept, - # and rows go in slices so the fp32 copies stay near 64M elements, not the whole [B, T, V]. - with torch.no_grad(): - flat = logits.reshape(-1, logits.shape[-1]) - out = torch.empty(flat.shape[0], dtype = torch.float32, device = logits.device) - step = max(1, (1 << 26) // max(int(flat.shape[-1]), 1)) - for lo in range(0, flat.shape[0], step): - log_p = torch.log_softmax(flat[lo : lo + step].to(torch.float32) / temperature, dim = -1) - out[lo : lo + step] = -(log_p.exp() * log_p).sum(-1) - del log_p - return out.reshape(logits.shape[:-1]) - - import math - total_samples = input_ids.shape[0] - batch_size = math.ceil(total_samples / B) - input_ids_chunks = [] - attention_mask_chunks = [] - completion_ids_chunks = [] - 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]) - completion_ids_chunks.append(completion_input_ids[start:end]) - - # Shared with the no-grad pass, so the two cannot slice the same tensors differently. - vision_chunks = _grpo_vision_chunks(vision_inputs, total_samples, batch_size) - - zipped_inputs = zip( - input_ids_chunks, - attention_mask_chunks, - vision_chunks, - completion_ids_chunks, - ) - - # Bound in the body, not at module scope, for the reason spelled out just below: this - # function's source is copied into the generated UnslothGRPOTrainer cache without - # unsloth_zoo's module imports, so a module-level import reaches the import path and - # not the one that actually runs in production. - from contextlib import nullcontext - - if not _autocast_on: - autocaster = nullcontext() - else: - autocaster = torch.amp.autocast(device_type = trainer.model.device.type, dtype = trainer._autocast_dtype) - - # PrefixGrouper grad path. This function's source is copied into the generated - # UnslothGRPOTrainer cache without unsloth_zoo's module imports, so bind names - # inside the body; the prefix_grouper import stays lazy + guarded (circular import, - # may be absent) and a failed import just leaves PG off. - from unsloth_zoo.temporary_patches.common import UNSLOTH_ENABLE_LOGGING - - # Memoize env gate + import once per process on the function object (which survives - # into the cache). Env gate checked first so =0 never imports PG code; () = PG off. - _pg_funcs = getattr(grpo_accumulated_loss, "_pg_funcs", None) - if _pg_funcs is None: - _pg_funcs = () - if os.environ.get("UNSLOTH_GRPO_PREFIX_GROUPER", "1").lower() not in ( - "0", "false", "no", "off", - ): - try: - 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, - ) - _pg_funcs = ( - _pg_build_layout, _pg_enabled_fn, _pg_verify_on, _pg_tol_ok, _PG_TOL_KILL, - ) - except Exception: - _pg_funcs = () - grpo_accumulated_loss._pg_funcs = _pg_funcs - # Skip PG under vLLM (fast_inference=True): rollout dominates the step, so the - # saving is small and the first-use self-verify is net overhead. - _pg_engage = bool(_pg_funcs) and not getattr(trainer, "use_vllm", False) and not _need_entropy - - # ---- PrefixGrouper (GRPO shared-prompt dedup; UNSLOTH_GRPO_PREFIX_GROUPER=0 disables) ---- - # Each prompt's G completions share the prefix; PG forwards it once + the G suffixes - # (FlexAttention shared-prefix mask), cutting G*(P+R) tokens to P+G*R. First-use - # self-verify vs the full-row packed new_logprobs; grads flow through the shared stream - # (prefix grad once = sum of G repeats, identical math). Off/failed/unverified -> - # full-row packed path runs as before. - _pg_result = None - _pg_use = False - _pg_skip_pack = False - _pg_num_gen = getattr(trainer, "num_generations", None) - # Runtime gate; broad except -> engage False. - if _pg_engage and _pg_funcs: - try: - _pg_build_layout, _pg_enabled_fn, _pg_verify_on, _pg_tol_ok, _PG_TOL_KILL = _pg_funcs - # Exclusions: softcap models (gemma2) - the FlexAttention kernel skips - # attn_logit_softcapping; hybrid SSM (FalconH1) and MoE (Qwen3-MoE) - their - # decoders do not thread prefix_seg_info, so state would leak across suffixes. - _pg_cfg = getattr(unwrapped_model, "config", None) - _pg_engage = ( - _pg_enabled_fn() - 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) - and not any( - getattr(_pg_cfg, _pg_a, None) is not None - for _pg_a in ("mamba_d_ssm", "mamba_d_state", "mamba_expand") - ) - and not any( - getattr(_pg_cfg, _pg_a, None) is not None - for _pg_a in ( - "num_experts", "num_experts_per_tok", "num_local_experts", - "n_routed_experts", "moe_intermediate_size", - ) - ) - ) - except Exception: - _pg_engage = False - else: - _pg_engage = False - _pg_layout = None - _pg_trusted = False # signature already verified -> skip the full-row forward this step - if _pg_engage: - try: - _pg_pad_id = trainer.processing_class.pad_token_id - # Build the layout from the left-packed input_ids with the original left-pad - # counts so the prefix/suffix split matches the packed path (_pack_cstart) and - # the verify is apples-to-apples. Cap the PG span at any sliding window, - # mirroring the packed _pack_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_id, _pg_num_gen, left_pad_tokens_per_prompt, - max_segment_cap = _pg_sw, - ) - _pg_unsafe = getattr(unwrapped_model, "_unsloth_prefix_grouper_grad_unsafe", None) - if _pg_unsafe is None: - _pg_unsafe = set() - if _pg_layout is not None and _pg_layout.signature in _pg_unsafe: - _pg_layout = None - elif _pg_layout is not None: - _pg_layout.W = logits_to_keep + max_left_pad - _pg_verified = getattr(unwrapped_model, "_unsloth_prefix_grouper_grad_verified", None) - # trust only if the verified envelope covers this batch's lengths - # (re-verify when T or the longest segment grows) - _pg_T = int(_pg_layout.flat_ids.shape[1]) - _pg_maxseg = int(_pg_layout.position_ids.max()) + 1 - _pg_env = ( - _pg_verified.get(_pg_layout.signature) - 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] - ): - _pg_trusted = True - _pg_skip_pack = True # trusted shape -> skip the full-row forward - except Exception as _pg_err: - _pg_layout = None - _pg_trusted = False - _pg_skip_pack = False - 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 (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: the exact per-row - # result, and it fixes the padded path's left-pad RoPE error. Loss/gradients flow through it. Self- - # verified against the per-row forward (shape/RoPE-aware, re-checked as T grows); falls back if a - # backend ignores packed_seq_lengths. lm_head runs on completion positions only. - new_logprobs = None - _pack_result = None - _pack_use = False - _pack_T = None - # True only while the packed forward and its scatter run, so an OOM elsewhere (setup, per-row verifier) - # is not read as a length limit - _pack_in_packed = False - _pack_enabled = os.environ.get("UNSLOTH_GRPO_SEQ_PACKING", "1").lower() not in ("0", "false", "no", "off") - _pack_ok = getattr(unwrapped_model, "_unsloth_seq_packing_grad_ok", None) - if (_pack_enabled and not _pg_skip_pack and pixel_values is None - and token_type_ids is None and mm_token_type_ids is None and _pack_ok is not False): - try: - _pack_pad_id = trainer.processing_class.pad_token_id - _pack_keep = input_ids != _pack_pad_id - _pack_lengths = _pack_keep.sum(dim = 1) - _pack_lengths_cpu = _pack_lengths.tolist() # single GPU->CPU sync, reused below - _pack_nz_cpu = [_n for _n in _pack_lengths_cpu if _n > 0] - _pack_T = sum(_pack_nz_cpu) - _pack_L = input_ids.shape[1] - _pack_W = logits_to_keep + max_left_pad - _pack_maxseg = max(_pack_nz_cpu) if _pack_nz_cpu else 0 - # sliding-window models lose the per-sequence local window in a packed stream - _pack_sw = getattr(getattr(unwrapped_model, "config", None), "sliding_window", None) - _pack_sw_ok = not (isinstance(_pack_sw, int) and _pack_sw > 0 and _pack_maxseg > _pack_sw) - _pack_unsafe = getattr(unwrapped_model, "_unsloth_seq_packing_grad_unsafe_T", None) - # skip the whole packed forward for a known-unsafe length region (a prior moderate mismatch or a - # packed OOM), before the flatten and the active-row read, so a capped batch costs one sync - _pack_capped = _pack_unsafe is not None and _pack_T >= _pack_unsafe - if not _pack_capped: - _pack_flat_ids = input_ids[_pack_keep].unsqueeze(0) - _pack_active = int((completion_mask.sum(dim = 1) > 0).sum()) - if (not _pack_capped) and _pack_T >= 2 and len(_pack_nz_cpu) > 0 and _pack_sw_ok\ - and (_pack_ok is True or _pack_active >= 2): - _pack_psl = torch.tensor(_pack_nz_cpu, dtype = torch.int32, device = input_ids.device) - # reset 0-based position_ids per segment - _pack_pos = (_pack_keep.cumsum(dim = 1) - 1)[_pack_keep].unsqueeze(0) - _pack_chunks = max(1, total_rows * multiplier) - _pack_nz_idx = _pack_keep.nonzero(as_tuple = False) # [T, 2] = (row, col) - _pack_within = _pack_nz_idx[1:, 0] == _pack_nz_idx[:-1, 0] # [T-1] - # completion start is per-row after left-packing: (L - logits_to_keep) minus that - # row's left-pad (matches create_completion_attention_mask exactly) - _pack_cstart = (_pack_L - logits_to_keep) - left_pad_tokens_per_prompt # [rows] - _pack_ctgt = (_pack_nz_idx[1:, 1] >= _pack_cstart[_pack_nz_idx[1:, 0]]) & _pack_within - _pack_in_packed = True - with autocaster: - # use_cache=False: a KV cache silently disables varlen packing - _pack_hidden = unwrapped_model( - input_ids = _pack_flat_ids, - position_ids = _pack_pos, - packed_seq_lengths = _pack_psl, - use_cache = False, - ).logits - _refresh_lm_head() - # `.logits` carries hidden states only when the forward is the - # Unsloth generated one honouring UNSLOTH_RETURN_HIDDEN_STATES; - # otherwise it is real [T, vocab] logits and the lm_head matmul - # dies. Dispatch on width, as the padded path already does. - _pack_h = _pack_hidden[0, :-1, :][_pack_ctgt].unsqueeze(0) - _pack_tid = _pack_flat_ids[0, 1:][_pack_ctgt].unsqueeze(0) - if _unsloth_grpo_returns_hidden_states(unwrapped_model, _pack_h, lm_head): - _fp32_bind() - if _need_entropy: - _pack_sel, _pack_ent = _hidden_logps_and_entropies(_pack_h, _pack_tid, _pack_chunks) - _pack_sel, _pack_ent = _pack_sel[0], _pack_ent[0] - else: - _pack_sel = chunked_hidden_states_selective_log_softmax( - _pack_h, lm_head, _pack_tid, _pack_chunks, - logit_scale_multiply, logit_scale_divide, logit_softcapping, temperature, - lm_head_bias = lm_head_bias, - )[0] - else: - # Raw logits: the forward already applied scale/softcap. - _pack_sel = chunked_selective_log_softmax( - _pack_h, _pack_tid, - temperature = temperature, chunks = _pack_chunks, - )[0] - if _need_entropy: - _pack_ent = _raw_logits_entropies(_pack_h)[0] - # GPT-OSS offload race guard (matches the padded loop) - device_synchronize() - # scatter each completion logprob back to its (row, col) so [:, -_pack_W:] matches padded - _pack_tgt = (_pack_nz_idx[1:, 0] * _pack_L + _pack_nz_idx[1:, 1])[_pack_ctgt] - _pack_result = torch.zeros( - total_rows * _pack_L, dtype = torch.float32, device = input_ids.device, - ).index_put((_pack_tgt,), _pack_sel.to(torch.float32)).view(total_rows, _pack_L)[:, -_pack_W:] - if _need_entropy: - _pack_entropies = torch.zeros( - total_rows * _pack_L, dtype = torch.float32, device = input_ids.device, - ).index_put((_pack_tgt,), _pack_ent.to(torch.float32)).view(total_rows, _pack_L)[:, -_pack_W:] - _pack_in_packed = False - # trust decision: re-verify when T or the longest segment grows past what was verified - # (a LongRoPE cache switch can change the result) - _pack_vT = int(getattr(unwrapped_model, "_unsloth_seq_packing_grad_verified_T", 0)) - _pack_vS = int(getattr(unwrapped_model, "_unsloth_seq_packing_grad_verified_seg", 0)) - _pack_force_verify = os.environ.get("UNSLOTH_GRPO_SEQ_PACKING_VERIFY", "0") == "1" - if (not _pack_force_verify) and _pack_ok is True and _pack_T <= _pack_vT and _pack_maxseg <= _pack_vS: - _pack_use = True # already verified for this shape - else: - # verify against the per-row clean forward (exact ground truth; no grad, value check) - _pack_ref = torch.zeros_like(_pack_result) - with torch.no_grad(), autocaster: - for _pack_i in range(total_rows): - _pack_ni = _pack_lengths_cpu[_pack_i] - if _pack_ni < 2: continue - _pack_rmask = _pack_keep[_pack_i] - _pack_real = input_ids[_pack_i][_pack_rmask].unsqueeze(0) - _pack_rpos = torch.arange(_pack_ni, device = input_ids.device).unsqueeze(0) - _pack_rh = unwrapped_model(input_ids = _pack_real, position_ids = _pack_rpos, use_cache = False).logits - _refresh_lm_head() - # same width dispatch as the packed call above: this forward - # returns raw logits whenever that one did, and the first - # packed batch always lands here - if _unsloth_grpo_returns_hidden_states(unwrapped_model, _pack_rh, lm_head): - _fp32_bind() - _pack_rsel = chunked_hidden_states_selective_log_softmax( - _pack_rh[:, :-1, :], lm_head, _pack_real[:, 1:], 1, - logit_scale_multiply, logit_scale_divide, logit_softcapping, temperature, - lm_head_bias = lm_head_bias, - )[0] - else: - _pack_rsel = chunked_selective_log_softmax( - _pack_rh[:, :-1, :], _pack_real[:, 1:], - temperature = temperature, chunks = 1, - )[0] - _pack_rcols = _pack_rmask.nonzero(as_tuple = False).squeeze(1)[1:] - (_pack_L - _pack_W) - _pack_rkeep = _pack_rcols >= 0 - _pack_ref[_pack_i, _pack_rcols[_pack_rkeep]] = _pack_rsel[_pack_rkeep].to(torch.float32) - device_synchronize() - # compare over the exact loss-mask region (same mask the loss uses; pure - # create_completion_attention_mask, before any tool_mask is applied) - _pack_cm = create_completion_attention_mask( - input_ids[:, -_pack_W:], left_pad_tokens_per_prompt, max_left_pad, _pack_pad_id - ).float() - _pack_diff = float(((_pack_result.detach() - _pack_ref).abs() * _pack_cm).max()) - if UNSLOTH_ENABLE_LOGGING: - print(f"[Unsloth] GRPO seq-packing (grad) verify: T={_pack_T} maxseg={_pack_maxseg} packed-vs-perrow max|d|={_pack_diff:.4f}", flush = True) - # floor ~0.25 through different kernels; cross-sample contamination is >= 2.4 - if _pack_diff < 7e-1: - unwrapped_model._unsloth_seq_packing_grad_ok = True - # only widen the trusted shape when >= 2 completion rows actually exercised - # cross-sample packing; a < 2 row pass proves nothing, so keep re-verifying - # larger shapes until a real multi-row batch clears them - if _pack_active >= 2: - unwrapped_model._unsloth_seq_packing_grad_verified_T = max(_pack_vT, _pack_T) - unwrapped_model._unsloth_seq_packing_grad_verified_seg = max(_pack_vS, _pack_maxseg) - _pack_ok = True - _pack_use = True - else: - _pack_use = False - if _pack_diff >= 1.5: - # large mismatch = contamination (attention ignores the packed mask, e.g. - # some MoE): disable packing for this model - unwrapped_model._unsloth_seq_packing_grad_ok = False - else: - # moderate mismatch -> likely a length boundary (LongRoPE): mark unsafe but - # keep packing for smaller shapes - unwrapped_model._unsloth_seq_packing_grad_unsafe_T = ( - _pack_T if _pack_unsafe is None else min(_pack_unsafe, _pack_T) - ) - if UNSLOTH_ENABLE_LOGGING: - print(f"[Unsloth] GRPO seq-packing (grad) fell back at T={_pack_T} (diff={_pack_diff:.3f})", flush = True) - except Exception as _pack_err: - # any failure -> drop intermediates, use the padded loop, do not retry - _pack_hidden = None - _pack_sel = None - _pack_result = None - _pack_use = False - if isinstance(_pack_err, torch.cuda.OutOfMemoryError): - torch.cuda.empty_cache() - if isinstance(_pack_err, torch.cuda.OutOfMemoryError) and _pack_in_packed and _pack_T is not None: - # an OOM in the packed forward means this length does not fit, not that packing is wrong - _pack_unsafe = getattr(unwrapped_model, "_unsloth_seq_packing_grad_unsafe_T", None) - _pack_unsafe = _pack_T if _pack_unsafe is None else min(_pack_unsafe, _pack_T) - unwrapped_model._unsloth_seq_packing_grad_unsafe_T = _pack_unsafe - if UNSLOTH_ENABLE_LOGGING: - print(f"[Unsloth] GRPO sequence-packing OOM at T={_pack_T}: packing capped below T={_pack_unsafe}, this batch uses the padded loop: {_pack_err!r}", flush = True) - else: - unwrapped_model._unsloth_seq_packing_grad_ok = False - if UNSLOTH_ENABLE_LOGGING: - print(f"[Unsloth] GRPO sequence-packing disabled (fell back to padded): {_pack_err!r}", flush = True) - # ---- PrefixGrouper resolution + first-use self-verify (grad) ---- - # Verify runs under no_grad, then a separate grad forward builds new_logprobs, - # so no inference tensors are saved for backward. - def _pg_grad_forward(): - _pg_chunks = max(1, total_rows * multiplier) - with autocaster: - _h = 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 - _refresh_lm_head() - # Same width dispatch as the packed path and compute_logprobs_chunk. - # `.logits` carries hidden states only when the forward is the Unsloth - # generated one honouring UNSLOTH_RETURN_HIDDEN_STATES; otherwise it is - # real [T, vocab] logits. extract_logps always calls its helper as - # (hidden, lm_head, ids, chunks, ...), so pass a raw-logits helper with - # that same signature, which skips the lm_head matmul and the scale / - # softcap the forward already applied. - if not _unsloth_grpo_returns_hidden_states(unwrapped_model, _h, lm_head): - def _pg_fn(_pg_h, _pg_lm, _pg_ids, _pg_n, _pg_lsm, _pg_lsd, _pg_lsc, _pg_t): - return chunked_selective_log_softmax( - _pg_h, _pg_ids, temperature = _pg_t, chunks = _pg_n, - ) - else: - _fp32_bind() - _pg_fn = chunked_hidden_states_selective_log_softmax - if lm_head_bias is not None: - from functools import partial - _pg_fn = partial(_pg_fn, lm_head_bias = lm_head_bias) - _pg_lp = _pg_layout.extract_logps( - _h, lm_head, _pg_fn, - _pg_chunks, logit_scale_multiply, logit_scale_divide, - logit_softcapping, temperature, - ) # [total_rows, W] with grad - # GPT-OSS offload race guard - device_synchronize() - return _pg_lp - - if _pg_layout is not None: - # A verify-phase OOM (packed graph co-resident) does not prove PG alone cannot fit; - # only an OOM after the packed graph is freed is worth marking unsafe. - _pg_phase_verify = False - try: - if not _pg_trusted: - # first use: verify vs the packed new_logprobs. < tol_ok -> trust; - # >= TOL_KILL -> unsafe forever; borderline -> fall back this shape. - if _pack_use and _pack_result is not None: - _pg_phase_verify = True # packed graph still co-resident - with torch.no_grad(): - _pg_ref = _pg_grad_forward() - _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, - trainer.processing_class.pad_token_id, - ).float() - _pg_a = _pg_ref[:, -_pg_W2:].float() - _pg_b = _pack_result.detach()[:, -_pg_W2:].float() - _pg_diff = float(((_pg_a - _pg_b).abs() * _pg_cm).max()) - if UNSLOTH_ENABLE_LOGGING: - print( - f"[Unsloth] GRPO PrefixGrouper (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_grad_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_grad_verified = _pg_v - _pg_trusted = True - else: - _pg_u = getattr(unwrapped_model, "_unsloth_prefix_grouper_grad_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_grad_unsafe = _pg_u - _pg_trusted = False - # else: no packed reference -> cannot verify -> fall back. - if _pg_trusted: - # free the packed graph BEFORE the grad forward: holding both can OOM when - # PG alone would fit, and on PG failure the padded loop recomputes anyway. - _pack_hidden = _pack_sel = _pack_result = None - _pg_phase_verify = False # packed freed: an OOM below is PG-alone - _pg_result = _pg_grad_forward() - _pg_use = True - except Exception as _pg_err2: - _pg_use = False - os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" - # untrust this signature so the next batch runs the packed path again - _pg_v = getattr(unwrapped_model, "_unsloth_prefix_grouper_grad_verified", None) - if isinstance(_pg_v, dict): - _pg_v.pop(_pg_layout.signature, None) - if isinstance(_pg_err2, torch.cuda.OutOfMemoryError): - # mark unsafe only for a PG-alone OOM (deterministic at these lengths); - # a verify-phase OOM (packed co-resident) proves nothing, just retry. - if not _pg_phase_verify: - _pg_u = getattr(unwrapped_model, "_unsloth_prefix_grouper_grad_unsafe", None) - if _pg_u is None: - _pg_u = set() - _pg_u.add(_pg_layout.signature) - unwrapped_model._unsloth_prefix_grouper_grad_unsafe = _pg_u - torch.cuda.empty_cache() - if UNSLOTH_ENABLE_LOGGING: - print(f"[Unsloth] GRPO PrefixGrouper (grad) forward failed -> packed/padded fallback: {_pg_err2!r}", flush = True) - - if _pg_use and _pg_result is not None: - new_logprobs = _pg_result # PrefixGrouper verified -> skip the loop - zipped_inputs = [] - elif _pack_use and _pack_result is not None: - new_logprobs = _pack_result # verified -> skip the loop - if _need_entropy: - all_entropies_list = [_pack_entropies] - zipped_inputs = [] - else: - # packing rejected/unused: drop the packed graph and its index buffers before the padded loop so - # neither co-resides with it - _pack_hidden = _pack_sel = _pack_result = None - _pack_ent = _pack_entropies = None - if _pack_T is not None: - _pack_keep = _pack_lengths = _pack_flat_ids = _pack_pos = _pack_psl = None - _pack_nz_idx = _pack_within = _pack_cstart = _pack_ctgt = _pack_tgt = None - _pack_h = _pack_tid = _pack_ref = _pack_cm = None - - def to_device(tensor, device, non_blocking=True): - if tensor is None: return None - return tensor.to(device, non_blocking=non_blocking) - - def _offload_device_module(tensor_or_device): - # Stream/Event module for the offload copy. torch.cuda is also the HIP - # backend, so ROCm reports is_cuda and needs no branch of its own; XPU has - # its own namespace, matching gradient_checkpointing.py. Anything else - # (CPU, MPS, ...) returns None and takes the pageable copy. - device = getattr(tensor_or_device, "device", tensor_or_device) - if device.type == "cuda": return torch.cuda - if device.type == "xpu": return getattr(torch, "xpu", None) - return None - - class Unsloth_Offloaded_Log_Softmax(torch.autograd.Function): - """Manual gradient checkpointing / CPU offloading for log softmax.""" - @staticmethod - def forward(ctx, hidden_states, lm_head, index, chunks, - logit_scale_multiply, logit_scale_divide, - logit_softcapping, temperature, lm_head_bias = None): - # Detach so we don't keep the graph (and extra memory) on CPU. - detached_hidden_states = hidden_states.detach().contiguous() - ctx.device = hidden_states.device - ctx.copy_event = None - # Backward runs outside autocast; recompute must match forward (torch.utils.checkpoint does the same). - ctx.autocast_kwargs = dict( - device_type = lm_head.device.type, - enabled = torch.is_autocast_enabled(lm_head.device.type), - dtype = torch.get_autocast_dtype(lm_head.device.type), - cache_enabled = torch.is_autocast_cache_enabled(), - ) - - # Always offload: this path only runs when the caller is already memory bound - # (long completions / large batches), so the win is overlapping the copy. - saved_hidden_states = None - device_module = _offload_device_module(detached_hidden_states) - if device_module is not None: - # Async D2H on a side stream; backward MUST wait on copy_event before - # the H2D reload or it races the copy. - try: - pinned_buffer = torch.empty_like(detached_hidden_states, device = "cpu", pin_memory = True) - if pinned_buffer is not None: - current_stream = device_module.current_stream(detached_hidden_states.device) - copy_stream = device_module.Stream(device = detached_hidden_states.device) - copy_stream.wait_stream(current_stream) - with device_module.stream(copy_stream): - pinned_buffer.copy_(detached_hidden_states, non_blocking = True) - # Keeps the GPU storage alive until the side-stream copy finishes. - detached_hidden_states.record_stream(copy_stream) - copy_event = device_module.Event() - copy_event.record(copy_stream) - saved_hidden_states = pinned_buffer - ctx.copy_event = copy_event - except (RuntimeError, OSError, AttributeError): - # Any accelerator that cannot do pinned side-stream copies falls - # back below; correctness never depends on this path. - saved_hidden_states = None - ctx.copy_event = None - if saved_hidden_states is None: - # No accelerator, or the async copy is unavailable: pageable copy. - saved_hidden_states = detached_hidden_states.to("cpu", non_blocking = True) - ctx.saved_hidden_states = saved_hidden_states - # Drop the clone before the log-softmax below. hidden_states is usually a - # [:, :-1, :] slice, so .contiguous() allocated a full copy; holding the - # reference across the forward would keep it resident alongside the chunk - # logits. record_stream still blocks reuse until the D2H lands, so the - # allocator reclaims it mid-compute rather than at the end of forward. - del detached_hidden_states - - ctx.lm_head = lm_head - ctx.lm_head_requires_grad = lm_head.requires_grad - ctx.lm_head_bias = lm_head_bias - ctx.bias_requires_grad = lm_head_bias is not None and lm_head_bias.requires_grad - ctx.index = index - ctx.args = (chunks, logit_scale_multiply, logit_scale_divide, logit_softcapping, temperature) - - bias_kwargs = {} if lm_head_bias is None else {"lm_head_bias": lm_head_bias} - with torch.no_grad(): - output = chunked_hidden_states_selective_log_softmax( - hidden_states, lm_head, index, *ctx.args, **bias_kwargs, - ) - - return output - - @staticmethod - def backward(ctx, grad_output): - if ctx.copy_event is not None: - # The offload copy must land before the H2D reload. - device_module = _offload_device_module(ctx.device) - ctx.copy_event.wait(device_module.current_stream(ctx.device)) - hidden_states = to_device(ctx.saved_hidden_states, ctx.device) - hidden_states.requires_grad_(True) - - lm_head = ctx.lm_head - if ctx.lm_head_requires_grad: - # Recompute against a private leaf. A Tensor.register_hook on the real - # lm_head fires for tensors named in autograd.grad's inputs, so reusing - # it here would run a user's grad mask / scaler once on this local - # gradient and again when the returned gradient reaches lm_head. - lm_head = lm_head.detach().requires_grad_(True) - lm_head_bias = ctx.lm_head_bias - if ctx.bias_requires_grad: - lm_head_bias = lm_head_bias.detach().requires_grad_(True) - index = ctx.index - - bias_kwargs = {} if lm_head_bias is None else {"lm_head_bias": lm_head_bias} - with torch.enable_grad(), torch.autocast(**ctx.autocast_kwargs): - output = chunked_hidden_states_selective_log_softmax( - hidden_states, lm_head, index, *ctx.args, **bias_kwargs, - ) - - # autograd.grad, not backward: backward writes into leaf .grad, which the - # outer AccumulateGrad would then double-count. - inputs = [hidden_states] - if ctx.lm_head_requires_grad: - inputs.append(lm_head) - if ctx.bias_requires_grad: - inputs.append(lm_head_bias) - grad_inputs = torch.autograd.grad(output, inputs, grad_output) - - return ( - grad_inputs[0], - grad_inputs[1] if ctx.lm_head_requires_grad else None, - None, - None, - None, - None, - None, - None, - grad_inputs[-1] if ctx.bias_requires_grad else None, - ) - - def efficient_log_softmax(hidden_states, lm_head, index, chunks=32, - logit_scale_multiply=0.0, logit_scale_divide=0.0, - logit_softcapping=0.0, temperature=1, batch_size=8, - lm_head_bias=None): - # The old chunked scorer keeps every chunk's logits, so long completions keep the offloaded scorer. - from unsloth_zoo.rl_replacements import grpo_scorer_keeps_logits as _keeps_logits - _plain = os.environ.get("UNSLOTH_GRPO_LOGPROB_RECOMPUTE", "0") != "1" - if not ((index.shape[1] <= 1024 and batch_size <= 8) or batch_size == 1)\ - and (bool(_support_kwargs) or (_plain and _keeps_logits())): - return Unsloth_Offloaded_Log_Softmax.apply( - hidden_states, lm_head, index, chunks, - logit_scale_multiply, logit_scale_divide, - logit_softcapping, temperature, lm_head_bias, - ) - return chunked_hidden_states_selective_log_softmax( - hidden_states, - lm_head, - index, - chunks, - logit_scale_multiply, - logit_scale_divide, - logit_softcapping, - temperature, - lm_head_bias = lm_head_bias, - ) - - def compute_logprobs_chunk(new_hidden_states_chunk, completion_ids, input_ids_chunk): - _refresh_lm_head() - # Hidden states -> lm_head matmul path; raw logits -> skip matmul and - # skip scale/softcap (model forward already applied them). - chunks = input_ids_chunk.shape[0] * multiplier - if _unsloth_grpo_returns_hidden_states(unwrapped_model, new_hidden_states_chunk, lm_head): - _fp32_bind() - if _need_entropy: - logps, entropies = _hidden_logps_and_entropies(new_hidden_states_chunk, completion_ids, chunks) - all_entropies_list.append(entropies) - return logps - return efficient_log_softmax( - new_hidden_states_chunk, - lm_head, - completion_ids, - chunks = chunks, - logit_scale_multiply = logit_scale_multiply, - logit_scale_divide = logit_scale_divide, - logit_softcapping = logit_softcapping, - temperature = temperature, - batch_size = B, - lm_head_bias = lm_head_bias, - ) - if _need_entropy: - all_entropies_list.append(_raw_logits_entropies(new_hidden_states_chunk)) - return chunked_selective_log_softmax( - new_hidden_states_chunk, - completion_ids, - temperature = temperature, - chunks = chunks, - ) - for ( - input_ids_chunk, - attention_mask_chunk, - vision_chunk, - completion_ids - ) in zipped_inputs: - with autocaster: - if pixel_values is None: - new_hidden_states_chunk = unwrapped_model( - input_ids = input_ids_chunk, - attention_mask = attention_mask_chunk, - **vision_chunk, - ).logits - - new_hidden_states_chunk = new_hidden_states_chunk[:, -(logits_to_keep + max_left_pad + 1): , :] - new_hidden_states_chunk = new_hidden_states_chunk[:, :-1, :] - logprobs_chunk = compute_logprobs_chunk(new_hidden_states_chunk, completion_ids, input_ids_chunk) - else: - new_hidden_states_chunk = unwrapped_model( - input_ids = input_ids_chunk, - attention_mask = attention_mask_chunk, - logits_to_keep = logits_to_keep + 1, - **vision_chunk, - ).logits - - new_hidden_states_chunk = new_hidden_states_chunk[:, :-1, :] - logprobs_chunk = compute_logprobs_chunk(new_hidden_states_chunk, completion_ids, input_ids_chunk) - # Avoids race conditions with GPT OSS offload_embbed=True; no measurable slowdown. - device_synchronize() - all_logprobs_list.append(logprobs_chunk) - - if new_logprobs is None: - # padded fallback (packing disabled / unsupported / not verified for this length) - new_logprobs = torch.cat(all_logprobs_list, dim=0) - - if kept_completion_rows is not None: - completion_mask = completion_mask * kept_completion_rows.to( - device = completion_mask.device, dtype = completion_mask.dtype, - ) - - entropy_mask = None - if _need_entropy: - new_entropies = torch.cat(all_entropies_list, dim = 0) - if _entropy_mask_on: - # TRL's own method, so its quantile population, ties and empty-batch rule match per version - # (multi-process TRL < 0.24 takes the padded gather so an empty rank cannot hang the rest). - entropy_mask = _grpo_high_entropy_mask( - trainer, new_entropies.detach(), completion_mask, 1 - _entropy_quantile, - sample_weight = kwargs.get("sample_weight", None), - ) - kwargs["entropy_mask"] = entropy_mask - - # DeepSpeed (no scaler) only: undo training_step's 1/GAS, which transformers >= 5.19 skips when - # loss_is_scaled_for_ga (TRL >= 1.15 GRPO) and older ones when compute_loss_func is set. - upstream_scale = None - if trainer.accelerator.scaler is None and getattr(trainer, "is_deepspeed_enabled", False): - import transformers as _unsloth_transformers - loss_is_scaled_for_ga = None - if hasattr(_unsloth_transformers.Trainer, "loss_is_scaled_for_ga"): - loss_is_scaled_for_ga = getattr(trainer, "loss_is_scaled_for_ga", None) - if loss_is_scaled_for_ga is None: - loss_is_scaled_for_ga = getattr(trainer, "compute_loss_func", None) is not None - upstream_scale = 1.0 - if not loss_is_scaled_for_ga: - upstream_scale = float(kwargs.get("current_gradient_accumulation_steps", 1)) - - _rollout_unscored = None - if _rollout_logps_as_old: - old_logps, _rollout_unscored = grpo_rollout_old_fill_unscored(old_logps, new_logprobs) - # Only on this path: grpo_compute_loss clamps pi_theta / pi_rollout's log to [-20, 20] like verl. - kwargs["rollout_logps_as_old"] = True - # The rollout_rs keep mask grpo_compute_loss applies, for the IS metric weights and the entropy bonus. - _rs_keep = None - if kwargs.get("rollout_rs", None) is not None and old_logps is not None and old_logps.shape == completion_mask.shape: - _rs_log_ratio = None - if _rollout_logps_as_old: - _rs_log_ratio = new_logprobs.detach() - old_logps - elif trainer.use_vllm and sampling_per_token_logps is not None and kwargs.get("vllm_importance_sampling_correction", False): - _rs_log_ratio = old_logps - sampling_per_token_logps - if _rs_log_ratio is not None: - from unsloth_zoo.rl_replacements import grpo_rollout_rs_keep as _grpo_rollout_rs_keep - _rs_keep = _grpo_rollout_rs_keep(_rs_log_ratio, completion_mask, kwargs["rollout_rs"]) - with autocaster: - loss, completion_length, mean_kl, delta, flat_is_ratio, coef_1 = UnslothEfficientGRPO.apply( - new_logprobs, - old_logps, - ref_logps, - sampling_per_token_logps, - lm_head, - completion_input_ids, - completion_mask, - advantages, - trainer.beta, - trainer.accelerator.scaler, - 1, - kwargs, - upstream_scale, - ) - _rollout_stats_args = None - if _rollout_logps_as_old: - # Unclamped log(pi_theta / pi_rollout) for the KL diagnostics; the loss ratio is clamped to [-20, 20]. - _raw_log_ratio = new_logprobs.detach().to(torch.float64) - old_logps.detach().to(torch.float64) - if _raw_log_ratio.shape != completion_mask.shape: - _raw_log_ratio = torch.log(coef_1.detach().to(torch.float64)) - delta, flat_is_ratio = grpo_rollout_as_old_ratio_metrics(coef_1, completion_mask, _raw_log_ratio) - if _rs_keep is not None and flat_is_ratio is not None: - # Rejected tokens get IS weight 0, as on the vLLM IS path: rejected_frac counts them, ESS drops them. - flat_is_ratio = torch.where(_rs_keep, flat_is_ratio, torch.zeros_like(flat_is_ratio)) - if _rollout_metrics_on and flat_is_ratio is not None: - _log_ratio = torch.nan_to_num(_raw_log_ratio, nan = 0.0, posinf = 670.0, neginf = -670.0) - _log_ratio = torch.where(completion_mask.to(torch.bool), _log_ratio, torch.zeros_like(_log_ratio)) - if _rollout_unscored is not None: - # Unscored tokens carry no mismatch information: NaN leaves them out of the sums. - _log_ratio = torch.where(_rollout_unscored, torch.full_like(_log_ratio, float("nan")), _log_ratio) - # An overflowed ratio is capped (finite squares) so ESS tends to its 1 / N limit, not nan. - _ess_weights = torch.clamp(flat_is_ratio.to(torch.float64), max = 1e100) - _rollout_stats_args = (_log_ratio, torch.zeros_like(_log_ratio), completion_mask, _ess_weights) - elif ( - _rollout_metrics_on and old_logps is not None and sampling_per_token_logps is not None - and flat_is_ratio is not None and flat_is_ratio.numel() > 0 - ): - _rollout_stats_args = (old_logps, sampling_per_token_logps, completion_mask, flat_is_ratio) - if _rollout_stats_args is not None: - from unsloth_zoo.rl_replacements import ( - grpo_rollout_mismatch_stats as _grpo_rollout_mismatch_stats, - grpo_rollout_mismatch_metrics as _grpo_rollout_mismatch_metrics, - ) - _rollout_stats = _grpo_rollout_mismatch_stats(*_rollout_stats_args) - _rollout_stats = trainer.accelerator.gather(_rollout_stats.reshape(1, -1)).reshape(-1, _rollout_stats.numel()) - _rollout_stats = torch.cat([_rollout_stats[:, :3].sum(0), _rollout_stats[:, 3:4].amax(0), _rollout_stats[:, 4:].sum(0)]) - _metrics = trainer._metrics - _bucket = "train" if _training else "eval" - if "train" in _metrics: - _metrics = _metrics[_bucket] - # Pool the sums over every micro-batch since the last log instead of averaging - # per-micro-batch ratios: TRL's log() clears _metrics, which restarts the window. - _rollout_acc = getattr(trainer, "_unsloth_rollout_mismatch_stats", None) - if not isinstance(_rollout_acc, dict): - _rollout_acc = {} - trainer._unsloth_rollout_mismatch_stats = _rollout_acc - _prev = _rollout_acc.get(_bucket) - if _prev is not None and "sampling/rollout_kl_k1/mean" in _metrics: - _prev = _prev.to(_rollout_stats.device) - _rollout_stats = torch.cat([ - _prev[:3] + _rollout_stats[:3], - torch.maximum(_prev[3:4], _rollout_stats[3:4]), - _prev[4:] + _rollout_stats[4:], - ]) - _rollout_acc[_bucket] = _rollout_stats.detach() - for _name, _value in _grpo_rollout_mismatch_metrics(_rollout_stats).items(): - _metrics[_name] = [_value] - if _entropy_bonus_on: - _entropy_weights = _grpo_entropy_bonus_weights( - completion_mask, entropy_mask, _entropy_coef, kwargs, training = _training, - ) - if _rs_keep is not None: - # Rejected tokens leave the entropy bonus too (verl drops them from response_mask); - # its normaliser keeps them, like the policy term's. - _entropy_weights = torch.where(_rs_keep.to(_entropy_weights.device), _entropy_weights, 0.0) - _scaler = trainer.accelerator.scaler - loss = loss + _UnslothGRPOEntropyBonus.apply( - new_entropies, _entropy_weights.to(new_entropies.device), - _scaler.get_scale() if _scaler is not None else 1.0, upstream_scale, - ) - finally: - if _unsloth_prior_hidden_states is None: - os.environ.pop("UNSLOTH_RETURN_HIDDEN_STATES", None) - else: - os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = _unsloth_prior_hidden_states - - # Force logits (not hidden states) again or output is gibberish. - os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "0" - - return loss, completion_length, mean_kl, delta, flat_is_ratio, coef_1, completion_mask - -from unsloth_zoo.temporary_patches.utils import torch_compile_with_fallback -@torch_compile_with_fallback(dynamic = True, fullgraph = True, options = torch_compile_options) -def grpo_compute_loss_slow( - ref, - new, - old, - sampling_per_token_logps, - input_ids, - mask, - beta, - advantages, - **kwargs -): - # All Unsloth Zoo code licensed under AGPL3 - # Optional argument defaults. - loss_type = kwargs.get("loss_type", "grpo") - epsilon_low = kwargs.get("epsilon_low", 0.2) - epsilon_high = kwargs.get("epsilon_high", 0.2) - max_completion_length = kwargs.get("max_completion_length", 8192) - delta = kwargs.get("delta", None) - importance_sampling_level = kwargs.get("importance_sampling_level", "token") - num_items_in_batch = kwargs.get("num_items_in_batch", None) - current_gradient_accumulation_steps = kwargs.get("current_gradient_accumulation_steps", 1) - steps_per_generation = kwargs.get("steps_per_generation", None) - num_processes = kwargs.get("num_processes", 1) - use_vllm = kwargs.get("use_vllm", False) - # The off-policy mask uses vLLM sampling logprobs whenever the batch supplies them (matching TRL); - # the vLLM importance-sampling ratio is applied to the loss only when this flag is on. - vllm_importance_sampling_correction = kwargs.get("vllm_importance_sampling_correction", False) - vllm_importance_sampling_mode = kwargs.get("vllm_importance_sampling_mode", "sequence_mask") - vllm_importance_sampling_cap = kwargs.get("vllm_importance_sampling_cap", 2.0) - vllm_importance_sampling_clip_min = kwargs.get("vllm_importance_sampling_clip_min", None) - vllm_importance_sampling_clip_max = kwargs.get("vllm_importance_sampling_clip_max", 3.0) - get_sapo_token_loss = kwargs.get("get_sapo_token_loss", None) - sapo_temperature_pos = kwargs.get("sapo_temperature_pos", 1.0) - sapo_temperature_neg = kwargs.get("sapo_temperature_neg", 1.05) - get_gamma_weights = kwargs.get("get_gamma_weights", None) - vespo_k_pos = kwargs.get("vespo_k_pos", 2.0) - vespo_lambda_pos = kwargs.get("vespo_lambda_pos", 3.0) - vespo_k_neg = kwargs.get("vespo_k_neg", 3.0) - vespo_lambda_neg = kwargs.get("vespo_lambda_neg", 2.0) - get_off_policy_mask = kwargs.get("get_off_policy_mask", None) - off_policy_mask_threshold = kwargs.get("off_policy_mask_threshold", None) - # TRL top_entropy_quantile: the policy term (not the KL) on high-entropy tokens only. - entropy_mask = kwargs.get("entropy_mask", None) - # Parsed (option, lower_log, upper) tuples from grpo_parse_rollout_rs, or None. - rollout_rs = kwargs.get("rollout_rs", None) - # Per-completion weights from unsloth's generation step (ScaleRL zero-variance exclusion / - # prompt-level aggregation), normalised over the whole generation batch; replaces the loss_type normaliser. - sample_weight = kwargs.get("sample_weight", None) - # Only direct callers see this fallback; the trainer always forwards an explicit value. - use_bias_correction_kl = kwargs.get("use_bias_correction_kl", False) - input_ids = input_ids.unsqueeze(-1) - - importance_sampling_ratio = None - # Bool (B, T) keep mask from rollout_rs; None when rejection is off. - rs_keep = None - - # exp(new - old) and exp(ref - new) below are taken before `mask` is applied. A sequence-packed - # logp path leaves the masked (prompt/pad) columns at 0 while a padded one fills them with a real - # logp, so when new and old/ref disagree there those ratios can overflow to inf and inf * 0 (the - # masked-out loss) becomes nan. Force new/old/ref to share 0 on the masked columns so both ratios - # are exp(0) = 1 there; every loss term below multiplies by `mask`, so this changes nothing. - if mask is not None: - _keep = mask.to(torch.bool) - new = torch.where(_keep, new, 0.0) - if old is not None: old = torch.where(_keep, old, 0.0) - if ref is not None: ref = torch.where(_keep, ref, 0.0) - - if advantages.dim() == 1: - advantages = advantages.unsqueeze(1) - - if off_policy_mask_threshold is not None: - # DeepSeek-V3.2 off-policy mask. The mismatch logprobs are sampling_per_token_logps (vLLM - # sampling logprobs) if present, else old, else new.detach() when both are absent - # (num_iterations == 1 with no vLLM). This mirrors TRL, which defaults old_per_token_logps to - # per_token_logps.detach() so get_off_policy_mask never receives None (it computes - # mismatch - per_token_logps.detach(), so new.detach() yields a zero-KL keep-all mask). The - # callable is a signature-stable adapter installed in grpo_accumulated_loss, so this stays - # fixed across TRL versions with no signature introspection inside this compiled function. - off_policy_mask = get_off_policy_mask( - advantages=advantages, - per_token_logps=new, - sampling_per_token_logps=sampling_per_token_logps if sampling_per_token_logps is not None else (old if old is not None else new.detach()), - mask=mask, - off_policy_threshold=off_policy_mask_threshold, - ) - - with torch.no_grad(): - if use_vllm and sampling_per_token_logps is not None and vllm_importance_sampling_correction: - # Filter out extra leading prompt tokens after left-padding input_ids. - # Match TRL: aggregate log-ratios then exp (product), not sum of exp ratios. - importance_sampling_ratio = (old - sampling_per_token_logps) * mask - # Unscored vLLM tokens arrive as nan and nan * 0 survives the mask: ratio 1, as TRL does. - importance_sampling_ratio = torch.nan_to_num(importance_sampling_ratio, nan = 0.0) - - if vllm_importance_sampling_mode in ["sequence_mask", "sequence_truncate"]: - importance_sampling_ratio = importance_sampling_ratio.sum(dim=-1, keepdim=True) - elif vllm_importance_sampling_mode in ["geometric_mask", "geometric_truncate"]: - # Geometric mean of the token ratios (verl seq_mean_k1): the product blows up on long rows. - importance_sampling_ratio = importance_sampling_ratio.sum(dim=-1, keepdim=True) / mask.sum(dim=-1, keepdim=True).clamp(min=1.0) - - importance_sampling_ratio = torch.exp(importance_sampling_ratio) - - if vllm_importance_sampling_mode in ["token_truncate", "sequence_truncate", "geometric_truncate"]: - importance_sampling_ratio = torch.clamp( - importance_sampling_ratio, - min=vllm_importance_sampling_clip_min, - max=vllm_importance_sampling_clip_max - ) - elif vllm_importance_sampling_mode in ["token_mask", "sequence_mask", "geometric_mask"]: - min_val = ( - vllm_importance_sampling_clip_min - if vllm_importance_sampling_clip_min is not None - else -math.inf - ) - - max_val = ( - vllm_importance_sampling_clip_max - if vllm_importance_sampling_clip_max is not None - else math.inf - ) - - invalid_mis_mask = (importance_sampling_ratio < min_val) | ( - importance_sampling_ratio > max_val - ) - - importance_sampling_ratio = importance_sampling_ratio.masked_fill( - invalid_mis_mask, value=0.0 - ) - else: - raise ValueError( - f"Unknown vLLM importance sampling mode: {vllm_importance_sampling_mode}. Possible values are 'token_truncate', 'token_mask', 'sequence_truncate', 'sequence_mask', 'geometric_truncate' and 'geometric_mask'." - ) - - # verl compute_rollout_rejection_mask: every option must pass. verl then replaces - # response_mask with response_mask * keep, so a rejected token leaves the policy term - # and the KL loss (both selected out below). It stays in the loss normaliser like - # TRL's *_mask modes, since num_items_in_batch is counted before the loss. - # Log ratio is pi_old / pi_rollout; with rollout logps as pi_old (verl bypass mode) it is - # pi_theta / pi_rollout, verl compute_policy_loss_bypass_mode passing old_log_prob=log_prob. - rs_log_ratio = None - if rollout_rs is not None: - if use_vllm and sampling_per_token_logps is not None and vllm_importance_sampling_correction: - rs_log_ratio = old - sampling_per_token_logps - elif kwargs.get("rollout_logps_as_old", False) and old is not None: - rs_log_ratio = new.detach() - old - if rs_log_ratio is not None: - rs_log_ratio = torch.clamp(torch.nan_to_num(rs_log_ratio * mask, nan = 0.0), min = -20.0, max = 20.0) - rs_mask = mask.to(rs_log_ratio.dtype) - rs_count = rs_mask.sum(dim=-1, keepdim=True).clamp(min=1.0) - rs_keep = torch.ones_like(rs_log_ratio) - for rs_option, rs_lower, rs_upper in rollout_rs: - if rs_option.endswith("k1"): - rs_stat = -rs_log_ratio - elif rs_option.endswith("k2"): - rs_stat = 0.5 * rs_log_ratio * rs_log_ratio - else: - rs_stat = torch.expm1(rs_log_ratio) - rs_log_ratio - rs_stat = rs_stat * rs_mask - if rs_option.startswith("seq_sum"): - rs_stat = rs_stat.sum(dim=-1, keepdim=True) - elif rs_option.startswith("seq_mean"): - rs_stat = rs_stat.sum(dim=-1, keepdim=True) / rs_count - elif rs_option.startswith("seq_max"): - # k2, k3 >= 0 and masked entries are 0, so this is the masked max (0 on an empty row). - rs_stat = rs_stat.amax(dim=-1, keepdim=True) - if rs_lower is not None: - rs_pass = (rs_stat >= rs_lower) & (rs_stat <= rs_upper) - else: - rs_pass = rs_stat <= rs_upper - rs_keep = rs_keep * rs_pass.to(rs_keep.dtype) - rs_keep = rs_keep > 0 - if importance_sampling_ratio is not None: - # Select, not multiply: an unclipped ratio can overflow to inf and inf * 0 is nan. - importance_sampling_ratio = torch.where(rs_keep, importance_sampling_ratio, 0.0) - pass - - # Must detach when old is None: exp(new - new.detach()) == 1 but keeps grads correct. - if old is not None: - log_ratio = new - old - else: - log_ratio = new - new.detach() - if kwargs.get("rollout_logps_as_old", False): - # Rollout gap can pass 88 nats (float32 exp overflow, 0 * inf backward); verl - # compute_policy_loss_vanilla clamps the log ratio to [-20, 20] too. - log_ratio = torch.clamp(log_ratio, min = -20.0, max = 20.0) - - if importance_sampling_level == "token": - log_importance_weights = log_ratio - elif importance_sampling_level == "sequence": - log_importance_weights = (log_ratio * mask).sum(-1) / mask.sum(-1).clamp(min=1.0) - log_importance_weights = log_importance_weights.unsqueeze(-1) - else: - raise ValueError( - f"Unknown importance sampling level: {importance_sampling_level}. Possible values are 'token' " - "and 'sequence'." - ) - - coef_1 = torch.exp(log_importance_weights) - - # Reverse KL: low-variance low-bias estimator as used in the GRPO paper. - if beta != 0.0: - # Keep expm1 finite, then cap the final KL term (TRL does not clamp). - kl_log_ratio = torch.clamp(ref - new, min = -20.0, max = 20.0) - kl_i = torch.expm1(kl_log_ratio) - kl_log_ratio - # TRL order: pre-clamp non-detached coef_1, before the loss_type dispatch. - if use_bias_correction_kl: - kl_i = kl_i * coef_1 - # Clamp after correction so a saturated KL also stops the ratio's gradient. - kl_i = torch.clamp(kl_i, min = -10.0, max = 10.0) - else: - # Zeros with the correct shape. - if importance_sampling_level == "sequence": - kl_i = new.new_zeros(new.size(0), 1) - else: - kl_i = torch.zeros_like(new) - - if loss_type == "cispo": - clamped_ratios = torch.clamp(coef_1, max=epsilon_high).detach() - loss_i = -clamped_ratios * advantages * new - elif loss_type in ["grpo", "bnpo", "dr_grpo", "dapo", "luspo"]: - coef_2 = torch.clamp(coef_1, 1 - epsilon_low, 1 + epsilon_high) - - if delta is not None: - loss_1 = torch.clamp(coef_1, max=delta) * advantages - else: - loss_1 = coef_1 * advantages - pass - loss_2 = coef_2 * advantages - loss_i = -torch.min(loss_1, loss_2) - elif loss_type == "sapo": - temperatures = torch.where(advantages > 0, sapo_temperature_pos, sapo_temperature_neg) - soft_coef_1 = torch.sigmoid(temperatures * (coef_1 - 1)) * 4 / temperatures - loss_i = -soft_coef_1 * advantages - elif loss_type == "vespo": - if get_gamma_weights is None: - raise Exception("vespo is only available in TRL 0.26.0+") - phi_seq = get_gamma_weights( - advantages=advantages, - log_ratio_per_token=log_ratio, - mask=mask, - importance_sampling_ratio=importance_sampling_ratio, - k_pos=vespo_k_pos, - lambda_pos=vespo_lambda_pos, - k_neg=vespo_k_neg, - lambda_neg=vespo_lambda_neg, - ) - loss_i = -phi_seq * advantages * new - else: - raise ValueError(f"Unknown loss type: {loss_type}") - - if off_policy_mask_threshold is not None: - loss_i = loss_i * off_policy_mask - - if entropy_mask is not None: - loss_i = loss_i * entropy_mask - - if use_vllm and sampling_per_token_logps is not None and vllm_importance_sampling_correction: - # vespo applies the IS ratio inside get_gamma_weights, so skip it here. - if loss_type != "vespo": - loss_i = loss_i * importance_sampling_ratio - # delta for the metric. - with torch.no_grad(): - delta = torch.abs(old - sampling_per_token_logps) - delta = delta * mask - # Zeroed rather than filtered like TRL, to keep the returned shape. - delta = torch.nan_to_num(delta, nan = 0.0, posinf = math.inf) - flat_is_ratio = importance_sampling_ratio * mask - else: - delta = torch.tensor([]).detach() - flat_is_ratio = torch.tensor([]).detach() - if beta != 0.0: - loss_i = loss_i + beta * kl_i - if rs_keep is not None: - # Rejected tokens drop out of every loss term, KL included (verl rollout_rs). A select, so - # an inf or nan there neither reaches the loss nor turns into a nan gradient. - loss_i = torch.where(rs_keep, loss_i, 0.0) - - mask = mask.to(torch.float32) - n_mask_per_reward = mask.sum(1) - - # https://github.com/huggingface/trl/blob/e8b8499f1f8d76838155b515e414ee98f757d6d5/trl/trainer/grpo_trainer.py#L1624 - if sample_weight is not None: - # Same process / accumulation rescale as the dapo branch, whose weights are 1 / num_items_in_batch. - normalizer = 1.0 / num_processes - if steps_per_generation: - normalizer = normalizer * current_gradient_accumulation_steps / steps_per_generation - loss = (loss_i * mask * sample_weight.to(loss_i.dtype).unsqueeze(-1)).sum() / normalizer - elif loss_type in ["grpo", "sapo"]: - loss = ((loss_i * mask).sum(-1) / mask.sum(-1).clamp(min=1.0)).mean() - loss = loss / current_gradient_accumulation_steps - elif loss_type == "bnpo": - # BNPO uses the local micro-batch token count even when TRL supplies a global count. - loss = (loss_i * mask).sum() / mask.sum().clamp(min=1.0) - loss = loss / current_gradient_accumulation_steps - elif loss_type == "dr_grpo": - loss = (loss_i * mask).sum() / (loss_i.size(0) * max_completion_length) - loss = loss / current_gradient_accumulation_steps - elif loss_type in ["cispo", "dapo", "vespo"]: - # Floor at 1 like TRL: a fully masked batch (mask_truncated_completions) is 0/0 = nan otherwise. - if torch.is_tensor(num_items_in_batch): - normalizer = num_items_in_batch.clamp(min = 1.0) / num_processes - else: - normalizer = max(float(num_items_in_batch), 1.0) / num_processes - # num_items_in_batch spans the whole generation batch; rescale to one accumulation window like TRL. - if steps_per_generation: - normalizer = normalizer * current_gradient_accumulation_steps / steps_per_generation - loss = (loss_i * mask).sum() / normalizer - elif loss_type == "luspo": - # loss_i is (B, T) unless sequence level with beta 0, so mask elementwise (TRL >= 1.10). - loss = (loss_i * mask).sum(-1).mean() - normalizer = current_gradient_accumulation_steps - loss = loss / normalizer - else: - raise ValueError(f"Unknown loss type: {loss_type}") - - # Folded metrics. - def masked_batch_mean(x): - with torch.inference_mode(): - completion_length = n_mask_per_reward.mean() - # Sequence level only if x is narrower than mask: one-token completions are (B, 1) too. - if x.shape[1] == 1 and mask.shape[1] != 1: - return completion_length, x.mean() - else: - # Rows with no tokens (mask_truncated_completions) are left out, as in TRL. - mean_kl_per_reward = (x * mask).sum(1) / n_mask_per_reward.clamp(min = 1.0) - kept_rows = (n_mask_per_reward > 0).sum() - mean_kl = torch.where( - kept_rows == n_mask_per_reward.numel(), - mean_kl_per_reward.mean(), - mean_kl_per_reward.sum() / kept_rows.clamp(min = 1), - ) - return completion_length, mean_kl - # Rejected tokens add nothing to the reported KL either; the denominator keeps every completion - # token, like the loss normaliser. - kl_metric = kl_i if rs_keep is None or beta == 0.0 else torch.where(rs_keep, kl_i, 0.0) - completion_length, mean_kl = masked_batch_mean(kl_metric) - return loss, completion_length, mean_kl, delta, flat_is_ratio, coef_1, mask - -def grpo_update_SamplingParams( - SamplingParams, - generation_kwargs, - vllm_sampling_params = None, -): - good_sampling_params_keys = inspect.signature(SamplingParams).parameters.keys() - - new_generation_kwargs = {} - for key in generation_kwargs.keys(): - if key in good_sampling_params_keys: - new_generation_kwargs[key] = generation_kwargs[key] - generation_kwargs = new_generation_kwargs - - if vllm_sampling_params is not None: - overwrites = getattr(vllm_sampling_params, "_set_kwargs", None) - if overwrites is None: - default_sampling_params = SamplingParams() - overwrites = {} - for key in good_sampling_params_keys: - if key.startswith("_") or not hasattr(vllm_sampling_params, key): - continue - overwrited_key = getattr(vllm_sampling_params, key) - if overwrited_key != getattr(default_sampling_params, key, None): - overwrites[key] = overwrited_key - for key, overwrited_key in overwrites.items(): - if key in good_sampling_params_keys and key not in ( - "seed", - "n", - "temperature", - "max_tokens", - "logprobs", - ): - generation_kwargs[key] = overwrited_key - return generation_kwargs - -def _get_inference_mode_context_manager(model: torch.nn.Module): - """A torchao-quantized state dict hits "Cannot set version_counter for inference tensor" on ops like aten.t() under inference mode, a PyTorch bug affecting all tensor subclasses (pytorch/pytorch#164872), so use `torch.no_grad()` in that case and `torch.inference_mode()` otherwise.""" - torchao_config = getattr(model, "torchao_config", None) - if torchao_config is not None and torchao_config.qat_scheme is None: - return torch.no_grad() - # FSDP keeps the parameters it all-gathers for one forward and reuses them in the next, so a - # gather under inference mode poisons the training step: "Inference tensors cannot be saved - # for backward" (unsloth#3551). TRL generates under no_grad for the same reason. - import os, sys - - if os.environ.get("ACCELERATE_USE_FSDP", "").strip().lower() in ("1", "true", "yes", "on"): - return torch.no_grad() - if os.environ.get("FSDP_VERSION", "").strip() not in ("", "0"): - return torch.no_grad() - # Unsloth's native FSDP2 (UNSLOTH_FSDP2=1 under torchrun) exports neither variable, so accelerate - # does not wrap the model again: read its launch flag, and the mark shard_model leaves on the root. - _fsdp2 = os.environ.get("UNSLOTH_FSDP2", "").strip().lower() - if _fsdp2 in ("1", "true", "yes", "on", "native"): - try: - if int(os.environ.get("WORLD_SIZE", "1")) > 1: - return torch.no_grad() - except ValueError: - pass - for _m in (model, getattr(model, "module", None), getattr(model, "_orig_mod", None)): - if _m is not None and getattr(_m, "_unsloth_fsdp2_native", False) is True: - return torch.no_grad() - accelerate_state = sys.modules.get("accelerate.state", None) - shared_state = getattr( - getattr(accelerate_state, "AcceleratorState", None), "_shared_state", None - ) - if isinstance(shared_state, dict) and "FSDP" in str(shared_state.get("distributed_type", "")): - return torch.no_grad() - # ZeRO-3 reuses gathered persistent params (same poisoning); a known stage below 3 does not. - deepspeed = os.environ.get("ACCELERATE_USE_DEEPSPEED", "").strip().lower() in ( - "1", - "true", - "yes", - "on", - ) or ( - isinstance(shared_state, dict) - and "DEEPSPEED" in str(shared_state.get("distributed_type", "")) - ) - if deepspeed: - # Inline, not a helper: rl_replacements copies only this function into the GRPO trainer. - stages = [] - hf_deepspeed = sys.modules.get("transformers.integrations.deepspeed", None) - try: - if hf_deepspeed is not None and hf_deepspeed.is_deepspeed_zero3_enabled(): - stages.append(3) - except Exception: - pass - plugins = shared_state.get("deepspeed_plugins") if isinstance(shared_state, dict) else None - if plugins is not None and not isinstance(plugins, dict): - plugins = {"default": plugins} - for plugin in (plugins or {}).values(): - if not getattr(plugin, "selected", True): - continue - # A config file's stage lives in hf_ds_config; zero_stage defaults to 2 regardless. - stage = getattr(getattr(plugin, "hf_ds_config", None), "_stage", None) - if not isinstance(stage, int) or stage < 0: - stage = getattr(plugin, "zero_stage", None) - stages.append(stage) - env_stage = os.environ.get("ACCELERATE_DEEPSPEED_ZERO_STAGE", "").strip() - if env_stage.isdigit(): - stages.append(int(env_stage)) - stages = [stage for stage in stages if isinstance(stage, int) and stage >= 0] - if not stages or 3 in stages: - return torch.no_grad() - try: - first = next(iter(model.parameters()), None) - except Exception: - first = None - if first is not None and hasattr(first, "ds_id"): - return torch.no_grad() - return torch.inference_mode() - -import os as _unsloth_os -UNSLOTH_ENABLE_LOGGING = _unsloth_os.environ.get('UNSLOTH_ENABLE_LOGGING', '0') in ('1', 'True', 'true') - -UNSLOTH_GRPO_SEQ_PACKING_ON = _unsloth_os.environ.get('UNSLOTH_GRPO_SEQ_PACKING', '1').lower() not in ('0', 'false', 'no', 'off') - -try: - import inspect as _unsloth_inspect - from unsloth_zoo.rl_replacements import RL_REPLACEMENTS as _unsloth_zoo_RL - UNSLOTH_ZOO_HAS_MASKED_COL_GUARD = 'torch.where(_keep, new' in _unsloth_inspect.getsource(_unsloth_zoo_RL['grpo_compute_loss']) -except Exception: - UNSLOTH_ZOO_HAS_MASKED_COL_GUARD = False - -try: - import inspect as _unsloth_inspect - from unsloth_zoo.rl_replacements import chunked_hidden_states_selective_log_softmax as _unsloth_chunked_lsm - UNSLOTH_ZOO_HAS_HEAD_BIAS = 'lm_head_bias' in _unsloth_inspect.signature(_unsloth_chunked_lsm).parameters -except Exception: - UNSLOTH_ZOO_HAS_HEAD_BIAS = False - -_pg_build_layout = _pg_enabled_fn = _pg_verify_on = _pg_tol_ok = _PG_TOL_KILL = None -UNSLOTH_GRPO_PREFIX_GROUPER_ON = _unsloth_os.environ.get('UNSLOTH_GRPO_PREFIX_GROUPER', '1').lower() not in ('0', 'false', 'no', 'off') -if UNSLOTH_GRPO_PREFIX_GROUPER_ON: - try: - 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 - except Exception: - UNSLOTH_GRPO_PREFIX_GROUPER_ON = False - -try: - from unsloth_zoo.device_map_planner import detect_logit_transforms -except Exception: - detect_logit_transforms = None -@dataclass -class UnslothGRPOConfig(GRPOConfig): - """ - - Configuration class for the [`GRPOTrainer`]. - - This class includes only the parameters that are specific to GRPO training. For a full list of training arguments, - please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may - differ from those in [`~transformers.TrainingArguments`]. - - Using [`~transformers.HfArgumentParser`] we can turn this class into - [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the - command line. - - Parameters: - > Parameters that control the model and reference model - - model_init_kwargs (`str` or `dict[str, Any]`, *optional*): - Keyword arguments for [`~transformers.AutoModelForCausalLM.from_pretrained`], used when the `model` - argument of the [`GRPOTrainer`] is provided as a string. The `revision` value is also used when loading - processing classes. - trust_remote_code (`bool`, *optional*, defaults to `False`): - Whether to allow loading models and tokenizers that ship custom Python code from the Hub. Forwarded to - [`~transformers.AutoModelForCausalLM.from_pretrained`] and [`~transformers.AutoProcessor.from_pretrained`]. - Also applied to reward-model and reward-tokenizer loads. - router_aux_loss_coef (`float`, *optional*): - Coefficient of the load-balancing auxiliary loss for Mixture-of-Experts (MoE) models, added to the training - loss with this weight. When not set, the value declared by the model config is used. Set to `0.0` to - disable it. Fails when used with a non-MoE model. - disable_dropout (`bool`, *optional*, defaults to `False`): - Whether to disable dropout in the model. This is useful for training with a reference model, as it prevents - the model from generating different logprobs for the same input. - cast_lm_head_to_fp32 (`bool`, *optional*, defaults to `False`): - Whether to cast the language modeling head of the policy and reference models to float32. As recommended by - the [ScaleRL](https://huggingface.co/papers/2510.13786) recipe. This flag is only supported when the model - has untied word embedding and language modeling head layers i.e. `tie_word_embeddings` in the model config - is False. - - > Parameters that control the data preprocessing - - remove_unused_columns (`bool`, *optional*, defaults to `False`): - Whether to only keep the column `"prompt"` in the dataset. If you use a custom reward function that - requires any column other than `"prompts"` and `"completions"`, you should keep this to `False`. - num_generations (`int`, *optional*, defaults to `8`): - Number of generations per prompt to sample. The effective batch size (num_processes * per_device_batch_size - * gradient_accumulation_steps) must be evenly divisible by this value. - num_generations_eval (`int`, *optional*): - Number of generations to sample during evaluation. This allows using fewer generations during evaluation to - save computation. If `None`, uses the value of `num_generations`. - max_completion_length (`int` or `None`, *optional*, defaults to `512`): - Maximum length of the generated completion. - ds3_gather_for_generation (`bool`, *optional*, defaults to `True`): - This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for generation, - improving generation speed. However, disabling this option allows training models that exceed the VRAM - capacity of a single GPU, albeit at the cost of slower generation. Disabling this option is not compatible - with vLLM generation. - shuffle_dataset (`bool`, *optional*, defaults to `True`): - Whether to shuffle the training dataset. - pad_to_multiple_of (`int`, *optional*): - If set, the prompts ids and completions ids will be padded to a multiple of this value. - - > Parameters that control generation - - generation_batch_size (`int`, *optional*): - Batch size to use for generation. If `None`, it defaults to the effective training batch size: - `per_device_train_batch_size * num_processes * steps_per_generation`. In other words, there is one - generation batch processed per optimization step. Mutually exclusive with `steps_per_generation`. - steps_per_generation (`int`, *optional*): - Number of steps per generation. If `None`, it defaults to `gradient_accumulation_steps`. Mutually exclusive - with `generation_batch_size`. - temperature (`float`, *optional*, defaults to `1.0`): - Temperature for sampling. The higher the temperature, the more random the completions. - top_p (`float`, *optional*, defaults to `1.0`): - Float that controls the cumulative probability of the top tokens to consider. Must be in (0, 1]. Set to - `1.0` to consider all tokens. - top_k (`int`, *optional*, defaults to `0`): - Number of highest probability vocabulary tokens to keep for top-k-filtering. If `0`, top-k-filtering is - disabled and all tokens are considered. - min_p (`float`, *optional*): - Minimum token probability, which will be scaled by the probability of the most likely token. It must be a - value between `0.0` and `1.0`. Typical values are in the `0.01-0.2` range. - generation_kwargs (`dict[str, Any]`, *optional*): - Additional keyword arguments to pass to [`~transformers.GenerationConfig`] (if using transformers) or - `SamplingParams` (if using vLLM) when sampling completions. This can be used to further customize the - generation behavior, such as setting `suppress_tokens`, `num_beams`, etc. If it contains keys that conflict - with the other generation parameters (like `min_p`, `top_p`, etc.), they will override them. - chat_template_kwargs (`dict[str, Any]`, *optional*): - Additional keyword arguments to pass to the `apply_chat_template` function when generating completions. - repetition_penalty (`float`, *optional*, defaults to `1.0`): - Float that penalizes new tokens based on whether they appear in the prompt and the generated text so far. - Values > `1.0` encourage the model to use new tokens, while values < `1.0` encourage the model to repeat - tokens. - cache_implementation (`str`, *optional*): - Implementation of the cache method for faster generation when `use_vllm` is set to `False`. - - > Parameters that control generation acceleration powered by vLLM - - use_vllm (`bool`, *optional*, defaults to `False`): - Whether to use vLLM for generating completions. If set to `True`, the trainer will use vLLM for generation - instead of the default model.generate(). Requires `vllm` to be installed. - vllm_mode (`str`, *optional*, defaults to `"colocate"`): - Mode to use for vLLM integration when `use_vllm` is set to `True`. Must be one of `"server"` or - `"colocate"`. - - - `"server"`: The trainer will send generation requests to a separate vLLM server. Make sure a vLLM server - is running (start with `vllm serve`). - - `"colocate"`: vLLM will run in the same process and share the training GPUs. This avoids the need for a - separate server but may cause resource contention with training. - vllm_model_impl (`str`, *optional*, defaults to `"vllm"`): - Model implementation to use for vLLM. Must be one of `"transformers"` or `"vllm"`. `"transformers"`: Use - the `transformers` backend for model implementation. `"vllm"`: Use the `vllm` library for model - implementation. - vllm_enable_sleep_mode (`bool`, *optional*, defaults to `False`): - Enable vLLM sleep mode to offload weights/cache during the optimizer step. Keeps GPU memory usage low, but - waking the engine adds host–device transfer latency. - vllm_structured_outputs_regex (`str`, *optional*): - Regex for vLLM structured outputs. If `None` (default), structured outputs is disabled. - - > Parameters that control the vLLM server (only used when `vllm_mode` is `"server"`) - - vllm_server_base_url (`str`, *optional*): - Base URL for the vLLM server (e.g., `"http://localhost:8000"`). If provided, `vllm_server_host` and - `vllm_server_port` are ignored. - vllm_server_host (`str`, *optional*, defaults to `"0.0.0.0"`): - Host of the vLLM server to connect to. Ignored if `vllm_server_base_url` is provided. - vllm_server_port (`int`, *optional*, defaults to `8000`): - Port of the vLLM server to connect to. Ignored if `vllm_server_base_url` is provided. - vllm_server_timeout (`float`, *optional*, defaults to `240.0`): - Total timeout duration in seconds to wait for the vLLM server to be up. If the server is not up after the - timeout, a `ConnectionError` is raised. - vllm_group_port (`int`, *optional*, defaults to `51216`): - Port number for the weight update group. This is used to communicate with the vLLM server. Unless the port - is occupied, there is no need to change it. - - > Parameters that control colocated vLLM execution (only used when `vllm_mode` is `"colocate"`) - - vllm_gpu_memory_utilization (`float`, *optional*, defaults to `0.3`): - Control the GPU memory utilization for vLLM. This setting only applies when `vllm_mode` is set to - `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when - launching the vLLM server via the `--vllm_gpu_memory_utilization` flag. - vllm_max_model_length (`int`, *optional*): - Context window for vLLM. Set it to at least the maximum prompt length in the dataset plus - `max_completion_length`; if omitted, it is inferred from the model config. - vllm_tensor_parallel_size (`int`, *optional*, defaults to `1`): - Control the tensor parallel size for vLLM. This setting only applies when `vllm_mode` is set to - `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when - launching the vLLM server via the `--vllm_tensor_parallel_size` flag. - - > Parameters that control the training - - beta (`float`, *optional*, defaults to `0.0`): - KL coefficient. If `0.0` (default), the reference model is not loaded, reducing memory usage and improving - training speed. [DeepSeek-R1 incentivizes reasoning in LLMs through reinforcement - learning](https://huggingface.co/papers/2501.12948) use a value of `0.001`. - num_iterations (`int`, *optional*, defaults to `1`): - Number of iterations per batch (denoted as μ in the algorithm). - epsilon (`float`, *optional*, defaults to `0.2`): - Epsilon value for clipping. - delta (`float`, *optional*): - Enables the upper clipping bound in two-sided GRPO loss when set to a float. If `None` (default), standard - GRPO clipping is used. Recommended to be greater than `1 + ε` when enabled. This method is introduced in - the [INTELLECT-2 tech report](https://huggingface.co/papers/2505.07291). - epsilon_high (`float`, *optional*): - Upper-bound epsilon value for clipping. If not specified, it defaults to the same value as the lower-bound - specified in argument `epsilon`. Paper [DAPO](https://huggingface.co/papers/2503.14476) recommends `0.28`. - When used with `loss_type='cispo'`, this corresponds to the ε_max param specified in the [ScaleRL - paper](https://huggingface.co/papers/2510.13786) and the recommended value is `5.0`. - sapo_temperature_neg (`float`, *optional*, defaults to `1.05`): - Temperature for tokens with non-positive advantage scores used in the `sapo` loss function. This parameter - is introduced in the [Soft Adaptive Policy Optimization paper](https://huggingface.co/papers/2511.20347). - sapo_temperature_pos (`float`, *optional*, defaults to `1.0`): - Temperature for tokens with positive advantage scores used in the `sapo` loss function. This parameter is - introduced in the [Soft Adaptive Policy Optimization paper](https://huggingface.co/papers/2511.20347). - vespo_k_pos (`float`, *optional*, defaults to `2.0`): - k parameter for positive advantages, it is the power exponent in the VESPO loss. Controls how aggressively - we down-weight samples with low importance weights (when the importance sampling ratio < 1). - vespo_lambda_pos (`float`, *optional*, defaults to `3.0`): - lambda parameter for positive advantages, it is the decay factor in the VESPO loss. Controls how - aggressively we down-weight samples with high importance weights (when the importance sampling ratio > 1). - vespo_k_neg (`float`, *optional*, defaults to `3.0`): - k parameter for negative advantages, it is the power exponent in the VESPO loss. Controls how aggressively - we down-weight samples with low importance weights (when the importance sampling ratio < 1). - vespo_lambda_neg (`float`, *optional*, defaults to `2.0`): - lambda parameter for negative advantages, it is the exponential decay factor in the VESPO loss. Controls - how aggressively we down-weight samples with high importance weights (when the importance sampling ratio > - 1). - importance_sampling_level (`str`, *optional*, defaults to `"token"`): - Controls whether importance sampling ratios are computed at the `"token"` or `"sequence"` level. `"token"` - keeps the raw per-token log-probability ratios (one weight per token). `"sequence"` averages the - log-probability ratios across valid tokens to produce a single ratio per sequence. The [GSPO - paper](https://huggingface.co/papers/2507.18071) shows that sequence-level sampling often yields more - stable training and better alignment with sequence-level rewards. - reward_weights (`list[float]`, *optional*): - Weights for each reward function. Must match the number of reward functions. If `None`, all rewards are - weighted equally with weight `1.0`. - multi_objective_aggregation (`str`, *optional*, defaults to `"sum_then_normalize"`): - Method to aggregate multiple reward functions. Supported values are: - - - `"sum_then_normalize"` (default): First sums the weighted rewards from each reward function, then applies - reward scaling/normalization as specified by `scale_rewards` (see `scale_rewards` for details). - - `"normalize_then_sum"`: First normalizes/scales each reward function across generations (within each - group), then sums the normalized rewards using the specified weights. The aggregated reward is then - normalized at the batch level when forming advantages. This is the suggested approach from the paper - [GDPO: Group reward-Decoupled Normalization Policy Optimization for Multi-reward RL - Optimization](https://huggingface.co/papers/2601.05242). - scale_rewards (`str` or `bool`, *optional*, defaults to `"group"`): - Specifies the scaling strategy for rewards. Supported values are: - - - `True` or `"group"` (default): rewards are scaled by the standard deviation within each group, ensuring - unit variance within a group. - - `"batch"`: rewards are scaled by the standard deviation across the entire batch, as recommended in the - [PPO Lite paper](https://huggingface.co/papers/2508.08221). - - `False` or `"none"`: no scaling is applied. The [Dr. GRPO - paper](https://huggingface.co/papers/2503.20783) recommends not scaling rewards, as scaling by the - standard deviation introduces a question-level difficulty bias. - loss_type (`str`, *optional*, defaults to `"dapo"`): - Specifies the loss formulation to use. Supported values are: - - - `"grpo"`: Aggregates token-level losses by normalizing over sequence length. Not recommended due to - length bias—this approach tends to prefer shorter completions with positive advantages and longer ones - with negative advantages. - - `"dr_grpo"`: Aggregates token-level losses by normalizing with a global constant. This method was - introduced in the [Dr. GRPO paper](https://huggingface.co/papers/2503.20783) to eliminate length bias. - The value of the constant corresponds to `max_completion_length`. - - `"dapo"` (default): Aggregates token-level losses by normalizing with the number of active token in the - global accumulated batch. This method was introduced in the [DAPO - paper](https://huggingface.co/papers/2503.14476) to eliminate length bias. - - `"bnpo"`: Aggregates token-level losses by normalizing with the number of active token in the local - batch. Note that normalization is performed over the local batch only, so results may slightly vary - depending on the local batch size, despite a constant effective batch size. When using - `per_device_train_batch_size==1`, the loss is equivalent to the GRPO loss. - - `"cispo"`: Clips the importance sampling weights instead of the advantage scaled importance weights. The - clipped weights are then multiplied with the advantages and policy model's log probs. Individual token - losses are aggregated by normalizing with the number of active tokens in the global accumulated batch. - This method was introduced in the [MiniMax-M1 paper](https://huggingface.co/papers/2506.13585). - - `"sapo"`: Soft Adaptive Policy Optimization loss, as introduced in the [Soft Adaptive Policy Optimization - paper](https://huggingface.co/papers/2511.20347). Replaces hard clipping with a smooth, - temperature-controlled gate that adaptively attenuates off-policy updates while preserving useful - learning signals. - - `"luspo"`: Length-Unbiased Sequence Policy Optimization loss. A sequence-level loss that scales each - sequence's loss by its length. This is a modification of GSPO and requires - `importance_sampling_level="sequence"`. Introduced in the [LUSPO - paper](https://huggingface.co/papers/2602.05261). - - `"vespo"`: Variational Sequence-Level Soft Policy Optimization. Replaces hard clipping with a smooth, - asymmetric Gamma weighting function applied directly to sequence-level importance weights. Introduced in - the [VESPO paper](https://huggingface.co/papers/2602.10693). - mask_truncated_completions (`bool`, *optional*, defaults to `False`): - When enabled, truncated completions are excluded from the loss calculation, preventing them from being - incorrectly penalized and introducing noise during training. According to the - [DAPO](https://huggingface.co/papers/2503.14476) paper, this is a good practice for training stability. - sync_ref_model (`bool`, *optional*, defaults to `False`): - Whether to synchronize the reference model with the active model every `ref_model_sync_steps` steps, using - the `ref_model_mixup_alpha` parameter. This synchronization originates from the - [TR-DPO](https://huggingface.co/papers/2404.09656) paper. - ref_model_mixup_alpha (`float`, *optional*, defaults to `0.6`): - α parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which controls the mix - between the current policy and the previous reference policy during updates. The reference policy is - updated according to the equation: `π_ref = α * π_θ + (1 - α) * π_ref_prev`. To use this parameter, you - must set `sync_ref_model=True`. - ref_model_sync_steps (`int`, *optional*, defaults to `512`): - τ parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which determines how - frequently the current policy is synchronized with the reference policy. To use this parameter, you must - set `sync_ref_model=True`. - top_entropy_quantile (`float`, *optional*, defaults to `1.0`): - ρ parameter from [Beyond the 80/20 Rule](https://huggingface.co/papers/2506.01939). Keeps in the policy - loss term only the top-ρ quantile of tokens by entropy of the probability distribution at each sequence - position, improving results. Range: `[0.0-1.0]`. A value of `0.0` masks all but the highest entropy token; - `1.0` keeps all tokens. The paper recommends a value of `0.2`. If used with - `mask_truncated_completions=True`, only tokens from non-truncated completions are considered. - entropy_coef (`float`, *optional*, defaults to `0.0`): - Coefficient of the entropy regularization term in the loss. A positive value adds an entropy bonus that - encourages exploration by keeping the policy from collapsing to near-deterministic outputs. The bonus is - always the mean per-token entropy regardless of `loss_type`; it is not rescaled to match a loss type's - policy normalization, so `entropy_coef` has the same meaning for every loss type. When - `use_adaptive_entropy=True`, this serves as the initial coefficient and is updated each optimizer step. Has - no effect when set to `0.0` (default). - use_adaptive_entropy (`bool`, *optional*, defaults to `False`): - Whether to use adaptive entropy control, introduced in - [Skywork-OR1](https://huggingface.co/papers/2505.22312). When enabled, the entropy coefficient - `entropy_coef` is updated each optimizer step: incremented by `entropy_coef_delta` when the current entropy - is below `entropy_target`, and decremented otherwise. The coefficient is only applied when entropy is at or - below `entropy_target`. - entropy_coef_min (`float`, *optional*, defaults to `0.0`): - Lower bound for the entropy coefficient when using adaptive entropy control. - entropy_coef_max (`float`, *optional*, defaults to `1.0`): - Upper bound for the entropy coefficient when using adaptive entropy control. - entropy_coef_delta (`float`, *optional*, defaults to `0.005`): - Step size for adjusting the entropy coefficient at each optimizer step during adaptive entropy control. - entropy_target (`float`, *optional*, defaults to `0.2`): - Target mean per-token entropy (in nats) used by adaptive entropy control. The coefficient is only applied - when the current entropy falls at or below this value. Measured over the same token set as the policy loss: - all completion tokens by default, or only the high-entropy subset when `top_entropy_quantile < 1.0`. - Typical language models have per-token entropies in the range 2–10 nats, so the default of `0.2` almost - never triggers regularization (only on near-complete entropy collapse); set it close to the entropy you - observe early in training (logged as the `entropy` metric) so the bonus engages before the policy collapses - (and account for the token subset when using `top_entropy_quantile`). - max_tool_calling_iterations (`int`, *optional*): - Maximum number of tool-calling turns when training an agent. If `None`, there is no limit and generation - stops when the model generates a response turn with no tool calls or when the total response length reaches - `max_model_length`. - vllm_importance_sampling_correction (`bool`, *optional*, defaults to `True`): - Whether to apply Importance Sampling (IS) to correct for the mismatch between vLLM completion logprobs and - recomputed training logprobs. If set to `False`, no IS is applied regardless of - `vllm_importance_sampling_mode`. When `True`, the selected mode determines how the IS ratios are computed - and constrained. - vllm_importance_sampling_mode (`str`, *optional*, defaults to `"sequence_mask"`): - Specifies how Importance Sampling is performed when `vllm_importance_sampling_correction=True`. Possible - values are: - - - `"token_truncate"`: Token-level truncated IS (default). Per-token ratios are clipped to - [C_min, C_max]. - - `"token_mask"`: Token-level masked IS. Per-token ratios outside [C_min, C_max] are set to zero. - - `"sequence_truncate"`: Sequence-level truncated IS. A single sequence ratio is clipped to - [C_min, C_max] and applied to all tokens in the sequence. - - `"sequence_mask"`: Sequence-level masked IS. Sequences with ratios outside [C_min, C_max] are masked - out. - vllm_importance_sampling_clip_max (`float`, *optional*, defaults to `3.0`): - Importance sampling upper bound C_max used by `vllm_importance_sampling_mode`. For `*_truncate` modes, - importance ratios are clipped from above at C_max. For `*_mask` modes, ratios larger than C_max are set to - zero. - vllm_importance_sampling_clip_min (`float`, *optional*): - Importance sampling lower bound C_min used by `vllm_importance_sampling_mode`. For `*_truncate` modes, - ratios are clipped from below at C_min. For `*_mask` modes, ratios below C_min are set to zero. To strictly - mask ratios below C_min without upper bound, set `vllm_importance_sampling_clip_max=None`. - off_policy_mask_threshold (`float`, *optional*): - Threshold for off-policy sequence masking. If `None`, off-policy sequence masking is disabled. When set, - sequences with negative advantages and high KL divergence are masked out to stabilize training. This - parameter corresponds to the `delta` threshold in Equation 9 of the [DeepSeek-V3.2 - paper](https://huggingface.co/papers/2512.02556). It expects a positive value (e.g., 0.5). - use_bias_correction_kl (`bool`, *optional*, defaults to `True`): - Whether to multiply the KL term by the importance sampling ratio, so that the KL gradient becomes the - unbiased reverse-KL gradient, as described in the [DeepSeek-V3.2 - paper](https://huggingface.co/papers/2512.02556). This changes the KL gradient whenever `beta != 0`, - including on-policy: the ratio is differentiable, so it affects the gradient even where its value is - exactly 1. The unbiased reverse-KL property holds for `importance_sampling_level="token"`; with - `"sequence"` a sequence-level weight is broadcast onto the per-token KL. - - > Parameters that control the logging - - log_completions (`bool`, *optional*, defaults to `False`): - Whether to log a sample of (prompt, completion) pairs every `logging_steps` steps. If `rich` is installed, - it prints the sample. If `wandb` and/or `trackio` logging is enabled, it logs it to `wandb` and/or - `trackio`. - log_multimodal (`bool`, *optional*, defaults to `True`): - Whether to log multimodal content (images, videos, etc.) together with completions. Disable this to reduce - log size when using high-resolution multimodal data. - num_completions_to_print (`int`, *optional*): - Number of completions to print with `rich`. If `None`, all completions are logged. - log_unique_prompts (`bool`, *optional*, defaults to `False`): - Whether to log unique prompts. If `True`, only unique prompts are logged. If `False`, all prompts are - logged. - log_completions_hub_repo (`str`, *optional*): - Hugging Face Hub repository to save the completions. Should be a complete repository name like - `'username/reponame'` or `'orgname/reponame'`, or just `'reponame'` in which case the repository will be - created in the currently-logged-in Hugging Face user's namespace. Note that this repository will be public - unless you set `hub_private_repo=True` or your organization's default is to create private repositories. - - > Parameters that control generation acceleration powered by transformers continuous batching - - use_transformers_continuous_batching (`bool`, *optional*, defaults to `False`): - Whether to use transformers' continuous batching engine for generating completions. Requires - `transformers>=5.8.0`. - transformers_continuous_batching_config (`dict`, *optional*): - Keyword arguments for [`~transformers.generation.ContinuousBatchingConfig`]. - - > Deprecated parameters - - use_transformers_paged: - - - - Parameter `use_transformers_paged` is deprecated and will be removed in version v2.0.0. Use - `use_transformers_continuous_batching` instead. - - - - vllm_importance_sampling_cap: - - - - Parameter `vllm_importance_sampling_cap` is deprecated and will be removed in version v2.0.0. Use - `vllm_importance_sampling_clip_max` instead. - - - - > [!NOTE] - > These parameters have default values different from [`~transformers.TrainingArguments`]: - > - `logging_steps`: Defaults to `10` instead of `500`. - > - `gradient_checkpointing`: Defaults to `True` instead of `False`. - > - `bf16`: Defaults to `True` if `fp16` is not set, instead of `False`. - > - `learning_rate`: Defaults to `1e-6` instead of `5e-5`. - - """ - _unsloth_accepts_adamw_flash = True - vllm_sampling_params: Optional[Any] = field( - default = None, - metadata = {'help': 'vLLM SamplingParams'}, - ) - unsloth_num_chunks : Optional[int] = field( - default = -1, - metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'}, - ) - unsloth_logit_chunk_multiplier : Optional[int] = field( - default = None, - metadata = {'help': 'Multiplier for chunked logit computations.'}, - ) - unsloth_grpo_mini_batch : Optional[int] = field( - default = None, - metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'}, - ) - - context_parallel_size : Optional[int] = field( - default = 1, - metadata = {'help': 'Ranks per context parallel group (training forwards only). 1 disables it.'}, - ) - context_parallel_backend : Optional[str] = field( - default = 'ulysses', - metadata = {'help': 'Context parallel attention backend; preference and RL trainers support ulysses.'}, - ) - def __init__( - self, - output_dir = None, - per_device_train_batch_size = 4, - num_train_epochs = 3.0, - max_steps = -1, - learning_rate = 5e-05, - lr_scheduler_type = 'linear', - lr_scheduler_kwargs = None, - warmup_steps = 0.1, - optim = 'adamw_8bit', - optim_args = None, - weight_decay = 0.001, - adam_beta1 = 0.9, - adam_beta2 = 0.999, - adam_epsilon = 1e-08, - optim_target_modules = None, - gradient_accumulation_steps = 2, - average_tokens_across_devices = True, - max_grad_norm = 1.0, - label_smoothing_factor = 0.0, - bf16 = False, - fp16 = False, - bf16_full_eval = False, - fp16_full_eval = False, - tf32 = None, - gradient_checkpointing = True, - gradient_checkpointing_kwargs = None, - torch_compile = False, - torch_compile_backend = None, - torch_compile_mode = None, - use_liger_kernel = False, - liger_kernel_config = None, - use_cache = False, - neftune_noise_alpha = None, - torch_empty_cache_steps = 250, - auto_find_batch_size = False, - logging_strategy = 'steps', - logging_steps = 1, - logging_first_step = False, - log_on_each_node = True, - logging_nan_inf_filter = False, - include_num_input_tokens_seen = False, - log_level = 'passive', - log_level_replica = 'warning', - disable_tqdm = None, - report_to = 'none', - run_name = None, - project = 'huggingface', - trackio_space_id = 'trackio', - eval_strategy = 'no', - eval_steps = None, - eval_delay = 0, - per_device_eval_batch_size = 4, - prediction_loss_only = False, - eval_on_start = False, - eval_do_concat_batches = True, - eval_use_gather_object = False, - eval_accumulation_steps = 2, - batch_eval_metrics = False, - save_only_model = False, - save_strategy = 'steps', - save_steps = 500, - save_on_each_node = False, - save_total_limit = None, - enable_jit_checkpoint = False, - push_to_hub = False, - hub_token = None, - hub_private_repo = None, - hub_model_id = None, - hub_strategy = 'every_save', - hub_always_push = False, - hub_revision = None, - load_best_model_at_end = False, - metric_for_best_model = None, - greater_is_better = None, - ignore_data_skip = False, - restore_callback_states_from_checkpoint = False, - full_determinism = False, - seed = 3407, - data_seed = None, - use_cpu = False, - accelerator_config = None, - parallelism_config = None, - dataloader_drop_last = False, - dataloader_num_workers = 0, - dataloader_pin_memory = True, - dataloader_persistent_workers = False, - dataloader_prefetch_factor = None, - remove_unused_columns = False, - label_names = None, - train_sampling_strategy = 'random', - length_column_name = 'length', - ddp_find_unused_parameters = None, - ddp_bucket_cap_mb = None, - ddp_broadcast_buffers = None, - ddp_backend = None, - ddp_timeout = 1800, - fsdp = None, - fsdp_config = None, - deepspeed = None, - debug = '', - skip_memory_metrics = True, - do_train = False, - do_eval = False, - do_predict = False, - resume_from_checkpoint = None, - warmup_ratio = None, - logging_dir = None, - local_rank = -1, - model_init_kwargs = None, - trust_remote_code = False, - router_aux_loss_coef = 0.0, - disable_dropout = False, - cast_lm_head_to_fp32 = False, - num_generations = 8, - num_generations_eval = None, - max_completion_length = 512, - ds3_gather_for_generation = True, - shuffle_dataset = True, - pad_to_multiple_of = None, - generation_batch_size = None, - steps_per_generation = None, - temperature = 1.0, - top_p = 1.0, - top_k = 0, - min_p = None, - generation_kwargs = {}, - chat_template_kwargs = None, - repetition_penalty = 1.0, - cache_implementation = None, - use_vllm = False, - vllm_mode = 'colocate', - vllm_model_impl = 'vllm', - vllm_enable_sleep_mode = False, - vllm_structured_outputs_regex = None, - vllm_server_base_url = None, - vllm_server_host = '0.0.0.0', - vllm_server_port = 8000, - vllm_server_timeout = 240.0, - vllm_group_port = 51216, - vllm_gpu_memory_utilization = 0.3, - vllm_max_model_length = None, - vllm_tensor_parallel_size = 1, - beta = None, - num_iterations = 1, - epsilon = 0.2, - delta = None, - epsilon_high = None, - sapo_temperature_neg = 1.05, - sapo_temperature_pos = 1.0, - vespo_k_pos = 2.0, - vespo_lambda_pos = 3.0, - vespo_k_neg = 3.0, - vespo_lambda_neg = 2.0, - importance_sampling_level = 'token', - reward_weights = None, - multi_objective_aggregation = 'sum_then_normalize', - scale_rewards = 'group', - loss_type = None, - mask_truncated_completions = None, - sync_ref_model = False, - ref_model_mixup_alpha = 0.6, - ref_model_sync_steps = 512, - top_entropy_quantile = 1.0, - entropy_coef = 0.0, - use_adaptive_entropy = False, - entropy_coef_min = 0.0, - entropy_coef_max = 1.0, - entropy_coef_delta = 0.005, - entropy_target = 0.2, - max_tool_calling_iterations = None, - vllm_importance_sampling_correction = False, - vllm_importance_sampling_mode = 'sequence_mask', - vllm_importance_sampling_clip_max = 3.0, - vllm_importance_sampling_clip_min = None, - off_policy_mask_threshold = None, - use_bias_correction_kl = True, - log_completions = False, - log_multimodal = True, - num_completions_to_print = None, - log_unique_prompts = False, - log_completions_hub_repo = None, - use_transformers_continuous_batching = False, - transformers_continuous_batching_config = None, - use_transformers_paged = False, - vllm_importance_sampling_cap = None, - vllm_sampling_params = None, - unsloth_num_chunks = -1, - unsloth_logit_chunk_multiplier = None, - unsloth_grpo_mini_batch = None, - context_parallel_size = 1,context_parallel_backend = 'ulysses', - **kwargs, - ): - if learning_rate < 1e-7: print(f'Unsloth: Your learning rate of `{learning_rate}` is too small and less than 1e-7! Consider increasing it, otherwise gradient updates will be close to 0!') - if learning_rate > 1: print(f'Unsloth: Your learning rate of `{learning_rate}` is way too larger > 1! Consider decreasing it to 1e-1, otherwise gradient updates will explode!') - _unsloth_fsdp2 = os.environ.get('FSDP_VERSION', '').strip() == '2' - _unsloth_fsdp_cfg = fsdp_config - if fsdp and isinstance(_unsloth_fsdp_cfg, str) and os.path.isfile(_unsloth_fsdp_cfg): - try: - import json as _unsloth_json - with open(_unsloth_fsdp_cfg, encoding = 'utf-8') as _unsloth_f: - _unsloth_fsdp_cfg = _unsloth_json.load(_unsloth_f) - except Exception: - _unsloth_fsdp_cfg = None - if fsdp and _unsloth_fsdp_cfg is None: - _unsloth_fsdp_cfg = {} - if fsdp and isinstance(_unsloth_fsdp_cfg, dict): - _unsloth_fsdp2 = _unsloth_fsdp2 or str(_unsloth_fsdp_cfg.get('fsdp_version', _unsloth_fsdp_cfg.get('version', 1))) == '2' - _unsloth_flash_adamw = str(getattr(optim, 'value', optim)).strip().lower() == 'adamw_flash' - if _unsloth_flash_adamw: optim = 'adamw_torch' - if _unsloth_fsdp2 and str(getattr(optim, 'value', optim)) == 'adamw_8bit': - print("Unsloth: FSDP2 shards parameters as DTensors, which bitsandbytes' 8-bit optimizers cannot step. Switching optim = 'adamw_8bit' to 'adamw_torch_fused'.") - optim = 'adamw_torch_fused' - if num_train_epochs is None: - num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override - if output_dir is None and save_strategy == 'steps' and save_steps == 500: - output_dir = 'unsloth_training_checkpoints' - save_strategy = 'no' - if os.environ.get('UNSLOTH_ENABLE_FLEX_ATTENTION', '0') == '1': - from unsloth_zoo.flex_attention import HAS_FLEX_ATTENTION - if HAS_FLEX_ATTENTION and pad_to_multiple_of is None: - from unsloth_zoo.flex_attention import FLEX_ATTENTION_BLOCK_SIZE - pad_to_multiple_of = FLEX_ATTENTION_BLOCK_SIZE - - _unsloth_default_loss_type = loss_type is None - if _unsloth_default_loss_type: - loss_type = 'dapo' - if beta is None: - print("Unsloth: GRPO now defaults to TRL's loss_type = 'dapo' with beta = 0.0. Set loss_type = 'bnpo', beta = 0.001 for Unsloth's previous default.") - if beta is None: - beta = 0.0 if str(loss_type).lower() in ('dapo', 'dr_grpo') else 0.001 - - if loss_type.lower() == 'dr_grpo': - loss_type = 'dr_grpo' - elif loss_type.lower() == 'dapo': - loss_type = 'dapo' - if loss_type.lower() == 'dr_grpo': - if scale_rewards == None: - scale_rewards = True - elif scale_rewards == True or scale_rewards == 'group': - print('Unsloth: The Dr GRPO paper recommends setting `scale_rewards` to False! Will override. Set it to `None` to keep scaling.') - scale_rewards = False - elif loss_type.lower() == 'dapo' and not _unsloth_default_loss_type: - if mask_truncated_completions is None: - print('Unsloth: The DAPO paper recommends `mask_truncated_completions = True` - we will set it.') - mask_truncated_completions = True - if epsilon_high is None: - print('Unsloth: The DAPO paper recommends `epsilon_high = 0.28` - we will set it.') - epsilon_high = 0.28 - if beta != 0.0: - print(f'[WARNING] Unsloth: The DAPO paper recommends setting `beta = 0.0` to remove the KL term - You have set it to {beta}.') - elif loss_type.lower() == 'cispo': - loss_type = 'cispo' - if epsilon_high is None: - print('Unsloth: CISPO caps the importance sampling weight at `epsilon_high`; the ScaleRL paper recommends `epsilon_high = 5.0` - we will set it.') - epsilon_high = 5.0 - elif epsilon_high < 1.0: - print(f'[WARNING] Unsloth: CISPO caps the importance sampling weight at `epsilon_high` itself, not 1 + epsilon_high, so {epsilon_high} caps every weight below 1. The ScaleRL paper uses 4 to 8.') - if mask_truncated_completions is None: - mask_truncated_completions = False - - if steps_per_generation is None and generation_batch_size is None: - ga = gradient_accumulation_steps - world_size = int(os.environ.get('WORLD_SIZE', '1')) - if (ga * world_size * per_device_train_batch_size) % num_generations != 0: - print('Unsloth: We now expect `per_device_train_batch_size` * `gradient_accumulation_steps` * `world_size` to be a multiple of `num_generations`.\nWe will change the batch size of ' + str(per_device_train_batch_size) + ' to the `num_generations` of ' + str(num_generations)) - per_device_train_batch_size = num_generations - - if eval_steps is not None and eval_strategy != 'steps': - print(f'Unsloth: `eval_steps = {eval_steps}` is ignored because `eval_strategy` is {getattr(eval_strategy, "value", eval_strategy)!r}. Set `eval_strategy = "steps"` to evaluate every `eval_steps` steps.') - - if temperature <= 0: - raise ValueError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.') - elif temperature >= 10: - raise ValueError('Unsloth: Please set a positive non-zero temperature less than 10, since sampling will be quite erratic.') - - if use_vllm and (top_k is None or top_k == 0): top_k = -1 - - # One dict so the filter sees the mirrored parameters AND `**kwargs`: - # filtering kwargs alone would double-bind any argument TRL renamed, - # since the new name is itself a mirrored parameter. - _unsloth_config_arguments = dict( - output_dir = output_dir, - per_device_train_batch_size = per_device_train_batch_size, - num_train_epochs = num_train_epochs, - max_steps = max_steps, - learning_rate = learning_rate, - lr_scheduler_type = lr_scheduler_type, - lr_scheduler_kwargs = lr_scheduler_kwargs, - warmup_steps = warmup_steps, - optim = optim, - optim_args = optim_args, - weight_decay = weight_decay, - adam_beta1 = adam_beta1, - adam_beta2 = adam_beta2, - adam_epsilon = adam_epsilon, - optim_target_modules = optim_target_modules, - gradient_accumulation_steps = gradient_accumulation_steps, - average_tokens_across_devices = average_tokens_across_devices, - max_grad_norm = max_grad_norm, - label_smoothing_factor = label_smoothing_factor, - bf16 = bf16, - fp16 = fp16, - bf16_full_eval = bf16_full_eval, - fp16_full_eval = fp16_full_eval, - tf32 = tf32, - gradient_checkpointing = gradient_checkpointing, - gradient_checkpointing_kwargs = gradient_checkpointing_kwargs, - torch_compile = torch_compile, - torch_compile_backend = torch_compile_backend, - torch_compile_mode = torch_compile_mode, - use_liger_kernel = use_liger_kernel, - liger_kernel_config = liger_kernel_config, - use_cache = use_cache, - neftune_noise_alpha = neftune_noise_alpha, - torch_empty_cache_steps = torch_empty_cache_steps, - auto_find_batch_size = auto_find_batch_size, - logging_strategy = logging_strategy, - logging_steps = logging_steps, - logging_first_step = logging_first_step, - log_on_each_node = log_on_each_node, - logging_nan_inf_filter = logging_nan_inf_filter, - include_num_input_tokens_seen = include_num_input_tokens_seen, - log_level = log_level, - log_level_replica = log_level_replica, - disable_tqdm = disable_tqdm, - report_to = report_to, - run_name = run_name, - project = project, - trackio_space_id = trackio_space_id, - eval_strategy = eval_strategy, - eval_steps = eval_steps, - eval_delay = eval_delay, - per_device_eval_batch_size = per_device_eval_batch_size, - prediction_loss_only = prediction_loss_only, - eval_on_start = eval_on_start, - eval_do_concat_batches = eval_do_concat_batches, - eval_use_gather_object = eval_use_gather_object, - eval_accumulation_steps = eval_accumulation_steps, - batch_eval_metrics = batch_eval_metrics, - save_only_model = save_only_model, - save_strategy = save_strategy, - save_steps = save_steps, - save_on_each_node = save_on_each_node, - save_total_limit = save_total_limit, - enable_jit_checkpoint = enable_jit_checkpoint, - push_to_hub = push_to_hub, - hub_token = hub_token, - hub_private_repo = hub_private_repo, - hub_model_id = hub_model_id, - hub_strategy = hub_strategy, - hub_always_push = hub_always_push, - hub_revision = hub_revision, - load_best_model_at_end = load_best_model_at_end, - metric_for_best_model = metric_for_best_model, - greater_is_better = greater_is_better, - ignore_data_skip = ignore_data_skip, - restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint, - full_determinism = full_determinism, - seed = seed, - data_seed = data_seed, - use_cpu = use_cpu, - accelerator_config = accelerator_config, - parallelism_config = parallelism_config, - dataloader_drop_last = dataloader_drop_last, - dataloader_num_workers = dataloader_num_workers, - dataloader_pin_memory = dataloader_pin_memory, - dataloader_persistent_workers = dataloader_persistent_workers, - dataloader_prefetch_factor = dataloader_prefetch_factor, - remove_unused_columns = remove_unused_columns, - label_names = label_names, - train_sampling_strategy = train_sampling_strategy, - length_column_name = length_column_name, - ddp_find_unused_parameters = ddp_find_unused_parameters, - ddp_bucket_cap_mb = ddp_bucket_cap_mb, - ddp_broadcast_buffers = ddp_broadcast_buffers, - ddp_backend = ddp_backend, - ddp_timeout = ddp_timeout, - fsdp = fsdp, - fsdp_config = fsdp_config, - deepspeed = deepspeed, - debug = debug, - skip_memory_metrics = skip_memory_metrics, - do_train = do_train, - do_eval = do_eval, - do_predict = do_predict, - resume_from_checkpoint = resume_from_checkpoint, - warmup_ratio = warmup_ratio, - logging_dir = logging_dir, - local_rank = local_rank, - model_init_kwargs = model_init_kwargs, - trust_remote_code = trust_remote_code, - router_aux_loss_coef = router_aux_loss_coef, - disable_dropout = disable_dropout, - cast_lm_head_to_fp32 = cast_lm_head_to_fp32, - num_generations = num_generations, - num_generations_eval = num_generations_eval, - max_completion_length = max_completion_length, - ds3_gather_for_generation = ds3_gather_for_generation, - shuffle_dataset = shuffle_dataset, - pad_to_multiple_of = pad_to_multiple_of, - generation_batch_size = generation_batch_size, - steps_per_generation = steps_per_generation, - temperature = temperature, - top_p = top_p, - top_k = top_k, - min_p = min_p, - generation_kwargs = generation_kwargs, - chat_template_kwargs = chat_template_kwargs, - repetition_penalty = repetition_penalty, - cache_implementation = cache_implementation, - use_vllm = use_vllm, - vllm_mode = vllm_mode, - vllm_model_impl = vllm_model_impl, - vllm_enable_sleep_mode = vllm_enable_sleep_mode, - vllm_structured_outputs_regex = vllm_structured_outputs_regex, - vllm_server_base_url = vllm_server_base_url, - vllm_server_host = vllm_server_host, - vllm_server_port = vllm_server_port, - vllm_server_timeout = vllm_server_timeout, - vllm_group_port = vllm_group_port, - vllm_gpu_memory_utilization = vllm_gpu_memory_utilization, - vllm_max_model_length = vllm_max_model_length, - vllm_tensor_parallel_size = vllm_tensor_parallel_size, - beta = beta, - num_iterations = num_iterations, - epsilon = epsilon, - delta = delta, - epsilon_high = epsilon_high, - sapo_temperature_neg = sapo_temperature_neg, - sapo_temperature_pos = sapo_temperature_pos, - vespo_k_pos = vespo_k_pos, - vespo_lambda_pos = vespo_lambda_pos, - vespo_k_neg = vespo_k_neg, - vespo_lambda_neg = vespo_lambda_neg, - importance_sampling_level = importance_sampling_level, - reward_weights = reward_weights, - multi_objective_aggregation = multi_objective_aggregation, - scale_rewards = scale_rewards, - loss_type = loss_type, - mask_truncated_completions = mask_truncated_completions, - sync_ref_model = sync_ref_model, - ref_model_mixup_alpha = ref_model_mixup_alpha, - ref_model_sync_steps = ref_model_sync_steps, - top_entropy_quantile = top_entropy_quantile, - entropy_coef = entropy_coef, - use_adaptive_entropy = use_adaptive_entropy, - entropy_coef_min = entropy_coef_min, - entropy_coef_max = entropy_coef_max, - entropy_coef_delta = entropy_coef_delta, - entropy_target = entropy_target, - max_tool_calling_iterations = max_tool_calling_iterations, - vllm_importance_sampling_correction = vllm_importance_sampling_correction, - vllm_importance_sampling_mode = vllm_importance_sampling_mode, - vllm_importance_sampling_clip_max = vllm_importance_sampling_clip_max, - vllm_importance_sampling_clip_min = vllm_importance_sampling_clip_min, - off_policy_mask_threshold = off_policy_mask_threshold, - use_bias_correction_kl = use_bias_correction_kl, - log_completions = log_completions, - log_multimodal = log_multimodal, - num_completions_to_print = num_completions_to_print, - log_unique_prompts = log_unique_prompts, - log_completions_hub_repo = log_completions_hub_repo, - use_transformers_continuous_batching = use_transformers_continuous_batching, - transformers_continuous_batching_config = transformers_continuous_batching_config, - use_transformers_paged = use_transformers_paged, - vllm_importance_sampling_cap = vllm_importance_sampling_cap,**kwargs) - super().__init__(**_unsloth_filter_config_init_kwargs(GRPOConfig, _unsloth_config_arguments, mirrored_from = __class__)) - self.vllm_sampling_params = vllm_sampling_params - self.unsloth_num_chunks = unsloth_num_chunks - if unsloth_grpo_mini_batch is not None: - if self.generation_batch_size >= unsloth_grpo_mini_batch: - self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch - else: - raise ValueError( - f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, " - f"which is self.per_device_train_batch_size * gradient_accumulation_steps." - ) - self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier - - self.context_parallel_size = context_parallel_size - self.context_parallel_backend = context_parallel_backend - # Unsloth: keep the reentrant checkpoint path - if getattr(self, 'gradient_checkpointing', False): - _gc_kwargs = getattr(self, 'gradient_checkpointing_kwargs', None) or {} - if _gc_kwargs.get('context_fn') is None and not _gc_kwargs.get('debug', False): - _gc_kwargs['use_reentrant'] = True - self.gradient_checkpointing_kwargs = _gc_kwargs - from unsloth.optimizers.flash_adamw import is_flash_adamw_marked, mark_flash_adamw, unmark_flash_adamw - if _unsloth_flash_adamw or is_flash_adamw_marked(self): - mark_flash_adamw(self) - else: - unmark_flash_adamw(self) - -pass - -class _UnslothGRPOTrainer(_BaseTrainer): - """""" - - loss_is_scaled_for_ga = True - - _tag_names = ["trl", "grpo"] - _name = "GRPO" - _paper = { - "title": "DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models", - "id": "2402.03300", - # docstyle-ignore - "citation": textwrap.dedent("""\ - @article{shao2024deepseekmath, - title = {{DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models}}, - author = {Zhihong Shao and Peiyi Wang and Qihao Zhu and Runxin Xu and Junxiao Song and Mingchuan Zhang and Y. K. Li and Y. Wu and Daya Guo}, - year = 2024, - eprint = {arXiv:2402.03300}, - }"""), - } - - def __init__( - self, - model: "str | PreTrainedModel | PeftModel", - reward_funcs: RewardFunc | list[RewardFunc] | None = None, - args: GRPOConfig | None = None, - train_dataset: Dataset | IterableDataset | None = None, - eval_dataset: Dataset - | IterableDataset - | DatasetDict - | IterableDatasetDict - | dict[str, Dataset | IterableDataset] - | None = None, - processing_class: PreTrainedTokenizerBase | ProcessorMixin | None = None, - reward_processing_classes: PreTrainedTokenizerBase | list[PreTrainedTokenizerBase] | None = None, - callbacks: list[TrainerCallback] | None = None, - optimizers: tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None] = (None, None), - quantization_config: "BitsAndBytesConfig | None" = None, - peft_config: "PeftConfig | None" = None, - tools: list[Callable] | None = None, - rollout_func: RolloutFunc | None = None, - environment_factory: EnvironmentFactory | dict[str, EnvironmentFactory] | None = None, - ): - - if hasattr(model, 'vllm_engine') and hasattr(args, 'use_vllm'): - _unsloth_fsdp_cfg = getattr(args, 'fsdp_config', None) - if getattr(args, 'fsdp', None) and isinstance(_unsloth_fsdp_cfg, dict) and str(_unsloth_fsdp_cfg.get('version', _unsloth_fsdp_cfg.get('fsdp_version', 1))) == '2': - raise NotImplementedError('Unsloth: `fast_inference = True` (vLLM) is not supported with FSDP2.\nLoad the model with `fast_inference = False` to train with FSDP2, or launch without FSDP2 to use vLLM.') - if (getattr(args, 'use_vllm', False) == False): - args.use_vllm = True - if getattr(args, 'top_k', -1) is None or getattr(args, 'top_k', -1) == 0: - args.top_k = -1 - args.vllm_mode='colocate' - _unsloth_esm = getattr(getattr(getattr(getattr(model.vllm_engine, 'llm_engine', None), 'vllm_config', None), 'model_config', None), 'enable_sleep_mode', None) - if (_unsloth_esm if _unsloth_esm is not None else os.environ.get('UNSLOTH_VLLM_STANDBY', '0') != '0'): - args.vllm_enable_sleep_mode=True - # Args - if args is None: - model_name = model if isinstance(model, str) else get_config_model_id(model.config) - model_name = model_name.split("/")[-1] - args = GRPOConfig(f"{model_name}-GRPO") - - # Model - # PEFT initializes the adapter weights randomly, so set_seed must be done before creating the model to ensure - # reproducibility. - set_seed(args.seed) - if isinstance(model, str): - model_init_kwargs = dict(args.model_init_kwargs or {}) # copy to avoid mutating model_init_kwargs - if quantization_config is not None: - if "quantization_config" in model_init_kwargs: - raise ValueError( - "You set `quantization_config` both as a trainer argument and in `args.model_init_kwargs`. " - "Please set it in only one place, preferably as a trainer argument." - ) - model_init_kwargs["quantization_config"] = quantization_config - # Distributed training requires device_map=None ["auto" fails] - if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: - model_init_kwargs["device_map"] = None - model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) - model_revision = model_init_kwargs.get("revision") - model = create_model_from_path(model, **model_init_kwargs) - else: - model_revision = None - if args.model_init_kwargs is not None: - logger.warning( - "You passed `model_init_kwargs` to the `GRPOConfig`, but your model is already instantiated. " - "The `model_init_kwargs` will be ignored." - ) - if quantization_config is not None: - logger.warning( - "You passed `quantization_config` to the trainer, but your model is already instantiated. The " - "`quantization_config` will be ignored." - ) - # Non-quantized models do not have the `is_loaded_in_{8,4}bit` attributes, whereas quantized models do - _is_quantized_model = getattr(model, "is_loaded_in_4bit", False) or getattr(model, "is_loaded_in_8bit", False) - - # Processing class - if processing_class is None: - processing_class = AutoProcessor.from_pretrained( - get_config_model_id(model.config), - revision=model_revision, - truncation_side="left", - padding_side="left", - trust_remote_code=args.trust_remote_code, - ) - - if args.use_transformers_continuous_batching and isinstance(processing_class, ProcessorMixin): - raise ValueError( - "`use_transformers_continuous_batching` does not support multimodal models. Use `use_vllm` instead." - ) - - # Handle pad token for processors or tokenizers - if isinstance(processing_class, ProcessorMixin): - self._tokenizer = processing_class.tokenizer - self._is_vlm = True - elif isinstance(processing_class, PreTrainedTokenizerBase): - self._tokenizer = processing_class - self._is_vlm = False - else: - raise TypeError("The `processing_class` must be either a `PreTrainedTokenizerBase` or a `ProcessorMixin`") - - if self._tokenizer.pad_token is None: - self._tokenizer.pad_token = self._tokenizer.eos_token - - # The model must agree with the tokenizer on the pad token from construction, so mirror it onto the model - # configs. - model.config.get_text_config().pad_token_id = self._tokenizer.pad_token_id - model.generation_config.pad_token_id = self._tokenizer.pad_token_id - - # Every eos id the model declares [e.g. an end-of-turn token] ends a completion, not only the tokenizer's - eos_token_ids = model.generation_config.eos_token_id - if eos_token_ids is None: - eos_token_ids = [] - elif isinstance(eos_token_ids, int): - eos_token_ids = [eos_token_ids] - if self._tokenizer.eos_token_id not in eos_token_ids: - eos_token_ids = [self._tokenizer.eos_token_id, *eos_token_ids] - self.eos_token_ids = eos_token_ids - - # Resolve vision placeholder token IDs once. Used by the forward pass to rebuild mm_token_type_ids - # when tool responses inject images into the completion [see _generate forward_kwargs block]. - self._image_pad_token_id = None - self._video_pad_token_id = None - if self._is_vlm: - for candidate in ("<|image_pad|>", "<|image|>"): - tid = self._tokenizer.convert_tokens_to_ids(candidate) - if tid != self._tokenizer.unk_token_id: - self._image_pad_token_id = tid - break - tid = self._tokenizer.convert_tokens_to_ids("<|video_pad|>") - if tid != self._tokenizer.unk_token_id: - self._video_pad_token_id = tid - - # PEFT - if False: - if not is_peft_available(): - raise ImportError( - "You passed `peft_config` but the `peft` library is not installed. " - "Install it with `pip install trl[peft]`." - ) - if not isinstance(peft_config, PeftConfig): - raise TypeError( - f"`peft_config` must be a `peft.PeftConfig` instance (e.g. `peft.LoraConfig`), " - f"got {type(peft_config).__name__}." - ) - if is_peft_model(model): - raise ValueError( - "You passed a `PeftModel` instance together with a `peft_config` to the trainer. Please first merge " - "and unload the existing adapter, save the resulting base model, and then pass that base model along " - "with the new `peft_config` to the trainer." - ) - # Create PEFT model - # ZeRO-3 + PEFT for non-quantized models: - # - PEFT's default autocast_adapter_dtype=True upcasts LoRA adapter params to fp32 even when the base model is bf16. - # - ZeRO-3's _allgather_params_coalesced allocates output buffers using the dtype of the first persistent parameter, - # so mixed-dtype persistent_parameters [bf16 base + fp32 LoRA] cause a TypeError on the first optimizer step. - # - Passing autocast_adapter_dtype=False keeps adapter params in the base model dtype [bf16], fixing the mismatch. - # - This is safe: the fp32 upcast is a QLoRA-specific concern [low-bit quantized base models], not needed for - # non-quantized bf16 training. - # - See: - # - TRL issue: https://github.com/huggingface/trl/issues/6089 - # - Upstream issue: https://github.com/deepspeedai/DeepSpeed/issues/8072 - get_peft_model_kwargs = {} - if args.deepspeed_plugin is not None and args.deepspeed_plugin.zero_stage == 3 and not _is_quantized_model: - get_peft_model_kwargs["autocast_adapter_dtype"] = False - model = get_peft_model(model, peft_config, **get_peft_model_kwargs) - # PEFT initialization logic removed via script for trl >= 1.4.0 - # PEFT + DeepSpeed ZeRO-3 requires reentrant checkpointing. For more details, see - # https://github.com/huggingface/trl/issues/2514#issuecomment-2692152703. - # Can be removed once https://github.com/deepspeedai/DeepSpeed/pull/8130 is merged and released. - if ( - is_peft_model(model) - and args.deepspeed_plugin is not None - and args.deepspeed_plugin.zero_stage == 3 - and args.gradient_checkpointing - ): - args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} - use_reentrant = args.gradient_checkpointing_kwargs.get("use_reentrant") - if use_reentrant is False: - logger.warning( - "You are using PEFT with DeepSpeed ZeRO-3 and gradient checkpointing with `use_reentrant=False`. " - "`use_reentrant` is forced to `True` in this configuration to ensure correct training. To remove " - "this warning, unset `use_reentrant` in `gradient_checkpointing_kwargs` or set it to `True`." - ) - args.gradient_checkpointing_kwargs["use_reentrant"] = True - - # When using gradient checkpointing with PEFT, we need to enable input gradients. transformers.Trainer normally - # handles this, but a bug currently prevents it; see https://github.com/huggingface/transformers/issues/42489 - if is_peft_model(model) and args.gradient_checkpointing: - model.enable_input_require_grads() - - # When using QLoRA, the PEFT adapter weights are converted to bf16 to follow the recommendations from the - # original paper [see https://huggingface.co/papers/2305.14314, paragraph 3]. Normally, this can be done by - # passing `autocast_adapter_dtype=False` to `get_peft_model`, but this option is not yet supported for - # quantized models. See: https://github.com/huggingface/peft/issues/2889 - # The DoRA magnitude vector is excluded: unlike LoRA A/B, its optimizer updates can be smaller than bf16 can - # represent, silently freezing it, see #7268. - if False: - for name, param in model.named_parameters(): - if param.requires_grad and "lora_magnitude_vector" not in name: - param.data = param.data.to(torch.bfloat16) - - # Reward functions - if reward_funcs is None: - reward_funcs = [] - elif not isinstance(reward_funcs, list): - reward_funcs = [reward_funcs] - self.reward_func_names = [] - reward_model_revisions = [None] * len(reward_funcs) - for i, reward_func in enumerate(reward_funcs): - if isinstance(reward_func, str): - model_init_kwargs = args.model_init_kwargs or {} - # Distributed training requires device_map=None ["auto" fails] - if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: - model_init_kwargs["device_map"] = None - model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) - reward_model_revisions[i] = model_init_kwargs.get("revision") - reward_funcs[i] = AutoModelForSequenceClassification.from_pretrained( - reward_func, num_labels=1, **model_init_kwargs - ) - if isinstance(reward_funcs[i], nn.Module): # Use Module over PretrainedModel for compat w/ compiled models - self.reward_func_names.append(get_config_model_id(reward_funcs[i].config).split("/")[-1]) - else: - self.reward_func_names.append(get_callable_name(reward_funcs[i])) - self.reward_funcs = reward_funcs - - # Reward weights - if args.reward_weights is not None: - if len(args.reward_weights) != len(reward_funcs): - raise ValueError( - f"Number of reward weights ({len(args.reward_weights)}) must match number of reward " - f"functions ({len(reward_funcs)})" - ) - self.reward_weights = torch.tensor(args.reward_weights, dtype=torch.float32) - else: - self.reward_weights = torch.ones(len(reward_funcs), dtype=torch.float32) - - # Reward processing class - if reward_processing_classes is None: - reward_processing_classes = [None] * len(reward_funcs) - elif not isinstance(reward_processing_classes, list): - reward_processing_classes = [reward_processing_classes] - if len(reward_processing_classes) != len(reward_funcs): - raise ValueError( - f"The number of reward processing classes ({len(reward_processing_classes)}) must match the number of " - f"reward functions ({len(reward_funcs)})." - ) - - for i, (reward_processing_class, reward_func) in enumerate( - zip(reward_processing_classes, reward_funcs, strict=True) - ): - if isinstance(reward_func, PreTrainedModel): - if reward_processing_class is None: - reward_processing_class = AutoTokenizer.from_pretrained( - get_config_model_id(reward_func.config), - revision=reward_model_revisions[i], - trust_remote_code=args.trust_remote_code, - ) - if reward_processing_class.pad_token_id is None: - reward_processing_class.pad_token = reward_processing_class.eos_token - # The reward model computes the reward for the latest non-padded token in the input sequence. - # So it's important to set the pad token ID to the padding token ID of the processing class. - reward_func.config.get_text_config().pad_token_id = reward_processing_class.pad_token_id - reward_processing_classes[i] = reward_processing_class - - self.reward_processing_classes = reward_processing_classes - - # Rollout function - if rollout_func is not None and os.environ.get("TRL_EXPERIMENTAL_SILENCE", "0") != "1": - warnings.warn( - "You are using 'rollout_func', which is an experimental feature. This API may change or be removed at " - "any time without prior notice. Silence this warning by setting environment variable " - "TRL_EXPERIMENTAL_SILENCE=1.", - UserWarning, - stacklevel=2, - ) - self.rollout_func = rollout_func - if environment_factory is not None and os.environ.get("TRL_EXPERIMENTAL_SILENCE", "0") != "1": - warnings.warn( - "You are using 'environment_factory', which is an experimental feature. This API may change or be " - "removed at any time without prior notice. Silence this warning by setting environment variable " - "TRL_EXPERIMENTAL_SILENCE=1.", - UserWarning, - stacklevel=2, - ) - - # Tools - if tools: - if not Version(transformers.__version__) >= Version("5.0.0"): - raise ImportError( - "Using tools with GRPOTrainer requires transformers version 5.0.0 or higher. Please upgrade " - "transformers with `pip install --upgrade transformers` to use this feature." - ) - if environment_factory: - if not Version(transformers.__version__) >= Version("5.2.0"): - raise ImportError( - "Using `environment_factory` with GRPOTrainer requires transformers version 5.2.0 or higher. " - "Please install transformers from the main branch with `pip install " - "git+https://github.com/huggingface/transformers.git@main` to use this feature." - ) - if tools or environment_factory: - # jmespath is only needed by the legacy `response_schema` parser, which is all transformers < 5.13 ships. - # The new-style `response_template` parser doesn't use it, so don't require it on newer versions. - if not _SUPPORTS_RESPONSE_TEMPLATE and not is_jmespath_available(): - raise ImportError( - "Using tools with GRPOTrainer on transformers below 5.13.0 requires the jmespath library for " - "response parsing. Please install it with `pip install jmespath`, or upgrade transformers to " - "5.13.0 or higher, which doesn't need it." - ) - if not supports_tool_calling(processing_class): - raise ValueError( - "The provided chat template does not support tool calling. The template must be able to render a " - "full tool-calling conversation (user -> assistant with tool_calls -> tool)." - ) - - # Set up the environments and extract their methods to be used as tools. - tools = tools or [] - self._standalone_tools = tools # tools that are not bound to an environment - - # Normalize `environment_factory` to a `{name: factory}` mapping. A single callable is the special case of one - # unnamed [`None`] environment shared by every example; a dict maps the `environment` field of each example to - # its factory. `_multi_environment` records which case we are in: only then is the `environment` field read. - self._multi_environment = isinstance(environment_factory, dict) - if environment_factory is None: - self.environment_factories = None - elif self._multi_environment: - self.environment_factories = environment_factory - else: - self.environment_factories = {None: environment_factory} - - if self.environment_factories is not None: - # The environment an example uses is only known at batch time [it depends on the data]. Here we just probe - # one instance of each environment to validate its `reset` method and extract its tool methods, used to - # render the per-example tool schema in the prompt. Instances are pooled and reused [reset] across batches; - # the probe seeds the pool so it is not wasted. The pool grows only when a batch needs more concurrent - # instances of an environment than have been created so far, preserving the "construct once, reset often" - # contract even when batches mix environments. - self._env_tools = {} # {environment name: tools exposed when this environment is selected} - self._environment_pool = {} # {environment name: list of reusable instances} - # `self.tools` is the union of every environment's tools, accumulated below as each environment is probed. - # Used where only the presence of a tool matters [chat template validation, async loop setup, tool metrics]; - # the per-example schema is rendered in `_tokenize_prompts`. - self.tools = list(tools) - env_reward_types = [] # env classes already given a reward column [dedup: same class under two names] - for name, factory in self.environment_factories.items(): - instance = factory() - has_reset = False - has_reward = False - methods = [] - # List on the class: getmembers on the instance evaluates properties - for member_name, _ in inspect.getmembers(type(instance), predicate=inspect.isfunction): - if member_name == "reset": - has_reset = True - elif member_name == "get_reward": - has_reward = True - elif not member_name.startswith("_"): - methods.append(getattr(instance, member_name)) - if not has_reset: - raise ValueError( - "Each environment instance returned by `environment_factory` must define a callable `reset`." - ) - self._env_tools[name] = tools + methods - self._environment_pool[name] = [instance] - self.tools += [method for method in methods if method not in self.tools] - - # If this environment owns its reward via `get_reward`, expose it as an extra reward source [named after - # the env class, weight 1]. One column per env class that defines `get_reward` [deduplicated, since a - # dict factory may map several names to the same class]; a rollout is scored only when its environment is - # exactly that class, so mixing an env that owns its reward with one that does not is safe [the latter's - # rollouts return `None`, turned into NaN and ignored]. Exact-type match [not `isinstance`] keeps this - # consistent with the async worker and avoids double-counting when one registered env subclasses another. - # `get_reward` may be async [e.g. an LLM judge]; wrap it accordingly so `_calculate_rewards` runs it on - # the daemon event loop like any other async reward func. - if has_reward and type(instance) not in env_reward_types: - env_type = type(instance) - env_reward_types.append(env_type) - if inspect.iscoroutinefunction(instance.get_reward): - - async def get_reward(environments, _env_type=env_type, **kwargs): - return [await e.get_reward() if type(e) is _env_type else None for e in environments] - - else: - - def get_reward(environments, _env_type=env_type, **kwargs): - return [e.get_reward() if type(e) is _env_type else None for e in environments] - - self.reward_funcs.append(get_reward) - self.reward_func_names.append(env_type.__name__) - self.reward_processing_classes.append(None) - self.reward_weights = torch.cat([self.reward_weights, torch.ones(1)]) - else: - self.tools = tools - - # At least one reward source is required: either `reward_funcs`, or an environment that owns the reward via a - # `get_reward` method. - if not self.reward_funcs: - raise ValueError( - "No reward source provided. Pass `reward_funcs`, or an `environment_factory` whose environment " - "defines a `get_reward` method." - ) - - # The per-rollout environment instances and tool dicts both depend on the batch [which environment each example - # selects], so they are built in `_generate_and_score_completions`, right before generation. Only - # `self.environments` needs a default here, because `_calculate_rewards` reads it for every batch [including - # batches with no environments]; the tool dicts are only read in the tool-calling loop, which runs after they - # have been [re]built. - self.environments = None - - # Check for async functions to start an event loop on a daemon thread - self._has_async_funcs = any(is_async_callable(func) for func in self.reward_funcs + self.tools) - - if self._has_async_funcs: - self.async_loop_thread, self.async_loop, self.async_loop_ready_event = start_event_loop_in_daemon( - name="GRPOTrainer-AsyncLoop" - ) - # wait until the event loop is running in the daemon thread - self.async_loop_ready_event.wait() - atexit.register(shutdown_event_loop_in_daemon, self.async_loop_thread, self.async_loop) - - # `add_response_schema` sets the response template [transformers >= 5.13] or legacy schema for known chat - # templates, so tool calls can be parsed. Skip if one is already set; warn if it's a migratable legacy schema. - if self.tools: - has_template = getattr(self._tokenizer, "response_template", None) is not None - has_schema = getattr(self._tokenizer, "response_schema", None) is not None - if not has_template and not has_schema: - processing_class = add_response_schema(processing_class) - elif has_schema and not has_template and _SUPPORTS_RESPONSE_TEMPLATE: - warnings.warn( - "The tokenizer has a legacy `response_schema` set but no `response_template`. The installed " - "transformers supports the new-style `response_template`; consider migrating, as `response_schema` " - "support will eventually be removed. See the Transformers response-parsing docs.", - FutureWarning, - ) - # In multi-turn training, the chat template *must* be prefix-preserving. If the tokenizer's original template - # isn't, we replace it at initialization with a training-safe, prefix-preserving template. - if self.tools and not is_chat_template_prefix_preserving(processing_class): - self.chat_template = get_training_chat_template(processing_class) - else: - self.chat_template = None - - # Training arguments - self.max_completion_length = args.max_completion_length # = |o_i| in the GRPO paper - self.num_generations = args.num_generations # = G in the GRPO paper - self.max_tool_calling_iterations = ( - args.max_tool_calling_iterations if args.max_tool_calling_iterations is not None else sys.maxsize - ) - self.num_generations_eval = args.num_generations_eval or self.num_generations - self.chat_template_kwargs = args.chat_template_kwargs or {} - self.temperature = args.temperature - self.top_p = args.top_p - self.top_k = args.top_k - self.min_p = args.min_p - self.repetition_penalty = args.repetition_penalty - self.use_transformers_continuous_batching = args.use_transformers_continuous_batching - if self.use_transformers_continuous_batching: - if not Version(transformers.__version__) >= Version("5.8.0"): - raise ImportError( - "Using `use_transformers_continuous_batching` requires transformers>=5.8.0. " - "Please upgrade with `pip install --upgrade transformers`." - ) - from transformers.generation import ContinuousBatchingConfig - - cb_kwargs = dict(args.transformers_continuous_batching_config or {}) - # The transformers default [0.9] leaves almost no VRAM for the training backward pass; - # use a training-aware default unless the user has set it explicitly. - cb_kwargs.setdefault("max_memory_percent", 0.5) - self.continuous_batching_config = ContinuousBatchingConfig(**cb_kwargs) - else: - self.continuous_batching_config = None - self.pad_to_multiple_of = args.pad_to_multiple_of - self.use_vllm = args.use_vllm - self.vllm_mode = args.vllm_mode - self.vllm_gpu_memory_utilization = args.vllm_gpu_memory_utilization # only applies to colocation mode - self.vllm_tensor_parallel_size = args.vllm_tensor_parallel_size # only applies to colocation mode - self.vllm_importance_sampling_correction = args.vllm_importance_sampling_correction - self.vllm_importance_sampling_mode = args.vllm_importance_sampling_mode - self.vllm_importance_sampling_clip_max = args.vllm_importance_sampling_clip_max - self.vllm_importance_sampling_clip_min = args.vllm_importance_sampling_clip_min - self.loss_type = args.loss_type - self.multi_objective_aggregation = args.multi_objective_aggregation - - # MoE load-balancing auxiliary loss. `output_router_logits` in the config means the model returns its router - # logits; `router_aux_loss_coef` is the architecture's own coefficient, 0.0 when the config declares none. - text_config = model.config.get_text_config() - coef = args.router_aux_loss_coef - self.router_aux_loss_coef = getattr(text_config, "router_aux_loss_coef", 0.0) if coef is None else coef - if coef and not hasattr(text_config, "output_router_logits"): - raise ValueError( - f"`router_aux_loss_coef` is set to {coef} but {type(model).__name__} is not a Mixture-of-Experts model " - f"that returns its router logits, so there is no auxiliary loss to weight." - ) - self.aux_loss_enabled = hasattr(text_config, "output_router_logits") and self.router_aux_loss_coef != 0.0 - if self.aux_loss_enabled: raise NotImplementedError("Unsloth GRPO does not compute the MoE router auxiliary loss; set router_aux_loss_coef = 0 (the Unsloth default).") - self.scale_rewards = args.scale_rewards - self.importance_sampling_level = args.importance_sampling_level - self.off_policy_mask_threshold = args.off_policy_mask_threshold - if is_peft_model(model) and isinstance(model.get_output_embeddings(), BaseTunerLayer): - # The log-probabilities are computed by multiplying the hidden states by `lm_head.weight` directly, so an - # adapter on the LM head would be ignored and never trained. - raise ValueError( - "Applying a PEFT adapter to `lm_head` is not supported: the log-probabilities are computed from " - "`lm_head.weight` directly, so the adapter would never be trained. Remove `'lm_head'` from your " - "`target_modules`." - ) - self.mask_truncated_completions = args.mask_truncated_completions - self.top_entropy_quantile = args.top_entropy_quantile - self.entropy_coef = args.entropy_coef - self.use_adaptive_entropy = args.use_adaptive_entropy - # Whether the entropy bonus is active. Constant for the run: entropy_coef only mutates in adaptive - # mode, where this is True regardless of its value [so the bonus can recover after being decremented]. - self._entropy_bonus_enabled = self.entropy_coef != 0.0 or self.use_adaptive_entropy - # Cached entropy from the last optimizer step; inf so the first accumulation window - # applies no bonus until a real measurement arrives [conservative default]. - self._last_world_entropy = float("inf") - # Running [entropy_sum, active_token_count] accumulated across the micro-batches of the current - # accumulation window, so the adaptive controller sees the exact window-global entropy [not just the - # last micro-batch]. Reset to None at each optimizer step. - self._entropy_window_stats = None - - # Datasets - self.shuffle_dataset = args.shuffle_dataset - - if train_dataset is None: - # A dataset is optional when an environment owns the data and returns the prompt from `reset[]`; then - # `max_steps` sets the length. Build a placeholder dataset of empty prompts to drive the loop — one - # generation round's worth of rows, cycled across steps. - if self.environment_factories is None: - raise ValueError("`train_dataset` is required unless an `environment_factory` is provided.") - if self._multi_environment: - raise ValueError( - "A `dict` `environment_factory` (multiple environments) requires a `train_dataset` with an " - "`environment` column to route each example to its environment. Provide a dataset, or pass a " - "single environment factory." - ) - if args.max_steps <= 0: - raise ValueError( - "When training without a `train_dataset` (the environment owns the data and returns the prompt " - "from `reset()`), `max_steps` must be set to a positive value to define the training length. Set " - "it via `GRPOConfig(max_steps=...)`." - ) - num_placeholder_rows = args.generation_batch_size // args.num_generations - train_dataset = Dataset.from_dict({"prompt": [[{"role": "user", "content": ""}]] * num_placeholder_rows}) - elif not isinstance(train_dataset, (Dataset, IterableDataset, list, tuple)): - raise TypeError( - f"`train_dataset` must be a `Dataset` or `IterableDataset`, got `{type(train_dataset).__name__}`." - ) - - # Iterable datasets can't be indexed, so the RepeatSampler can't be attached to them. Instead, the sampler's - # ordering is reproduced by streaming [see `get_train_dataloader`/`get_eval_dataloader` and - # `repeat_iterable_dataset`]. This requires `dispatch_batches=False`: with the default dispatch path, batches - # are collated on the main process and Accelerate tries to concatenate the string `prompt` column, which fails; - # `dispatch_batches=False` also lets each process shard the stream into contiguous slices, as the sampler does. - # See https://github.com/huggingface/trl/issues/3213 - uses_iterable_dataset = ( - isinstance(train_dataset, IterableDataset) - or isinstance(eval_dataset, IterableDataset) - or ( - isinstance(eval_dataset, dict) and any(isinstance(ds, IterableDataset) for ds in eval_dataset.values()) - ) - ) - if uses_iterable_dataset: - if args.accelerator_config.dispatch_batches: - raise ValueError( - "Iterable datasets require `dispatch_batches=False`, but it is set to `True` in " - "`accelerator_config`. Please set it to `False`." - ) - args.accelerator_config.dispatch_batches = False - # An iterable train set bakes prompt repeats into the stream, so it must be read by a single worker: multiple - # workers would shard and interleave it, splitting the num_generations groups. Map-style train keeps its workers. - if isinstance(train_dataset, IterableDataset) and args.dataloader_num_workers != 0: - logger.warning( - f"Iterable datasets require `dataloader_num_workers=0` to preserve prompt grouping; overriding the " - f"provided value ({args.dataloader_num_workers})." - ) - args.dataloader_num_workers = 0 - - if args.loss_type == "luspo" and args.importance_sampling_level != "sequence": - logger.warning( - "When using `'luspo'` loss, `importance_sampling_level` should be set to `'sequence'` to mirror the " - "paper's setup." - ) - - if args.loss_type == "vespo" and args.importance_sampling_level != "token": - logger.warning( - "VESPO computes sequence-level importance weights internally. `importance_sampling_level` should be " - "set to `'token'` (the default)." - ) - - if args.importance_sampling_level == "sequence" and args.loss_type in ["bnpo", "dr_grpo", "dapo", "cispo"]: - logger.warning( - f"When using `importance_sampling_level='sequence'`, the `'{args.loss_type}'` loss sums per-token " - "contributions, which effectively weights each sequence by its completion length instead of " - "optimizing the per-sequence objective. To reproduce the GSPO paper's setup, set `loss_type='grpo'` " - "(see https://huggingface.co/docs/trl/main/en/paper_index#group-sequence-policy-optimization]." - ) - - if self.loss_type == "vespo" and self.use_vllm and self.vllm_importance_sampling_correction: - if self.vllm_importance_sampling_mode not in ["token_truncate", "token_mask"]: - raise ValueError( - f"VESPO loss requires `vllm_importance_sampling_mode` to be either 'token_truncate' or " - f"'token_mask'. Got: {self.vllm_importance_sampling_mode}." - ) - - # Multi-step - self.num_iterations = args.num_iterations # = 𝜇 in the GRPO paper - self.epsilon_low = args.epsilon - self.epsilon_high = args.epsilon_high if args.epsilon_high is not None else args.epsilon - # Tracks the number of iterations [forward + backward passes], including those within a grad accum cycle - self._step = 0 - # Buffer the batch to reuse generated outputs across multiple updates. For more details, see - # `_get_train_sampler` and `_prepare_inputs`. - self._buffered_inputs = None - - # Transformers explicitly set use_reentrant=True in the past to silence a PyTorch warning, but the default was - # never updated once PyTorch switched to recommending use_reentrant=False. Until that change lands upstream - # [see https://github.com/huggingface/transformers/pull/43203] and is released [most likely in 5.0.0], we - # default to the recommended non-reentrant behavior here, while preserving any user-provided value. - if args.gradient_checkpointing and Version(transformers.__version__) < Version("5.0.0"): - args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} - args.gradient_checkpointing_kwargs.setdefault("use_reentrant", False) - - super().__init__( - model=model, - args=args, - data_collator=identity, # No data collation is needed in GRPO - train_dataset=train_dataset, - eval_dataset=eval_dataset, - processing_class=processing_class, - callbacks=callbacks, - optimizers=optimizers, - ) - - # With several GPUs visible and no distributed launcher, `Trainer` wraps the model in `nn.DataParallel`, whose - # replicas would all run the forward bound to the original model. Checked after `Trainer.__init__`, which sets - # `n_gpu` to 1 for a model split across devices with `device_map`. - if args.n_gpu > 1: - raise ValueError( - f"{type(self).__name__} does not support `nn.DataParallel`, which `Trainer` uses when several GPUs are " - "visible to a single process. Launch the script with `accelerate launch` or `torchrun`, or make a " - "single GPU visible with `CUDA_VISIBLE_DEVICES`." - ) - - # Reference model - self.beta = args.beta - if self.beta == 0.0: - # If beta is 0.0, the reference model is not needed - self.ref_model = None - elif is_peft_model(model): - # If PEFT is used, the reference model is not needed since the adapter can be disabled - # to revert to the initial model. - self.ref_model = None - else: - # For deepspeed, fsdp or non-distributed models, create a reference model from scratch - model_init_kwargs = dict(args.model_init_kwargs or {}) # copy to avoid mutating model_init_kwargs - if quantization_config is not None: - model_init_kwargs["quantization_config"] = quantization_config - # Distributed training requires device_map=None ["auto" fails] - if self.args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: - model_init_kwargs["device_map"] = None - model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) - self.ref_model = create_model_from_path(get_config_model_id(self.model.config), **model_init_kwargs) - - # Disable dropout in the models - if args.disable_dropout: - disable_dropout_in_model(model) - if self.ref_model is not None: - disable_dropout_in_model(self.ref_model) - - # Cast LM Head To FP32 - if args.cast_lm_head_to_fp32: - - def _cast_lm_head_to_fp32(target_model: PreTrainedModel): - """Cast lm_head to fp32 while preserving embedding output dtype if tied.""" - - if is_peft_available() and isinstance(target_model.lm_head, BaseTunerLayer): - raise ValueError( - "`cast_lm_head_to_fp32=True` is not supported when the lm_head carries a PEFT " - "adapter (e.g. `target_modules` includes `lm_head`). Remove `lm_head` from the " - "adapter's target modules, or set `cast_lm_head_to_fp32=False`." - ) - - original_dtype_local = target_model.lm_head.weight.dtype - lm_head = target_model.lm_head.float() - target_model.lm_head = lm_head - - def cast_forward_to_fp32(hidden_states): - return nn.functional.linear( - hidden_states.to(torch.float32), - lm_head.weight.to(torch.float32), - None if lm_head.bias is None else lm_head.bias.to(torch.float32), - ) - - lm_head.forward = cast_forward_to_fp32 - - if target_model.config.tie_word_embeddings: - - def cast_outputs_to_original_dtype(module, args, output): - return output.to(original_dtype_local) - - # Only cast activations; weights are now fp32 [intentional for numerical stability of logits] - target_model.model.embed_tokens.register_forward_hook(cast_outputs_to_original_dtype) - - _cast_lm_head_to_fp32(model) - if self.ref_model is not None: - _cast_lm_head_to_fp32(self.ref_model) - - # Liger's fused linear cross-entropy replaces `model.forward` when training starts, which would drop the fused - # LM head, so only its layer kernels are applied - if args.use_liger_kernel: - warnings.warn( - "`use_liger_kernel=True` is deprecated and will be removed in v2.0.0. Use the Hub kernels instead, " - 'with `model_init_kwargs={"use_kernels": True}`.', - FutureWarning, - stacklevel=2, - ) - liger_kernel_config = args.liger_kernel_config or {} - if liger_kernel_config.get("fused_linear_cross_entropy"): - raise ValueError( - '`liger_kernel_config={"fused_linear_cross_entropy": True}` is not supported: it replaces the ' - "model's forward, which this trainer patches with a fused LM head." - ) - args.liger_kernel_config = {**liger_kernel_config, "fused_linear_cross_entropy": False} - - # Compute the per-token log-probabilities in chunks, without materializing the full logits - # `entropy` feeds the `entropy` metric, the high-entropy token mask and the entropy bonus - add_fused_lm_head( - self.model.get_base_model() if is_peft_model(self.model) else self.model, - temperature=self.temperature, - cast_lm_head_to_fp32=args.cast_lm_head_to_fp32, - outputs=("log_probs", "entropy"), - ) - if self.ref_model is not None: - add_fused_lm_head( - self.ref_model, temperature=self.temperature, cast_lm_head_to_fp32=args.cast_lm_head_to_fp32 - ) - - # Initialize the metrics - self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} - self._total_train_tokens = 0 - self._current_train_step_time = 0.0 - self.log_completions = args.log_completions - self.log_multimodal = args.log_multimodal - self.log_unique_prompts = args.log_unique_prompts - self.num_completions_to_print = args.num_completions_to_print - # Keep logs sized to the generation batch to record only outputs from the latest model update. - self._logs = { - "images": deque(maxlen=args.generation_batch_size), - "prompt": deque(maxlen=args.generation_batch_size), - "completion": deque(maxlen=args.generation_batch_size), - "rewards": defaultdict(lambda: deque(maxlen=args.generation_batch_size)), - "advantages": deque(maxlen=args.generation_batch_size), - "extra": defaultdict(lambda: deque(maxlen=args.generation_batch_size)), - } - # Buffers for user-logged data from reward functions, flushed after gathering - self._pending_extra_logs = defaultdict(list) - self._pending_metrics = defaultdict(list) - - # Ensure each process receives a unique seed to prevent duplicate completions when generating with - # transformers if num_generations exceeds per_device_train_batch_size. We could skip it if we use vLLM, but - # it's safer to set it in all cases. - set_seed(args.seed, device_specific=True) - - if self.use_vllm: - self.vllm_generation = VLLMGeneration( - model=self.model, - accelerator=self.accelerator, - processing_class=self.processing_class, - mode=args.vllm_mode, - structured_outputs_regex=args.vllm_structured_outputs_regex, - server_base_url=args.vllm_server_base_url, - server_host=args.vllm_server_host, - server_port=args.vllm_server_port, - group_port=args.vllm_group_port, - server_timeout=args.vllm_server_timeout, - tensor_parallel_size=args.vllm_tensor_parallel_size, - gpu_memory_utilization=args.vllm_gpu_memory_utilization, - max_model_length=args.vllm_max_model_length, - max_num_seqs=args.per_device_train_batch_size - * args.vllm_tensor_parallel_size - * args.steps_per_generation, - enable_sleep_mode=args.vllm_enable_sleep_mode, - model_impl=args.vllm_model_impl, - trust_remote_code=args.trust_remote_code, - repetition_penalty=self.repetition_penalty, - temperature=self.temperature, - top_p=self.top_p, - top_k=self.top_k, - min_p=self.min_p, - max_completion_length=self.max_completion_length, - logprobs=0, - generation_kwargs=args.generation_kwargs, - ) - self._last_loaded_step = -1 - else: - generation_kwargs = { - "max_new_tokens": self.max_completion_length, - "do_sample": True, - "pad_token_id": self._tokenizer.pad_token_id, - "bos_token_id": self._tokenizer.bos_token_id, - "eos_token_id": self.eos_token_ids, - "temperature": self.temperature, - "top_p": self.top_p, - "top_k": self.top_k, - "min_p": self.min_p, - "repetition_penalty": self.repetition_penalty, - "cache_implementation": args.cache_implementation, - } - if args.generation_kwargs is not None: - generation_kwargs.update(args.generation_kwargs) - self.generation_config = GenerationConfig(**generation_kwargs, disable_compile=True) - # Keep training-specific generation kwargs to overwrite model's original generation config - self.generation_kwargs = generation_kwargs - - self._dist = DistributedBackend(self.accelerator) - - # Add tags to the model - self.model.add_model_tags(self._tag_names) - - if self.ref_model is not None: - if self.is_deepspeed_enabled: - self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator) - elif self.is_fsdp_enabled: - self.ref_model = prepare_fsdp(self.ref_model, self.accelerator) - else: - self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True) - - if args.sync_ref_model: - if self.beta == 0.0: - raise ValueError( - "You passed `sync_ref_model=True` while `beta=0.0`, which means the reference model is not used " - "during training. Consequently, GRPOTrainer does not create a `ref_model` instance, and there is " - "nothing to synchronize. Please set `sync_ref_model=False`, or set `beta` to a non-zero value." - ) - if is_peft_model(model): - raise NotImplementedError( - "You passed `sync_ref_model=True` while using a PEFT model, which is currently not supported. " - "With PEFT, GRPOTrainer does not keep a separate reference model in memory; instead, it recovers " - "reference behavior by temporarily disabling the adapter. As a result, there is no standalone " - "`ref_model` instance to synchronize. Use `sync_ref_model=False`, or opt for full fine-tuning if " - "you need a synced reference model. If you need `sync_ref_model` to work with PEFT, please open a " - "feature request at https://github.com/huggingface/trl/issues." - ) - self.add_callback(SyncRefModelCallback(ref_model=self.ref_model, accelerator=self.accelerator)) - - for i, reward_func in enumerate(self.reward_funcs): - if isinstance(reward_func, PreTrainedModel): - if self.is_deepspeed_enabled: - self.reward_funcs[i] = prepare_deepspeed(reward_func, self.accelerator) - else: - # set device placement to True to make `prepare_model` move `reward_func` to device when using fsdp - self.reward_funcs[i] = self.accelerator.prepare_model( - reward_func, evaluation_mode=True, device_placement=True - ) - - if self.accelerator.is_main_process and self.log_completions: - os.makedirs(os.path.join(self.args.output_dir, "completions"), exist_ok=True) - if self.args.log_completions_hub_repo is not None: - repo_id = self.args.log_completions_hub_repo - create_repo(repo_id, private=self.args.hub_private_repo, repo_type="dataset", exist_ok=True) - template_path = pkg_resources.files("trl").joinpath("templates/completions_dataset_card.md") - card_data = DatasetCardData( - pretty_name="TRL Completion logs", - tags=["trl", "trl-logs", "completions"], - ) - card = DatasetCard.from_template( - card_data=card_data, - template_path=str(template_path), - repo_id=repo_id, - hub_model_id=self.args.hub_model_id, - ) - card.push_to_hub(repo_id) - self.commit_scheduler = CommitScheduler( - repo_id=repo_id, - repo_type="dataset", - folder_path=f"{self.args.output_dir}/completions", - every=2, # minutes - allow_patterns=["*.parquet"], - ) - - def _set_signature_columns_if_needed(self): - # If `self.args.remove_unused_columns` is True, non-signature columns are removed. - # By default, this method sets `self._signature_columns` to the model's expected inputs (usually, "input_ids" - # and "attention_mask"). In GRPOTrainer, we preprocess data, so using the model's signature columns doesn't - # work. Instead, we set them to the columns expected by the `training_step` method, hence the override. - if self._signature_columns is None: - self._signature_columns = ["prompt", "image", "images"] - - # This method overrides `Trainer.get_train_dataloader` to support our custom batching strategy. - # Instead of returning a standard per-step batch (i.e., `per_device_batch_size), our dataloader loads an - # *generation* batch (i.e., `per_device_batch_size × steps_per_generation`). This allows us to generate completions - # once every steps_per_generation step—rather than once per accumulation step—which is significantly more - # efficient. Thus, `_prepare_inputs` is called with this *generation* batch, and it handles the splitting internally. - # Maintenance note: this method is a copy-paste of the original `Trainer.get_train_dataloader` with two changes: the - # batch size is multiplied by `steps_per_generation`, and iterable datasets are wrapped (see `repeat_iterable_dataset`). - def get_train_dataloader(self): - dataset = self.train_dataset - if isinstance(dataset, IterableDataset): - # Iterable datasets can't be indexed, so RepeatSampler can't be attached. Reproduce its ordering by - # transforming the stream instead (see `repeat_iterable_dataset`). The full permutation done by - # RepeatSampler becomes a buffered shuffle here. - if self.shuffle_dataset: - dataset = dataset.shuffle(seed=self.args.seed) - dataset = repeat_iterable_dataset( - dataset, - mini_repeat_count=self.num_generations, - batch_size=self.args.generation_batch_size // self.num_generations, - repeat_count=self.num_iterations * self.args.steps_per_generation, - ) - return self._get_dataloader( - dataset=dataset, - description="Training", - batch_size=self._train_batch_size * self.args.steps_per_generation, # < this is the change - sampler_fn=self._get_train_sampler, - is_training=True, - ) - - def _get_train_sampler(self, dataset: Dataset | None = None) -> Sampler: - # Returns a sampler that - # 1. ensures each prompt is repeated across multiple processes. This guarantees that identical prompts are - # distributed to different GPUs, allowing rewards to be computed and normalized correctly within each prompt - # group. Using the same seed across processes ensures consistent prompt assignment, preventing discrepancies - # in group formation. - # 2. repeats the batch multiple times to allow reusing generations across multiple updates. Refer to - # _prepare_inputs to see how the generations are stored and reused. - - # In the following figure, the values are the prompt indices. Each row shows the per-step batch - # returned by `_prepare_inputs`; rows within a `steps_per_generation` block are slices of the same - # generated batch. When `num_iterations > 1`, that block is reused for multiple optimization passes - # before regenerating. - # - # | GPU 0 | GPU 1 | - # - # global_step step <-───> num_generations=2 - # <-───────> per_device_train_batch_size=3 - # grad_accum ▲ ▲ 0 0 0 0 1 1 2 2 <- Generate for the first `steps_per_generation` (prompts 0 to 11); store the completions; use the first slice to compute the loss - # =2 ▼ | 0 1 3 3 4 4 5 5 <- Take the stored generations and use the second slice to compute the loss - # | - # | 1 2 6 6 7 7 8 8 <- Take the stored generations and use the third slice to compute the loss - # steps_per_gen=4 ▼ 1 3 9 9 10 10 11 11 <- Take the stored generations and use the fourth slice to compute the loss - # - # 2 4 12 12 13 13 14 14 <- Generate for the second `steps_per_generation` (prompts 12 to 23); store the completions; use the first slice to compute the loss - # 2 5 15 15 16 16 17 17 <- Take the stored generations and use the second slice to compute the loss - # ... - if dataset is None: - dataset = self.train_dataset - return RepeatSampler( - data_source=dataset, - mini_repeat_count=self.num_generations, - batch_size=self.args.generation_batch_size // self.num_generations, - repeat_count=self.num_iterations * self.args.steps_per_generation, - shuffle=self.shuffle_dataset, - seed=self.args.seed, - ) - - def _get_eval_sampler(self, eval_dataset) -> Sampler: - # See _get_train_sampler for an explanation of the sampler. - return RepeatSampler( - data_source=eval_dataset, - mini_repeat_count=self.num_generations_eval, - seed=self.args.seed, - ) - - # This method overrides `Trainer.get_eval_dataloader` to wrap iterable eval datasets, reproducing the - # RepeatSampler ordering that can't be attached to them (see `get_train_dataloader`). Map-style datasets keep the - # default path via `_get_eval_sampler`, which shuffles with `seed`, so the iterable wrap shuffles too (buffered) - # to walk prompts in a matching order. - # Maintenance note: this method is a copy-paste of the original `Trainer.get_eval_dataloader`, with the iterable - # wrapping as the only addition. - def get_eval_dataloader(self, eval_dataset: str | Dataset | IterableDataset | None = None) -> DataLoader: - if eval_dataset is None and self.eval_dataset is None: - raise ValueError("Trainer: evaluation requires an eval_dataset.") - - # If we have persistent workers, don't do a fork bomb especially as eval datasets - # don't change during training - dataloader_key = eval_dataset if isinstance(eval_dataset, str) else "eval" - if ( - hasattr(self, "_eval_dataloaders") - and dataloader_key in self._eval_dataloaders - and self.args.dataloader_persistent_workers - ): - return self._eval_dataloaders[dataloader_key] - - eval_dataset = ( - self.eval_dataset[eval_dataset] - if isinstance(eval_dataset, str) - else eval_dataset - if eval_dataset is not None - else self.eval_dataset - ) - - if isinstance(eval_dataset, IterableDataset): - # Apply the `__init__` iterable config here too - if self.args.accelerator_config.dispatch_batches: - raise ValueError( - "Iterable datasets require `dispatch_batches=False`, but it is set to `True` in " - "`accelerator_config`. Please set it to `False`." - ) - self.accelerator.dataloader_config.dispatch_batches = False - eval_dataset = eval_dataset.shuffle(seed=self.args.seed) - eval_dataset = repeat_iterable_dataset(eval_dataset, mini_repeat_count=self.num_generations_eval) - # Force a single worker for this loader only, without persisting the change - num_workers = self.args.dataloader_num_workers - self.args.dataloader_num_workers = 0 - - try: - return self._get_dataloader( - dataset=eval_dataset, - description="Evaluation", - batch_size=self.args.eval_batch_size, - sampler_fn=self._get_eval_sampler, - dataloader_key=dataloader_key, - ) - finally: - if isinstance(eval_dataset, IterableDataset): - self.args.dataloader_num_workers = num_workers - - def get_high_entropy_mask(self, entropies: torch.Tensor, mask: torch.Tensor, threshold: float) -> torch.Tensor: - """ - Returns a binary mask identifying tokens whose entropy exceeds a given quantile threshold. - - Args: - entropies (`torch.Tensor`): - Tensor of shape (batch_size, seq_len) with per-token entropy values. - mask (`torch.Tensor`): - Binary mask of the same shape as `entropies`, where `1` indicates valid tokens and `0` padding. - threshold (`float`): - Quantile threshold between `0.0` and `1.0` to select high-entropy tokens. - - Returns: - `torch.Tensor`: - Boolean mask of shape (batch_size, seq_len), where `True` indicates tokens with entropy >= threshold - and `False` otherwise. - """ - local = entropies[mask.bool()].float() - - # Use a negative pad_value as a sentinel because entropy values are always >= 0. - # This guarantees that the sentinel cannot collide with any real entropy value. - pad_value = -1e9 - - # Pad across processes so that every rank has the same tensor length - padded = self.accelerator.pad_across_processes(local, dim=0, pad_index=pad_value) - gathered = self.accelerator.gather(padded) - - # Drop sentinel values (safe because no entropy can be negative) - gathered = gathered[gathered != pad_value] - - if gathered.numel() == 0: - return torch.zeros_like(entropies, dtype=torch.bool) - - entropy_threshold = torch.quantile(gathered, threshold) - masked_entropies = entropies * mask.float() - entropy_mask = masked_entropies >= entropy_threshold - return entropy_mask & mask.bool() # ensure padding tokens are always masked out - - 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 in this function is licensed under AGPL3. - if compute_efficient: - # No scoring forward may run (vLLM, beta = 0): latch precision here, not from the shared env. - _unsloth_grpo_autocast_unscored(self) - # The current policy is then scored only in the zoo loss: refuse the kept set there too. - _unsloth_grpo_refuse_unscored_support(self) - return None, None - else: - _unsloth_grpo_autocast(self) - - compute_aux_loss = kwargs.get("compute_aux_loss", None) - - # Body-local: this source is copied out without this module's imports. #6960. - _grpo_vision_chunks = None - try: - from unsloth_zoo.rl_replacements import grpo_vision_chunks as _grpo_vision_chunks - except Exception: - pass - # Collected even without the zoo: an older one must cost only image slicing. - vision_inputs = _unsloth_grpo_vision_inputs(kwargs) - if _grpo_vision_chunks is None and vision_inputs.get("pixel_values", None) is not None: - raise RuntimeError( - "Unsloth: vision GRPO needs an unsloth_zoo build that exports " - "grpo_vision_chunks, the shared multimodal key tuple and chunker " - "used by both GRPO logprob paths. Please upgrade unsloth_zoo to " - "2026.9.5 or newer: pip install -U unsloth_zoo" - ) - pixel_values = vision_inputs.get("pixel_values", None) - image_grid_thw = vision_inputs.get("image_grid_thw", None) - video_grid_thw = vision_inputs.get("video_grid_thw", None) - if pixel_values is None: - # Only a sentinel below (the chunks carry the tensors): video rows must take the - # vision path too, since left and sequence packing move M-RoPE video positions. - pixel_values = vision_inputs.get("pixel_values_videos", None) - num_images = vision_inputs.get("num_images", None) - # Transformers 5.x needs token_type_ids/mm_token_type_ids for some vision models. - token_type_ids = vision_inputs.get("token_type_ids", None) - mm_token_type_ids = vision_inputs.get("mm_token_type_ids", None) - if ( - mm_token_type_ids is not None - or image_grid_thw is not None - or video_grid_thw is not None - ): - mm_token_type_ids = _unsloth_fix_mm_token_type_ids( - self.processing_class, input_ids, mm_token_type_ids - ) - vision_inputs["mm_token_type_ids"] = mm_token_type_ids - - unwrapped_model = self.accelerator.unwrap_model(model, keep_fp32_wrapper = False) - - from unsloth_zoo.rl_replacements import ( - chunked_hidden_states_selective_log_softmax, - chunked_selective_log_softmax, - ) - - # Old / ref logps must use the policy pass's scorer, else the ratio starts away from 1 and the KL is biased. - try: - if os.environ.get("UNSLOTH_GRPO_LOGPROB_RECOMPUTE", "0") == "1": - from unsloth_zoo.rl_replacements import grpo_truncated_support_kwargs - if not grpo_truncated_support_kwargs(self): - from unsloth_zoo.rl_replacements import ( - chunked_hidden_states_selective_log_softmax_recompute as chunked_hidden_states_selective_log_softmax, - ) - else: - from unsloth_zoo.rl_replacements import ( - grpo_chunked_hidden_states_logprobs as chunked_hidden_states_selective_log_softmax, - ) - except ImportError: - pass - - output_embeddings = unwrapped_model.get_output_embeddings() - lm_head = output_embeddings.weight - lm_head_bias = getattr(output_embeddings, "bias", None) - if lm_head_bias is not None and not UNSLOTH_ZOO_HAS_HEAD_BIAS: - import warnings - warnings.warn( - "Unsloth: this model's output head has a bias, which GRPO scoring needs " - "a newer unsloth_zoo to include. Please upgrade: pip install -U unsloth_zoo", - stacklevel = 2, - ) - lm_head_bias = None - _head_log_softmax = chunked_hidden_states_selective_log_softmax - # ZeRO-3 released the head after the forward: score with one gathered read-only copy. - _zero3_head = None - try: - from unsloth_zoo.fused_losses.cross_entropy_loss import ( - unsloth_zero3_gathered, - unsloth_zero3_partitioned, - ) - except (ImportError, NotImplementedError): - # Older zoo, or zoo refusing to import without an accelerator (CPU gloo workers). - pass - else: - if unsloth_zero3_partitioned(lm_head) or unsloth_zero3_partitioned( - getattr(output_embeddings, "bias", None) - ): - _zero3_head = ( - unsloth_zero3_gathered(lm_head), - unsloth_zero3_gathered(getattr(output_embeddings, "bias", None)), - ) - lm_head = _zero3_head[0] - if lm_head_bias is not None: - lm_head_bias = _zero3_head[1] - - def chunked_hidden_states_selective_log_softmax( - hidden_states, _lm_head, *args, **kwargs - ): - # FSDP2's forward hooks replace Parameters. Read the current head - # at projection time, including the bias, instead of retaining a shard. - if _zero3_head is not None: - head_weight, head_bias = _zero3_head - else: - head_weight = output_embeddings.weight - head_bias = getattr(output_embeddings, "bias", None) - if UNSLOTH_ZOO_HAS_HEAD_BIAS: - kwargs["lm_head_bias"] = head_bias - return _head_log_softmax( - hidden_states, - head_weight, - *args, - **kwargs, - ) - - # UNSLOTH_GRPO_TRUNCATED_SUPPORT (=1 or zoo's automatic gate): score over the rollout's - # top_k / top_p / min_p kept set, matching vLLM's processed_logprobs. The zoo loss refuses to mix - # normalisations unless this marker says the old / ref logps were scored the same way. - try: - from unsloth_zoo.rl_replacements import grpo_truncated_support_kwargs - _support_kwargs = grpo_truncated_support_kwargs(self) - except Exception: - _support_kwargs = {} - _support_mismatch = _unsloth_grpo_support_mismatch(self) if _support_kwargs else [] - if _support_mismatch: - # Left unmarked, the zoo loss drops the kept set for the current policy as well. - if not getattr(self, "_unsloth_truncated_support_refused", False): - import warnings - warnings.warn( - "Unsloth: not scoring over the rollouts' top_k / top_p / min_p kept set " - "(UNSLOTH_GRPO_TRUNCATED_SUPPORT=1, or automatic for truncated vLLM " - "rollouts feeding the IS correction or off-policy mask) because " - f"{', '.join(_support_mismatch)}: those logprobs come from weights other " - "than the rollout's, whose top_k / top_p / min_p kept set differs from the " - "sampler's, so per-policy kept sets give unnormalised PPO ratios and KL. " - "Scoring over the full vocabulary instead.", - stacklevel = 2, - ) - # Also silences the zoo's refusal, which would blame the unsloth version. - self._unsloth_truncated_support_refused = True - _support_kwargs = {} - if _support_kwargs: - from functools import partial - - chunked_hidden_states_selective_log_softmax = partial( - chunked_hidden_states_selective_log_softmax, - **_support_kwargs, - ) - chunked_selective_log_softmax = partial( - chunked_selective_log_softmax, - **_support_kwargs, - ) - self._unsloth_truncated_support_scored = True - - # Size on the dtype the forward actually runs in: with autocast off that is the model's own dtype. - forward_dtype = ( - self._autocast_dtype if getattr(self, "_autocast_enabled", True) else lm_head.dtype - ) - dtype_bytes = 16 if forward_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 - - import math - - total_samples = input_ids.shape[0] - batch_size = math.ceil(total_samples / B) - - input_ids_chunks = [] - attention_mask_chunks = [] - 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]) - - # One chunker shared with the gradient pass, so the two cannot disagree. - if _grpo_vision_chunks is None: - # Image-indexed keys already raised above, so only per-sample ones are left. - vision_chunks = [] - for _start in range(0, total_samples, batch_size): - _end = min(_start + batch_size, total_samples) - vision_chunks.append( - { - _key: vision_inputs[_key][_start:_end] - for _key in ("token_type_ids", "mm_token_type_ids") - if vision_inputs.get(_key, None) is not None - } - ) - else: - vision_chunks = _grpo_vision_chunks(vision_inputs, total_samples, batch_size) - - temperature = self.temperature - model_config = _unsloth_get_model_config(model) - if detect_logit_transforms is not None: - # model_config, not model: under DDP/Accelerate `model` is a wrapper that does not forward .config, so the helper would report zeros. - _transforms = detect_logit_transforms(model_config) - logit_softcapping = _transforms["logit_softcapping"] - logit_scale_multiply = _transforms["logit_scale_multiply"] - logit_scale_divide = _transforms["logit_scale_divide"] - else: - logit_softcapping = _unsloth_get_final_logit_softcapping(model) - logit_scale_multiply, logit_scale_divide = _unsloth_resolve_logit_scales( - model_config - ) - - zipped_inputs = zip( - input_ids_chunks, - attention_mask_chunks, - vision_chunks, - ) - os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" - - # Sequence packing (default on; UNSLOTH_GRPO_SEQ_PACKING=0 disables): one varlen [1, sum L] forward replaces the padded [B, Lmax] loop and fixes the left-pad RoPE error. Self-verified against the per-row forward, re-checked as T grows, and falls back if a backend ignores packed_seq_lengths. - logprobs = None - - # PrefixGrouper (GRPO shared-prompt dedup, default ON): G completions share the prompt, so storing it once behind a FlexAttention shared-prefix mask cuts the trunk forward from G*(P+R) to P+G*R tokens. Gated by UNSLOTH_GRPO_PREFIX_GROUPER, a tok_r auto-gate and a first-use self-verify, so a mask/isolation regression cannot ship silently. - _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: the rollout dominates the step, so PG saves little and its 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 shared-prefix isolation, so a decoder that does not forward prefix_seg_info leaks 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 = DEVICE_TYPE_TORCH, - dtype = self._autocast_dtype, - enabled = getattr(self, "_autocast_enabled", True), - ): - _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, 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 every step just re-pays it. Disable PG persistently; the packed/padded path below still gives 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; UNSLOTH_GRPO_SEQ_PACKING=0 disables): one varlen block-diagonal forward replaces the padded loop exactly and fixes its left-pad RoPE error. Self-verified, re-checked as T grows, falls back if a backend ignores packed_seq_lengths, and lm_head runs on completion positions only. - _pk_result = None - _pk_use = False - _pk_T = None - # True only while the packed forward and its scatter run, so an OOM elsewhere (setup, per-row verifier) is not read as a length limit. - _pk_in_packed = 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_T = sum(_pk_nz_cpu) - _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) - # Skip the packed forward entirely at known-unsafe lengths, avoiding a wasted pass or 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 - # Decided from the host lengths alone, so a capped batch skips the flatten and the active-row sync below. - _pk_fits = ( - _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) - ) - if _pk_fits: - _pk_flat = input_ids[_pk_keep].unsqueeze(0) - # 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()) - if _pk_fits and (_pk_ok is True or _pk_active >= 2): - _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, matching 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 - _pk_in_packed = True - with _get_inference_mode_context_manager(model): - with torch.amp.autocast( - device_type = DEVICE_TYPE_TORCH, - dtype = self._autocast_dtype, - enabled = getattr(self, "_autocast_enabled", True), - ): - # 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_out = _pk_hidden[0, :-1, :][_pk_ctgt].unsqueeze(0) - _pk_ids = _pk_flat[0, 1:][_pk_ctgt].unsqueeze(0) - # Hidden states or logits? Logits mean the forward already applied scaling/softcapping. - if _unsloth_grpo_returns_hidden_states( - unwrapped_model, _pk_out, lm_head - ): - _pk_sel = chunked_hidden_states_selective_log_softmax( - _pk_out, - lm_head, - _pk_ids, - _pk_chunks, - logit_scale_multiply, - logit_scale_divide, - logit_softcapping, - temperature, - )[0] - else: - # Model returned logits directly: scaling/softcapping already applied by the model forward. - _pk_sel = chunked_selective_log_softmax( - _pk_out, - _pk_ids, - temperature, - _pk_chunks, - )[0] - # GPT-OSS offload race guard, matching the padded loop. - device_synchronize() - # Scatter each logprob back to its (row, col) so [:, -_pk_W:] matches the padded path. - _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:] - ) - _pk_in_packed = False - # Re-verify when T or the longest segment grows past the verified envelope; 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 = DEVICE_TYPE_TORCH, - dtype = self._autocast_dtype, - enabled = getattr(self, "_autocast_enabled", True), - ): - 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_rout = _pk_rh[:, :-1, :] - # Hidden states or logits? Logits mean the forward already applied scaling/softcapping. - if _unsloth_grpo_returns_hidden_states( - unwrapped_model, _pk_rout, lm_head - ): - _pk_rsel = chunked_hidden_states_selective_log_softmax( - _pk_rout, - lm_head, - _pk_real[:, 1:], - 1, - logit_scale_multiply, - logit_scale_divide, - logit_softcapping, - temperature, - )[0] - else: - # Model returned logits directly: scaling/softcapping already applied by the model forward. - _pk_rsel = chunked_selective_log_softmax( - _pk_rout, - _pk_real[:, 1:], - temperature, - 1, - )[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 is ~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 at least 2 completion rows exercised cross-sample packing; a single row proves 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 and use the padded loop for this batch. - _pk_hidden = None - _pk_sel = None - _pk_result = None - _pk_use = False - if isinstance(_pk_err, torch.cuda.OutOfMemoryError): - torch.cuda.empty_cache() - if ( - isinstance(_pk_err, torch.cuda.OutOfMemoryError) - and _pk_in_packed - and _pk_T is not None - ): - # An OOM in the packed forward says this flattened length does not fit, not that packing is wrong: keep packing smaller batches. - _pk_unsafe = getattr( - unwrapped_model, "_unsloth_seq_packing_nograd_unsafe_T", None - ) - _pk_unsafe = _pk_T if _pk_unsafe is None else min(_pk_unsafe, _pk_T) - unwrapped_model._unsloth_seq_packing_nograd_unsafe_T = _pk_unsafe - if UNSLOTH_ENABLE_LOGGING: - print( - f"[Unsloth] GRPO sequence-packing (no-grad) OOM at T={_pk_T}: packing capped below T={_pk_unsafe}, this batch uses the padded loop: {_pk_err!r}", - flush = True, - ) - else: - 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 packed result over the completion mask. Below tol_ok trust the structure, at or above TOL_KILL mark it unsafe forever, borderline falls back for 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, - ) - # No packed reference (packing off or failed) means this cannot be verified, so 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: - _pk_hidden = _pk_sel = _pk_result = _pk_ref = None - # Index buffers of a declined packed batch would otherwise live through the padded loop. - _pk_keep = _pk_len = _pk_flat = _pk_pos = _pk_cmask = None - _pk_nz_idx = _pk_within = _pk_cstart = _pk_ctgt = _pk_tgt = _pk_ids = _pk_cm = None - - with _get_inference_mode_context_manager(model): - for ( - input_ids_chunk, - attention_mask_chunk, - vision_chunk, - ) in zipped_inputs: - with torch.amp.autocast( - device_type = DEVICE_TYPE_TORCH, - dtype = self._autocast_dtype, - enabled = getattr(self, "_autocast_enabled", True), - ): - if pixel_values is None: - outputs = unwrapped_model( - input_ids = input_ids_chunk, - attention_mask = attention_mask_chunk, - **vision_chunk, - ) - - logits_chunk = outputs.logits - del outputs - - 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, :] - # Hidden states or logits? Logits mean the forward already applied scaling/softcapping. - if _unsloth_grpo_returns_hidden_states( - unwrapped_model, logits_chunk, lm_head - ): - 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 the model forward. - logprobs_chunk = chunked_selective_log_softmax( - logits_chunk, - completion_input_ids_chunk, - temperature, - input_ids_chunk.shape[0] * multiplier, - ) - else: - # VLMs do not take the optimized path in models/, so they never hit the Flash Attn left-padding issue. - outputs = unwrapped_model( - input_ids = input_ids_chunk, - attention_mask = attention_mask_chunk, - logits_to_keep = logits_to_keep + 1, - **vision_chunk, - ) - - logits_chunk = outputs.logits - del outputs - - logits_chunk = logits_chunk[:, :-1, :] - completion_input_ids_chunk = input_ids_chunk[:, -logits_to_keep:] - # Hidden states or logits? Logits mean the forward already applied scaling/softcapping. - if _unsloth_grpo_returns_hidden_states( - unwrapped_model, logits_chunk, lm_head - ): - 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: - logprobs_chunk = chunked_selective_log_softmax( - logits_chunk, - completion_input_ids_chunk, - temperature, - ) - # Avoids a race with GPT OSS offload_embbed=True; it does not appear to slow models down. - 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: off by default (router_aux_loss_coef = 0 in models/rl.py) and explicit opt-in is rejected at trainer init, so it is always None. Kept for TRL >= 1.7.0's 3-tuple. - aux_loss = None - return logprobs.detach(), entropies, aux_loss # logps, entropies, aux_loss - # transformers <= 4.48 does not support logits_to_keep, so drop the logits here; see huggingface/trl#2770. - - def training_step(self, model, inputs, num_items_in_batch): - time_before = time.perf_counter() - output = super().training_step(model, inputs, num_items_in_batch) - self._step += 1 - time_after = time.perf_counter() - self._current_train_step_time += time_after - time_before - if self._step % self.current_gradient_accumulation_steps == 0: - self._metrics["train"]["step_time"].append(self._current_train_step_time) - self._current_train_step_time = 0.0 - return output - - @profiling_decorator - def _prepare_inputs(self, generation_batch: dict[str, torch.Tensor | Any]) -> dict[str, torch.Tensor | Any]: - # Prepares inputs for model training/evaluation by managing completion generation and batch handling. - # During training: - # - Receives the local generation batch (Per-GPU batch size × steps per generation) - # from the modified training dataloader instead of the standard local batch - # - Generates completions once for the entire generation batch and splits it into batches of size - # `per_device_train_batch_size` - # - Buffers these completions and returns the appropriate slice for the current accumulation step - # - Optimizes by regenerating completions only periodically (every steps_per_generation * num_iterations) - # During evaluation: - # - The input is treated as a standard local batch (no accumulation, no multiple iterations) - # - Completions are generated for each batch without buffering or reuse - # Returns a single local batch in both cases. - - mode = "train" if self.model.training else "eval" - if mode == "train": - generate_every = self.args.steps_per_generation * self.num_iterations - if self._step % generate_every == 0 or self._buffered_inputs is None: - # self._buffered_inputs=None can occur when resuming from a checkpoint - generation_batch = self._generate_and_score_completions(generation_batch) - generation_batch = split_pixel_values_by_grid(generation_batch) - generation_batch = _unsloth_grpo_split_vision_by_sample(generation_batch) - - try: generation_batch = shuffle_sequence_dict(generation_batch) - - except: pass - generation_batches = split_tensor_dict(generation_batch, self.args.steps_per_generation) - self._buffered_inputs = [_unsloth_grpo_unsplit_vision(unsplit_pixel_values_by_grid(batch)) for batch in generation_batches] - inputs = self._buffered_inputs[self._step % self.args.steps_per_generation] - else: - # In evaluation, there is neither batch grouping for generation, nor multiple iterations, hence - # local generation batch == local eval batch - inputs = self._generate_and_score_completions(generation_batch) - return inputs - - def _log_completion_extra(self, column: str, values: list): - """ - Log extra columns to the completions table. Called from reward functions via the `log_extra` kwarg. - - Args: - column (`str`): - Name of the column to add. - values (`list`): - Values for the column, one per sample in the batch. - """ - self._pending_extra_logs[column].extend(values) - - def _log_metric(self, name: str, value: float): - """ - Log a scalar metric from a reward function. Called via the `log_metric` kwarg. Values are averaged over each - logging step and reported alongside built-in metrics like `kl` and `entropy`. - - Args: - name (`str`): - Name of the metric. - value (`float`): - Scalar value for this batch. - """ - self._pending_metrics[name].append(value) - - @profiling_decorator - def _calculate_rewards(self, inputs, prompts, completions, completion_ids_list): - device = self.accelerator.device - rewards_per_func = torch.zeros(len(prompts), len(self.reward_funcs), device=device) - - # Repeat all input columns (but "prompt", "completion", and "completion_ids") to match the num of generations - keys = [key for key in inputs[0] if key not in ["prompt", "completion", "completion_ids"]] - reward_kwargs = {key: [example[key] for example in inputs] for key in keys} - - # This allows for dynamic reward shaping based on training progress. - reward_kwargs["trainer_state"] = self.state - - # Allow reward functions to log extra columns to the completions table. - reward_kwargs["log_extra"] = self._log_completion_extra - - # Allow reward functions to log additional scalar metrics. - reward_kwargs["log_metric"] = self._log_metric - - # Expose the per-completion environment instances to reward functions (both sync and async paths). - if self.environments is not None: - reward_kwargs["environments"] = self.environments - - async_funcs_info = [] # async custom functions for asyncio.gather - - for i, (reward_func, reward_processing_class, reward_func_name) in enumerate( - zip(self.reward_funcs, self.reward_processing_classes, self.reward_func_names, strict=True) - ): - if isinstance(reward_func, nn.Module): # Module (no PretrainedModel) for compat with compiled models - with profiling_context(self, reward_func_name): - if is_conversational({"prompt": prompts[0]}): - messages = [{"messages": p + c} for p, c in zip(prompts, completions, strict=True)] - texts = [ - apply_chat_template(x, reward_processing_class, **self.chat_template_kwargs)["text"] - for x in messages - ] - else: - texts = [p + c for p, c in zip(prompts, completions, strict=True)] - reward_inputs = reward_processing_class( - text=texts, return_tensors="pt", padding=True, padding_side="right", add_special_tokens=False - ) - reward_inputs = super()._prepare_inputs(reward_inputs) - with torch.inference_mode(): - rewards_per_func[:, i] = reward_func(**reward_inputs).logits[:, 0] # Shape (B*G,) - elif is_async_callable(reward_func): # Separate async reward funcs to run them in parallel later - async_funcs_info.append((i, reward_func, reward_func_name)) - else: - # Run synchronous reward function - with profiling_context(self, reward_func_name): - output_reward_func = reward_func( - prompts=prompts, completions=completions, completion_ids=completion_ids_list, **reward_kwargs - ) - # Convert None values to NaN - output_reward_func = [reward if reward is not None else torch.nan for reward in output_reward_func] - if len(output_reward_func) != len(prompts): - raise ValueError( - f"The reward function '{reward_func_name}' returned {len(output_reward_func)} rewards, but " - f"{len(prompts)} were expected (one per prompt-completion pair). Make sure the reward " - f"function returns exactly one reward per completion." - ) - rewards_per_func[:, i] = torch.tensor(output_reward_func, dtype=torch.float32, device=device) - - # Execute async custom functions in parallel using asyncio.gather - if async_funcs_info: - - async def _invoke_async(index, func, func_name): - with profiling_context(self, func_name): - output = await func( - prompts=prompts, completions=completions, completion_ids=completion_ids_list, **reward_kwargs - ) - output = [r if r is not None else torch.nan for r in output] - if len(output) != len(prompts): - raise ValueError( - f"The reward function '{func_name}' returned {len(output)} rewards, but {len(prompts)} " - f"were expected (one per prompt-completion pair). Make sure the reward function returns " - f"exactly one reward per completion." - ) - return index, output - - async def _run_async_funcs(): - coros = [_invoke_async(i, func, func_name) for (i, func, func_name) in async_funcs_info] - return await asyncio.gather(*coros) - - async_results = asyncio.run_coroutine_threadsafe(_run_async_funcs(), self.async_loop).result() - for idx, output_reward_func in async_results: - rewards_per_func[:, idx] = torch.tensor(output_reward_func, dtype=torch.float32, device=device) - - # If all reward functions return None for a given row, issue a detailed warning - if torch.isnan(rewards_per_func).all(dim=1).any(): - nan_row_idx = torch.isnan(rewards_per_func).all(dim=1).nonzero(as_tuple=True)[0][0] - row_reward_kwargs = { - key: value[nan_row_idx] - for key, value in reward_kwargs.items() - if key not in ("trainer_state", "log_extra", "log_metric") - } - row_reward_kwargs["prompt"] = prompts[nan_row_idx] - row_reward_kwargs["completion"] = completions[nan_row_idx] - logger.warning( - f"All reward functions returned None for the following kwargs:\n{row_reward_kwargs}\n" - "Please ensure that at least one reward function returns a valid reward." - ) - - # Gather the reward per function: this part is crucial, because the rewards are normalized per group and the - # completions may be distributed across processes - rewards_per_func = gather(rewards_per_func) - return rewards_per_func - - def _tokenize_prompts(self, prompts: list): - """Tokenize prompts and extract images/multimodal fields for generation.""" - if is_conversational({"prompt": prompts[0]}): - # Normalize string content to content blocks for VLM processors that don't handle plain strings. - if self._is_vlm: - prompts = [prepare_multimodal_messages(prompt) for prompt in prompts] - - # Extract images from messages for VLM support - images = [] - has_images = False - for prompt in prompts: - prompt_images = [] - for message in prompt: - if isinstance(message["content"], list): - for part in message["content"]: - if part["type"] == "image": - prompt_images.append(part["image"]) - has_images = True - images.append(prompt_images if prompt_images else None) - images = images if has_images else None - - if self.environment_factories is not None: - # Tools differ per example across environments, so render each prompt with its own tool schema. - prompt_ids = [] - multimodal_fields = {} - for prompt, name in zip(prompts, self._batch_environments, strict=True): - tokenized = self.processing_class.apply_chat_template( - conversation=[prompt], - tools=self._env_tools[name] or None, # `or None`: Llama bug: tool boilerplate for tools=[] - chat_template=self.chat_template, - add_generation_prompt=True, - tokenize=True, - return_dict=True, - **self.chat_template_kwargs, - ) - prompt_ids.append(tokenized["input_ids"][0]) - # For VLMs, the processor returns extra multimodal fields (pixel_values, image_grid_thw, etc.) - for k, v in tokenized.items(): - if k not in ("input_ids", "attention_mask"): - multimodal_fields.setdefault(k, []).append(v) - # Merge the per-prompt fields so the result matches the single batched `apply_chat_template` call used - # when there are no environments: image fields (pixel_values, image_grid_thw, ...) are flattened over - # patches/images and concatenated; per-token fields (e.g. token_type_ids) stay batched per prompt. - multimodal_fields = { - k: torch.cat(v) if isinstance(v[0], torch.Tensor) else [row for prompt_v in v for row in prompt_v] - for k, v in multimodal_fields.items() - } - return prompt_ids, images, multimodal_fields - - # Workaround for a bug in transformers 5.3.0 where some processors (e.g. Qwen2.5-VL) crash on - # batched unpadded input (transformers#44514). - # Fixed in transformers 5.4.0 (transformers#44563). - needs_padding_workaround = Version("5.3.0") <= Version(transformers.__version__) < Version("5.4.0") - tokenized = self.processing_class.apply_chat_template( - conversation=prompts, - tools=self.tools or None, # `or None`: Llama bug: it renders tool boilerplate for tools=[] - chat_template=self.chat_template, - add_generation_prompt=True, - tokenize=True, - return_dict=True, - **({"padding": True} if needs_padding_workaround else {}), - **self.chat_template_kwargs, - ) - if needs_padding_workaround: - # Unpad input_ids: remove padding tokens using attention_mask to get per-sequence lists - prompt_ids = [ - [tok for tok, m in zip(ids, mask, strict=True) if m] - for ids, mask in zip(tokenized["input_ids"], tokenized["attention_mask"], strict=True) - ] - else: - prompt_ids = tokenized["input_ids"] - # For VLMs, the processor returns extra multimodal fields (pixel_values, image_grid_thw, etc.) - multimodal_fields = {k: v for k, v in tokenized.items() if k not in ("input_ids", "attention_mask")} - else: - prompt_ids = self.processing_class(text=prompts)["input_ids"] - images = None - multimodal_fields = {} - return prompt_ids, images, multimodal_fields - - def _generate_single_turn(self, prompt_ids, images, multimodal_fields, num_generations, has_tool_images=False): - device = self.accelerator.device - - # Generate completions using either vLLM or regular generation - if self.use_vllm: - # Sync weights if training step changed - if self.state.global_step != self._last_loaded_step: - if not getattr(getattr(self.vllm_generation, 'llm', None), 'shared_weights', False): - with profiling_context(self, 'sync_weights'): - self.vllm_generation.sync_weights() - self._last_loaded_step = self.state.global_step - - # Generate using vLLM with raw token IDs - _, completion_ids, logprobs, _ = self.vllm_generation.generate( - prompts=prompt_ids, - images=images, - num_generations=num_generations, - profiler=profiling_context(self, "vLLM.generate"), - ) - # vLLM returns per-token top-k logprobs; keep only the top-1 (sampled token) logprob - logprobs = [[lp[0] for lp in seq] for seq in logprobs] - - elif self.use_transformers_continuous_batching: - with ( - profiling_context(self, "transformers.generate_batch"), - unwrap_model_for_generation( - self.model_wrapped, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation - ) as unwrapped_model, - torch.no_grad(), - self._dist.summon_full_params(self.model_wrapped, recurse=False), - ): - # Cast to the appropriate dtype based on training configuration - if self.args.bf16: - unwrapped_model.to(torch.bfloat16) - elif self.args.fp16: - unwrapped_model.to(torch.float16) - if self.args.cast_lm_head_to_fp32: - unwrapped_model.lm_head.to(torch.float32) - all_outputs = unwrapped_model.generate_batch( - prompt_ids, - generation_config=self.generation_config, - continuous_batching_config=self.continuous_batching_config, - progress_bar=False, - ) - unwrapped_model.train() - completion_ids = [output.generated_tokens for output in all_outputs.values()] - logprobs = None - - else: - # Regular generation path: left-pad token IDs into tensors - prompt_tensors = [torch.tensor(ids) for ids in prompt_ids] - padded_ids = pad(prompt_tensors, padding_value=self._tokenizer.pad_token_id, padding_side="left") - attention_mask = pad([torch.ones_like(t) for t in prompt_tensors], padding_value=0, padding_side="left") - generate_inputs = {"input_ids": padded_ids, "attention_mask": attention_mask} - # For VLMs, include multimodal fields as tensors (pixel_values, image_grid_thw, etc.) - for k, v in multimodal_fields.items(): - if isinstance(v, torch.Tensor): - generate_inputs[k] = v - elif isinstance(v, list) and v and isinstance(v[0], list): - # Per-token field (e.g., token_type_ids): left-pad like input_ids - generate_inputs[k] = pad([torch.tensor(x) for x in v], padding_value=0, padding_side="left") - else: - generate_inputs[k] = torch.tensor(np.array(v)) - - # For VLM tool images: build token type IDs from the padded input IDs. - if self._is_vlm and self.tools and has_tool_images: - mm_ids = torch.zeros_like(padded_ids) - if self._image_pad_token_id is not None: - mm_ids[padded_ids == self._image_pad_token_id] = 1 - if self._video_pad_token_id is not None: - mm_ids[padded_ids == self._video_pad_token_id] = 2 - - # Use the same key the model expects: token_type_ids for models like Gemma, - # mm_token_type_ids for models like Qwen. - if "image_grid_thw" in generate_inputs: - generate_inputs["mm_token_type_ids"] = mm_ids - else: - generate_inputs["token_type_ids"] = mm_ids - - generate_inputs = super()._prepare_inputs(generate_inputs) - 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 - - with ( - profiling_context(self, "transformers.generate"), - unwrap_model_for_generation( - self.model_wrapped, - self.accelerator, - gather_deepspeed3_params=self.args.ds3_gather_for_generation, - generation_kwargs=self.generation_kwargs, # Override model.generation_config with generation_kwargs to fix transformers#42762 - ) as unwrapped_model, - torch.no_grad(), - self._dist.summon_full_params(self.model_wrapped, recurse=False), - ): - prompt_completion_ids = unwrapped_model.generate( - **generate_inputs, generation_config=self.generation_config - ) - # Compute prompt length and extract completion ids - prompt_length = generate_inputs["input_ids"].size(1) - completion_ids = prompt_completion_ids[:, prompt_length:] - - # Mask everything after the first EOS token - is_eos = torch.isin(completion_ids, torch.tensor(self.eos_token_ids, device=device)) - eos_idx = torch.full((is_eos.size(0),), is_eos.size(1), dtype=torch.long, device=device) - eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)] - sequence_indices = torch.arange(is_eos.size(1), device=device).expand(is_eos.size(0), -1) - completion_mask = (sequence_indices <= eos_idx.unsqueeze(1)).int() - completion_ids = [ - c[m].tolist() for c, m in zip(completion_ids.cpu(), completion_mask.bool().cpu(), strict=True) - ] - logprobs = None # not used in this case - - return completion_ids, logprobs - - def _get_tool_suffix_ids(self, tool_messages): - """Get token IDs for tool result formatting by using a minimal dummy conversation.""" - # Use the real tool name instead of a dummy: some templates (e.g. GPT-OSS) derive the tool response - # header from the assistant's tool call name. - dummy_tool_calls = [{"type": "function", "function": {"name": tool_messages[0]["name"], "arguments": {}}}] - dummy_messages = [ - {"role": "user", "content": "dummy"}, - { - "role": "assistant", - # "content" is required here because VLM processors crash on tokenize=True without it - # (KeyError in processing_utils.py). See huggingface/transformers#45290. - "content": "", - "tool_calls": dummy_tool_calls, - }, - ] - if self._is_vlm: - dummy_messages = prepare_multimodal_messages(dummy_messages) - tool_messages = prepare_multimodal_messages(tool_messages) - - prefix_ids = self.processing_class.apply_chat_template( - dummy_messages, - add_generation_prompt=False, - tokenize=True, - chat_template=self.chat_template, - return_dict=False, - **self.chat_template_kwargs, - ) - full_ids = self.processing_class.apply_chat_template( - dummy_messages + tool_messages, - add_generation_prompt=True, - tokenize=True, - chat_template=self.chat_template, - return_dict=False, - **self.chat_template_kwargs, - ) - # VLM processors return batched output (list of lists), unbatch for single conversation - if self._is_vlm: - prefix_ids = prefix_ids[0] - full_ids = full_ids[0] - - # Some chat templates (notably Qwen3/Qwen3.5) render "...<|im_end|>\n" after an assistant/tool block. - # When we compute `suffix_ids` by slicing `full_ids`, we must align the slicing boundary to - # EOS (not EOS + newline). Templates that don't use EOS as end-of-turn (e.g. Gemma uses - # ) skip this trimming. - eos_positions = [i for i, tok_id in enumerate(prefix_ids) if tok_id == self._tokenizer.eos_token_id] - if eos_positions: - prefix_ids = prefix_ids[: eos_positions[-1] + 1] - - if full_ids[: len(prefix_ids)] != prefix_ids: - raise ValueError("Unexpected tokenization: the EOS-trimmed prefix IDs are not a prefix of the full IDs.") - return full_ids[len(prefix_ids) :] - - def _tool_call_loop(self, prompts, prompt_ids, completion_ids, completions, logprobs, images, multimodal_fields): - # Tool execution loop: execute tools, then regenerate completions with tool results appended to the prompt - tool_calls = [completion[0].get("tool_calls") for completion in completions] - idxs_with_tool = [idx for idx, tool_call in enumerate(tool_calls) if tool_call] - tool_calls = [tool_calls[idx] for idx in idxs_with_tool] - tool_mask = [[1] * len(ids) for ids in completion_ids] # 0 for tool result tokens, 1 elsewhere - # Collect images from multimodal tool responses for the forward pass - tool_images = [[] for _ in completion_ids] - tool_call_count = 0 - tool_failure_count = 0 - iteration_num = 0 - - while idxs_with_tool and iteration_num < self.max_tool_calling_iterations: - prompt_completion_tools = [prompts[i] for i in idxs_with_tool] # select only prompts that need tool calls - # Snapshot state so we can rollback tool results that would exceed max_completion_length - completions_len_before = [len(completions[i]) for i in idxs_with_tool] - tool_images_len_before = [len(tool_images[i]) for i in idxs_with_tool] - prompts_len_before = [len(prompts[i]) for i in idxs_with_tool] - - # Call the tools, and build the new prompt for generation - for idx in range(len(idxs_with_tool)): - idx_with_tool = idxs_with_tool[idx] - tool_call_list = tool_calls[idx] - prompt_completion_tool = prompt_completion_tools[idx] - sync_tool_dict = self._sync_tool_dicts[idx_with_tool] - async_tool_dict = self._async_tool_dicts[idx_with_tool] - # Append the last assistant message (which triggered tool_calls) to the prompt - prompt_completion_tool.append(completions[idx_with_tool][-1]) - async_coros = [] - tool_call_results = [] - for tool_call in tool_call_list: - tool_call_count += 1 - if tool_call["type"] == "function": - function = tool_call["function"] - name = function["name"] - try: - if name in sync_tool_dict: - tool_call_results.append((name, sync_tool_dict[name](**function["arguments"]))) - elif name in async_tool_dict: - async_coros.append((name, async_tool_dict[name](**function["arguments"]))) - else: - raise ValueError(f"Tool {name} not found.") - except Exception as e: - tool_failure_count += 1 - result = {"error": str(e)} - tool_call_results.append((name, result)) - else: - tool_failure_count += 1 - name = tool_call.get("name", "unknown") - tool_call_results.append((name, {"error": f"Unsupported tool call type: {tool_call['type']}"})) - - if async_coros: - - async def _run_async_tools(async_coros): - coros = [coro for _, coro in async_coros] - results = await asyncio.gather(*coros, return_exceptions=True) - return [(name, result) for (name, _), result in zip(async_coros, results, strict=False)] - - async_results = asyncio.run_coroutine_threadsafe( - _run_async_tools(async_coros), self.async_loop - ).result() - - for name, result in async_results: - if isinstance(result, Exception): - tool_failure_count += 1 - tool_call_results.append((name, {"error": str(result)})) - else: - tool_call_results.append((name, result)) - - for name, result in tool_call_results: - # Support multimodal tool responses: if the tool returns a list of content blocks - # (e.g., [{"type": "image", "image": ...}, {"type": "text", "text": "..."}]), - # pass them through directly so _tokenize_prompts can extract images for VLMs. - content = result if isinstance(result, list) else str(result) - tool_message = {"role": "tool", "name": name, "content": content} - # Collect images from multimodal tool responses - if isinstance(content, list): - for part in content: - if isinstance(part, dict) and part.get("type") == "image": - tool_images[idx_with_tool].append(part["image"]) - prompt_completion_tool.append(tool_message) - completions[idx_with_tool].append(tool_message) - - # Build token IDs by concatenation: prompt + completion + tool_suffix. - prompt_completion_tool_ids = [] - for idx in range(len(idxs_with_tool)): - idx_with_tool = idxs_with_tool[idx] - # Extract trailing tool messages from completions - tool_messages = [] - for message in reversed(completions[idx_with_tool]): - if message["role"] == "tool": - tool_messages.insert(0, message) - else: - break - suffix_ids = self._get_tool_suffix_ids(tool_messages) - prompt_completion_tool_ids.append( - prompt_ids[idx_with_tool] + completion_ids[idx_with_tool] + suffix_ids - ) - - # Drop tool results whose addition would push the sequence past max_completion_length (the completion - # budget) or past the backend context ceiling (vLLM and transformers will error out on inputs longer than - # the model's max length). The sample exits the loop with its completion as-is, and the tool - # messages/images appended this iteration are rolled back so completions and tool_images stay consistent - # with completion_ids. - if self.use_vllm and self.vllm_mode == "colocate": - max_model_len = self.vllm_generation.llm.llm_engine.model_config.max_model_len - else: - config = self.model.config.text_config if self._is_vlm else self.model.config - max_model_len = config.max_position_embeddings - overlong = [ - len(pct) - len(prompt_ids[i]) > self.max_completion_length or len(pct) >= max_model_len - for i, pct in zip(idxs_with_tool, prompt_completion_tool_ids, strict=True) - ] - for idx in range(len(idxs_with_tool)): - if overlong[idx]: - idx_with_tool = idxs_with_tool[idx] - del completions[idx_with_tool][completions_len_before[idx] :] - del tool_images[idx_with_tool][tool_images_len_before[idx] :] - del prompts[idx_with_tool][prompts_len_before[idx] :] - # Keep only non-overlong items for further processing - idxs_with_tool = [idx for idx, o in zip(idxs_with_tool, overlong, strict=True) if not o] - prompt_completion_tool_ids = [ - pct for pct, o in zip(prompt_completion_tool_ids, overlong, strict=True) if not o - ] - if not idxs_with_tool: - break # all overlong, exit tool loop - - # Filter images and multimodal fields to match the current subset (index into full batch). - # Merge tool response images so the model can see visual feedback during generation. - merged_images = images - if any(imgs for imgs in tool_images): - if merged_images is None: - merged_images = [imgs if imgs else None for imgs in tool_images] - else: - merged_images = [ - (existing or []) + new for existing, new in zip(merged_images, tool_images, strict=True) - ] - loop_images = [merged_images[i] for i in idxs_with_tool] if merged_images else None - if self._is_vlm and self.tools and any(imgs for imgs in tool_images): - flat_loop_images = [img for img_list in loop_images if img_list for img in img_list] - if flat_loop_images: - image_inputs = self.processing_class.image_processor(images=flat_loop_images, return_tensors="pt") - image_inputs = super()._prepare_inputs(image_inputs) - loop_multimodal_fields = dict(image_inputs) - else: - loop_multimodal_fields = {} - elif multimodal_fields: - if "num_images" not in multimodal_fields and images is not None: - multimodal_fields = { - **multimodal_fields, - "num_images": [len(img_list) if img_list else 0 for img_list in images], - } - split_fields = split_pixel_values_by_grid(multimodal_fields) - loop_multimodal_fields = {} - for k, v in split_fields.items(): - selected = [v[i] for i in idxs_with_tool] - # Per-token type-id fields may arrive as tensors (e.g. via the padding workaround); - # convert them to lists so they follow the same left-padding path as list-valued - # per-token fields below, keeping them aligned with the left-padded input_ids. - if k in ("token_type_ids", "mm_token_type_ids") and isinstance(selected[0], torch.Tensor): - selected = [s.tolist() for s in selected] - # Per-token fields (e.g. token_type_ids) need zero-padding to match extended prompt length - if isinstance(selected[0], list): - selected = [ - s + [0] * (len(pct) - len(s)) - for s, pct in zip(selected, prompt_completion_tool_ids, strict=True) - ] - elif isinstance(v, torch.Tensor): - selected = torch.stack(selected) - loop_multimodal_fields[k] = selected - loop_multimodal_fields = unsplit_pixel_values_by_grid(loop_multimodal_fields) - loop_multimodal_fields.pop("num_images", None) - if "pixel_values" in loop_multimodal_fields and loop_multimodal_fields["pixel_values"].numel() == 0: - for key in ["pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw"]: - loop_multimodal_fields.pop(key, None) - else: - loop_multimodal_fields = {} - - # Generate new completions after tool execution (using concatenated IDs, no re-tokenization) - post_tool_ids, post_tool_logprobs = self._generate_single_turn( - prompt_completion_tool_ids, - loop_images, - loop_multimodal_fields, - num_generations=1, # each sample has its own history, so generate one completion per sample - has_tool_images=any(imgs for imgs in tool_images), - ) - - # Truncate so that pct[len(prompt_ids[idx]) :] + post_tool does not exceed max_completion_length. - # The pre-regen check guarantees len(completion_tool_ids) <= max_completion_length, so any - # excess can only come from post_tool_ids. post_tool_ids is model-generated text and never - # contains image tokens, so a plain slice is safe. - for idx in range(len(idxs_with_tool)): - idx_with_tool = idxs_with_tool[idx] - completion_tool_length = len(prompt_completion_tool_ids[idx]) - len(prompt_ids[idx_with_tool]) - excess_length = completion_tool_length + len(post_tool_ids[idx]) - self.max_completion_length - if excess_length > 0: - new_len = len(post_tool_ids[idx]) - excess_length - post_tool_ids[idx] = post_tool_ids[idx][:new_len] - if logprobs is not None: - post_tool_logprobs[idx] = post_tool_logprobs[idx][:new_len] - - # Update tool_mask: the tool result should be 0 and the post-tool 1 - for idx in range(len(idxs_with_tool)): - idx_with_tool = idxs_with_tool[idx] - prompt_completion_tool_length = len(prompt_completion_tool_ids[idx]) - prompt_length = len(prompt_ids[idx_with_tool]) - completion_length = len(completion_ids[idx_with_tool]) - post_tool_length = len(post_tool_ids[idx]) - tool_length = prompt_completion_tool_length - prompt_length - completion_length - tool_mask[idx_with_tool] += [0] * tool_length + [1] * post_tool_length - if logprobs is not None: - logprobs[idx_with_tool] += [0.0] * tool_length + post_tool_logprobs[idx] - - # Update completion_ids with the new completions (after tool execution) - for idx in range(len(idxs_with_tool)): - idx_with_tool = idxs_with_tool[idx] - prompt_length = len(prompt_ids[idx_with_tool]) - pct = prompt_completion_tool_ids[idx] # = prompt-completion-tool - completion_ids[idx_with_tool] = pct[prompt_length:] + post_tool_ids[idx] - - # Decode post-tool completions - post_tool_completions = [ - parse_response(self._tokenizer, ids, prefix=prompt_completion_tool_ids[idx]) if ids else {} - for idx, ids in enumerate(post_tool_ids) - ] - - # Add post-tool completions to the existing completions - for idx in range(len(idxs_with_tool)): - idx_with_tool = idxs_with_tool[idx] - if post_tool_completions[idx]: # {} if post-tool completions completely truncated - completions[idx_with_tool].append(post_tool_completions[idx]) - - # Check for further tool calls - tool_calls = [completion.get("tool_calls") for completion in post_tool_completions] - idxs_with_tool = [idx for idx, tool_call in zip(idxs_with_tool, tool_calls, strict=True) if tool_call] - tool_calls = [tool_call for tool_call in tool_calls if tool_call] - iteration_num += 1 - - return tool_mask, completions, completion_ids, logprobs, tool_call_count, tool_failure_count, tool_images - - def _generate(self, prompts: list): - device = self.accelerator.device - mode = "train" if self.model.training else "eval" - - # Copy the prompts to avoid modifying the original list - prompts = copy.deepcopy(prompts) - - if self.rollout_func is not None: - # Keep vLLM weights in sync for custom rollouts that rely on vLLM utilities. - if self.use_vllm and self.state.global_step != self._last_loaded_step: - if not getattr(getattr(self.vllm_generation, 'llm', None), 'shared_weights', False): - with profiling_context(self, 'sync_weights'): - self.vllm_generation.sync_weights() - self._last_loaded_step = self.state.global_step - - # Pass prompts to rollout_func preserving structured messages. - # Chat templating must happen inside rollout_func, at the backend boundary, so that - # multimodal content (images, typed content blocks) is not lost before rollout logic runs. - output = self.rollout_func(prompts, self) - required_keys = {"prompt_ids", "completion_ids", "logprobs"} - missing_keys = required_keys - output.keys() - if missing_keys: - missing_keys_list = sorted(missing_keys) - raise ValueError(f"rollout_func must return keys {missing_keys_list} in its output dict.") - extra_fields = {k: v for k, v in output.items() if k not in required_keys} - prompt_ids, completion_ids, logprobs = output["prompt_ids"], output["completion_ids"], output["logprobs"] - images = None - multimodal_fields = {} - else: - prompt_ids, images, multimodal_fields = self._tokenize_prompts(prompts) - num_generations = self.num_generations if mode == "train" else self.num_generations_eval - completion_ids, logprobs = self._generate_single_turn( - prompt_ids, images, multimodal_fields, num_generations - ) - extra_fields = {} - - # Decode completions. It's important to use `parse_response` when possible, because it handles tool calls. - if is_conversational({"prompt": prompts[0]}): - if Version(transformers.__version__) >= Version("5.0.0") and ( # parse_response added in v5 - getattr(self._tokenizer, "response_template", None) is not None # new-style - or getattr(self._tokenizer, "response_schema", None) is not None # old-style - ): - completions = [ - [parse_response(self._tokenizer, ids, prefix=prompt_ids[i])] - for i, ids in enumerate(completion_ids) - ] - else: - contents = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) - completions = [[{"role": "assistant", "content": content}] for content in contents] - else: - completions = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) - - # Extract tool calls from the completions and (possibly) execute them - tool_images = [] - if self.tools: - ( - tool_mask, - completions, - completion_ids, - logprobs, - tool_call_count, - tool_failure_count, - tool_images, - ) = self._tool_call_loop( - prompts, prompt_ids, completion_ids, completions, logprobs, images, multimodal_fields - ) - # Merge tool response images into the images list for the forward pass - if any(imgs for imgs in tool_images): - if images is None: - images = [imgs if imgs else None for imgs in tool_images] - else: - images = [(existing or []) + new for existing, new in zip(images, tool_images, strict=True)] - else: - # Support custom env_mask from rollout_func (e.g., for environment feedback masking) - # Internally treated as tool_mask - marks model tokens (1) vs external tokens (0) - tool_mask = extra_fields.pop("env_mask", None) - - # Get completion length per sequence, used for logging - prompt_lengths = torch.tensor([len(ids) for ids in prompt_ids], device=device) - if tool_mask is not None: # count only model-generated tokens (tool_mask=1) - completion_lengths = torch.tensor([sum(mask) for mask in tool_mask], device=device) - else: - completion_lengths = torch.tensor([len(ids) for ids in completion_ids], device=device) - agg_prompt_lengths = self.accelerator.gather(prompt_lengths) - agg_completion_lengths = self.accelerator.gather(completion_lengths) - # Fail clearly if the generation backend returned no completions (avoids a cryptic min() error below). - if agg_completion_lengths.numel() == 0: - raise RuntimeError( - "No completions were generated. This usually means the generation backend failed to return any " - "results; see the generation logs above for the underlying error." - ) - total_prompt_tokens = agg_prompt_lengths.sum() - - # Log the metrics - if mode == "train": - self.state.num_input_tokens_seen += (total_prompt_tokens + agg_completion_lengths.sum()).item() - self._metrics[mode]["num_tokens"] = [self.state.num_input_tokens_seen] - - # Log completion lengths, mean, min, max - self._metrics[mode]["completions/mean_length"].append(agg_completion_lengths.float().mean().item()) - self._metrics[mode]["completions/min_length"].append(agg_completion_lengths.float().min().item()) - self._metrics[mode]["completions/max_length"].append(agg_completion_lengths.float().max().item()) - - # Identify sequences that terminated with EOS and log their lengths - eos_and_pad = [*self.eos_token_ids, self._tokenizer.pad_token_id] - is_truncated = torch.tensor([ids[-1] not in eos_and_pad for ids in completion_ids], device=device) - agg_is_truncated = self.accelerator.gather(is_truncated) - self._metrics[mode]["completions/clipped_ratio"].append(agg_is_truncated.float().mean().item()) - term_completion_lengths = agg_completion_lengths[~agg_is_truncated] - if len(term_completion_lengths) == 0: # edge case where no terminated sequences are found - term_completion_lengths = torch.zeros(1, device=device) - self._metrics[mode]["completions/mean_terminated_length"].append(term_completion_lengths.float().mean().item()) - self._metrics[mode]["completions/min_terminated_length"].append(term_completion_lengths.float().min().item()) - self._metrics[mode]["completions/max_terminated_length"].append(term_completion_lengths.float().max().item()) - - if self.tools: - agg_tool_call_count = self.accelerator.gather(torch.tensor(tool_call_count, device=device)).sum() - tool_call_frequency = (agg_tool_call_count / len(agg_prompt_lengths)).item() - self._metrics[mode]["tools/call_frequency"].append(tool_call_frequency) - agg_tool_failure_count = self.accelerator.gather(torch.tensor(tool_failure_count, device=device)).sum() - failure_frequency = ( - (agg_tool_failure_count / agg_tool_call_count).item() if agg_tool_call_count > 0 else 0.0 - ) - self._metrics[mode]["tools/failure_frequency"].append(failure_frequency) - - return prompt_ids, completion_ids, tool_mask, completions, logprobs, extra_fields, images, tool_images - - def _generate_and_score_completions( - self, inputs: list[dict[str, torch.Tensor | Any]] - ) -> dict[str, torch.Tensor | Any]: - device = self.accelerator.device - mode = "train" if self.model.training else "eval" - - # `prompt` is optional only when an environment owns the data (e.g. a multi-environment routing dataset that - # carries only an `environment` column); each rollout's `reset()` then supplies it. Default it here rather than - # writing it back onto the row, so the placeholder stays out of the `reset()` kwargs built below, while every - # prompt-derived check downstream (conversational detection, multimodal handling) stays consistent. Without an - # environment, a missing `prompt` is a malformed dataset and must still fail fast. - if self.environment_factories is not None: - prompts = [x.get("prompt", [{"role": "user", "content": ""}]) for x in inputs] - else: - prompts = [x["prompt"] for x in inputs] - # Unsloth: Extract per-sample chat_template_kwargs before metadata is lost - _ct_ = getattr(self.processing_class, 'chat_template', None) or '' - _sk_ = {'prompt', 'chosen', 'rejected', 'completion', 'messages', 'label', - 'images', 'image', 'videos', 'video', 'audios', 'audio'} - self._unsloth_batch_chat_kwargs = [] - for _inp_ in inputs: - _kw_ = {} - if isinstance(_inp_, dict): - for _k_ in _inp_.keys() - _sk_: - if _k_ in _ct_ and isinstance(_inp_[_k_], str): - _kw_[_k_] = _inp_[_k_] - self._unsloth_batch_chat_kwargs.append(_kw_) - prompts = _unsloth_grpo_clean_video_prompts(prompts, self) - # Resolve each example's environment and draw one reusable instance per rollout from the pool, creating more - # only when this batch needs more concurrent instances of an environment than exist. `_batch_environments` - # records each example's environment so `_tokenize_prompts` can render the matching tool schema. - if self.environment_factories is not None: - self._batch_environments = [x.get("environment") if self._multi_environment else None for x in inputs] - if self._multi_environment: - for name in set(self._batch_environments): - if name not in self.environment_factories: - raise ValueError( - f"Example has `environment={name!r}`, which is not among the environments passed to " - f"`environment_factory`. Expected one of: {list(self.environment_factories)}." - ) - self.environments = [] - pool_cursor = {} # how many instances of each environment have been handed out so far this batch - for name in self._batch_environments: - pool = self._environment_pool[name] - index = pool_cursor.get(name, 0) - if index == len(pool): - pool.append(self.environment_factories[name]()) - pool_cursor[name] = index + 1 - self.environments.append(pool[index]) - - # Build the per-rollout tool dicts for this batch: the standalone tools plus, for each rollout, the methods of - # its environment. Done here (not at init) because each example's environment, hence its tools, is data-dependent. - if self.tools: - self._sync_tool_dicts = [] - self._async_tool_dicts = [] - for i in range(len(inputs)): - methods = [] - if self.environments: - environment = self.environments[i] - methods = [ - getattr(environment, member_name) - for member_name, _ in inspect.getmembers(type(environment), predicate=inspect.isfunction) - if member_name not in ("reset", "get_reward") and not member_name.startswith("_") - ] - sync_tool_dict, async_tool_dict = {}, {} - for tool in self._standalone_tools + methods: - if inspect.iscoroutinefunction(tool): - async_tool_dict[tool.__name__] = tool - else: - sync_tool_dict[tool.__name__] = tool - self._sync_tool_dicts.append(sync_tool_dict) - self._async_tool_dicts.append(async_tool_dict) - - if self.environments: - for i, (prompt, environment, x) in enumerate(zip(prompts, self.environments, inputs, strict=True)): - # `environment` is a control field in multi-environment mode, so it is not forwarded to `reset`. - reset_kwargs = {k: v for k, v in x.items() if k != "environment"} if self._multi_environment else x - observation = environment.reset(**reset_kwargs) - if observation is None: - continue - content = prompt[-1]["content"] - if isinstance(observation, list) and isinstance(content, str): - content = [{"type": "text", "text": content}] - if isinstance(observation, str) and isinstance(content, list): - observation = [{"type": "text", "text": observation}] - # Rebuild the last message rather than mutating it in place, so the input example is left untouched. - prompts[i] = prompt[:-1] + [{**prompt[-1], "content": content + observation}] - - if "images" in inputs[0]: - images = [example.get("images") for example in inputs] - elif "image" in inputs[0]: - images = [_unsloth_grpo_image_cell(example.get("image")) for example in inputs] - else: - images = None - # Transformers requires at least one image in the batch, otherwise it throws an error - if images is not None and all(img_list == [] for img_list in images): - images = None - - # If the prompts are conversational and the inputs contain images, we need to convert the prompts from - # [{"role": "user", "content": "What color is the sky?"}] to - # [{"role": "user", "content": [{"type": "image", "image": }, {"type": "text", "text": "What color is the sky?"}]}] - if images is not None: - if not is_conversational(inputs[0]): - raise ValueError( - "Multimodal training requires conversational prompts. It looks like the dataset contains " - "non-conversational inputs, likely because a chat template was applied before passing the dataset " - "to the trainer. Please provide the raw conversational prompts and let the trainer apply the chat " - "template internally." - ) - prompts = [ - prepare_multimodal_messages(prompt, images=image_list) - for prompt, image_list in zip(prompts, images, strict=True) - ] - - dataset_images = images # preserve dataset images before _generate may overwrite - ( - prompt_ids_list, - completion_ids_list, - tool_mask_list, - completions, - sampling_per_token_logps_list, - extra_fields, - images, - tool_images, - ) = self._generate(prompts) - - _unsloth_clear_stateful_mrope( - self.accelerator.unwrap_model(self.model, keep_fp32_wrapper = False) - ) - if images is None: - images = dataset_images # restore dataset images (rollout_func path returns None) - - # Convert lists of token IDs to padded tensors - prompt_ids = [torch.tensor(ids) for ids in prompt_ids_list] - prompt_mask = [torch.ones_like(ids, dtype=torch.long) for ids in prompt_ids] - prompt_ids = pad( - prompt_ids, - padding_value=self._tokenizer.pad_token_id, - padding_side="left", - pad_to_multiple_of=self.pad_to_multiple_of, - ).to(device=device) - prompt_mask = pad( - prompt_mask, padding_value=0, padding_side="left", pad_to_multiple_of=self.pad_to_multiple_of - ).to(device=device) - completion_ids = [torch.tensor(ids) for ids in completion_ids_list] - completion_mask = [torch.ones_like(ids, dtype=torch.long) for ids in completion_ids] - completion_ids = pad( - completion_ids, - padding_value=self._tokenizer.pad_token_id, - padding_side="right", - pad_to_multiple_of=self.pad_to_multiple_of, - ).to(device=device) - completion_mask = pad( - completion_mask, padding_value=0, padding_side="right", pad_to_multiple_of=self.pad_to_multiple_of - ).to(device=device) - if sampling_per_token_logps_list is not None: - # vLLM replaces a NaN token logprob with `None` (see `extract_logprobs`); map it back to NaN so the - # tensor builds (`torch.tensor([..., None, ...])` raises "Could not infer dtype of NoneType"). The NaN - # positions are neutralized in the importance-sampling ratio below so they contribute no correction. - sampling_per_token_logps = [ - torch.tensor([float("nan") if x is None else x for x in logps]) - for logps in sampling_per_token_logps_list - ] - sampling_per_token_logps = pad( - sampling_per_token_logps, - padding_value=0.0, - padding_side="right", - pad_to_multiple_of=self.pad_to_multiple_of, - ).to(device=device) - else: - sampling_per_token_logps = None - if tool_mask_list is not None: - tool_mask = [torch.tensor(mask) for mask in tool_mask_list] - tool_mask = pad( - tool_mask, padding_value=1, padding_side="right", pad_to_multiple_of=self.pad_to_multiple_of - ).to(device=device) - else: - tool_mask = None - - # If mask_truncated_completions is enabled, zero out truncated completions for attention and loss masking - if self.mask_truncated_completions: - eos_and_pad = [*self.eos_token_ids, self._tokenizer.pad_token_id] - is_truncated = torch.tensor([ids[-1] not in eos_and_pad for ids in completion_ids_list], device=device) - # Mask completion_mask for attention masking - completion_mask = completion_mask * (~is_truncated).unsqueeze(1).int() - # Also mask tool_mask for consistency in multi-turn training - if tool_mask is not None: - tool_mask = tool_mask * (~is_truncated).unsqueeze(1).int() - - loss_mask = completion_mask if tool_mask is None else completion_mask * tool_mask - num_items_in_batch = self.accelerator.gather(loss_mask.sum()).sum() - - # Concatenate prompt_mask with completion_mask for logit computation - prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1) # (B, P+C) - attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) # (B, P+C) - - logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens - - max_left_pad = None - batch_size = self.args.per_device_train_batch_size if mode == "train" else self.args.per_device_eval_batch_size - # Which name carries "this batch has images" moved twice, and at the bottom of the - # declared window neither name exists: 0.20.0 through 0.23.1 bind has_images, 0.24.0 - # and up bind images, and 0.18.2 / 0.19.1 have no vision path at all, so every batch - # there is text only. Probing has_images and falling back to images without a third - # branch made the floor raise NameError out of the except handler and killed training. - try: - _unsloth_text_only = not has_images - except NameError: - try: - _unsloth_text_only = images is None - except NameError: - _unsloth_text_only = True - try: - _unsloth_video_images = images - except NameError: - _unsloth_video_images = None - try: - _unsloth_video_prompts = prompts - except NameError: - _unsloth_video_prompts = [] - _unsloth_video_kwargs = _unsloth_grpo_video_inputs( - self, _unsloth_video_prompts, _unsloth_video_images - ) - if _unsloth_video_kwargs is not None: - _unsloth_text_only = False - if _unsloth_text_only: - # 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) - - num_images = [len(img_list) if img_list else 0 for img_list in images] if images is not None else None - - # Get forward_kwargs for models with multimodal inputs. - # When tool images are present (from _tool_call_loop), use image_processor directly and build - # mm_token_type_ids from prompt_completion_ids. Otherwise, use the full processor pipeline - # which returns model-specific keys (image_sizes, pixel_attention_mask, etc.). - if self.tools and any(imgs for imgs in tool_images) and self._is_vlm: - flat_images = [img for img_list in images if img_list for img in img_list] - image_inputs = self.processing_class.image_processor(images=flat_images, return_tensors="pt") - image_inputs = super()._prepare_inputs(image_inputs) - forward_kwargs = dict(image_inputs) - elif images is not None: - if self.environment_factories is not None: - per_prompt_tools = [self._env_tools[name] for name in self._batch_environments] - else: - per_prompt_tools = [self.tools] * len(prompts) - prompts_text = [ - apply_chat_template( - {"prompt": prompt}, self.processing_class, tools=tools, **self.chat_template_kwargs - )["prompt"] - for prompt, tools in zip(prompts, per_prompt_tools, strict=True) - ] - prompt_inputs = self.processing_class(images=images, text=prompts_text, padding=True, return_tensors="pt") - prompt_inputs = super()._prepare_inputs(prompt_inputs) - forward_kwargs = {k: v for k, v in prompt_inputs.items() if k not in ["input_ids", "attention_mask"]} - else: - forward_kwargs = {} - if _unsloth_video_kwargs is not None: - forward_kwargs = {**forward_kwargs, **_unsloth_video_kwargs} - - # Recover LFM2-VL tile counts; the full processor drops row/column metadata. - num_tiles = None - if images is not None and "spatial_shapes" in forward_kwargs: - image_info = self.processing_class.image_processor( - images=images, return_tensors="pt", return_row_col_info=True - ) - tiles_per_image = image_info["image_rows"] * image_info["image_cols"] - if self.processing_class.image_processor.use_thumbnail: - tiles_per_image = tiles_per_image + (tiles_per_image > 1).to(tiles_per_image.dtype) - num_tiles = [group.sum().item() for group in torch.split(tiles_per_image, num_images)] - # Same for InternVL, whose pixel_values is tile-indexed ([total_tiles, channels, height, width]). - elif ( - images is not None - and forward_kwargs["pixel_values"].ndim == 4 - and forward_kwargs["pixel_values"].size(0) != sum(num_images) - ): - num_patches = self.processing_class.image_processor( - images=images, crop_to_patches=True, return_tensors="pt" - )["num_patches"] - num_tiles = [group.sum().item() for group in torch.split(num_patches, num_images)] - - # If token_type_ids are used, extend them with zeros for the completion part - if "token_type_ids" in forward_kwargs: - token_type_ids = forward_kwargs["token_type_ids"] - if self.pad_to_multiple_of is not None: - # Needed only with pad_to_multiple_of: otherwise prompt_ids and token_type_ids must have equal len - padding_size = prompt_ids.size(1) - token_type_ids.size(1) - if padding_size > 0: - token_type_ids = torch.cat( - [token_type_ids.new_zeros((token_type_ids.size(0), padding_size)), token_type_ids], dim=1 - ) - forward_kwargs["token_type_ids"] = torch.cat( - [token_type_ids, token_type_ids.new_zeros(completion_ids.shape)], dim=1 - ) - # If mm_token_type_ids are used, extend them with zeros for the completion part - if "mm_token_type_ids" in forward_kwargs: - mm_token_type_ids = forward_kwargs["mm_token_type_ids"] - if self.pad_to_multiple_of is not None: - # Needed only with pad_to_multiple_of: otherwise prompt_ids and mm_token_type_ids must have equal len - padding_size = prompt_ids.size(1) - mm_token_type_ids.size(1) - if padding_size > 0: - mm_token_type_ids = torch.cat( - [mm_token_type_ids.new_zeros((mm_token_type_ids.size(0), padding_size)), mm_token_type_ids], - dim=1, - ) - forward_kwargs["mm_token_type_ids"] = torch.cat( - [mm_token_type_ids, mm_token_type_ids.new_zeros(completion_ids.shape)], dim=1 - ) - if ( - "mm_token_type_ids" in forward_kwargs - or "image_grid_thw" in forward_kwargs - or "video_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 - - # For VLM tool images: build token type IDs from the full prompt_completion_ids. - # This must happen AFTER the token_type_ids/mm_token_type_ids extension blocks above, - # because our version already covers the full sequence (images are in the completion, - # not just the prompt). - if self.tools and any(imgs for imgs in tool_images) and self._is_vlm: - mm_ids = torch.zeros_like(prompt_completion_ids) - if self._image_pad_token_id is not None: - mm_ids[prompt_completion_ids == self._image_pad_token_id] = 1 - if self._video_pad_token_id is not None: - mm_ids[prompt_completion_ids == self._video_pad_token_id] = 2 - - # Use the same key the model expects: token_type_ids for models like Gemma, - # mm_token_type_ids for models like Qwen. - image_grid_thw = forward_kwargs.get("image_grid_thw") - if image_grid_thw is not None: - forward_kwargs["mm_token_type_ids"] = mm_ids - else: - forward_kwargs["token_type_ids"] = mm_ids - - # Truncation safety (Qwen-style models with image_grid_thw only): if - # max_completion_length truncated some image tokens, the number of image pad tokens - # in input_ids won't match pixel_values features. Check per-sample and drop ALL - # images for any sample with a mismatch (safe fallback). - if image_grid_thw is not None and num_images is not None: - merge_length = getattr(self.processing_class.image_processor, "merge_size", 2) ** 2 - img_offset = 0 - has_mismatch = False - for b in range(mm_ids.shape[0]): - sample_tokens = (mm_ids[b] == 1).sum().item() - sample_features = 0 - for i in range(num_images[b]): - grid_idx = img_offset + i - if grid_idx < image_grid_thw.shape[0]: - sample_features += image_grid_thw[grid_idx].prod().item() // merge_length - if sample_tokens != sample_features: - has_mismatch = True - break - img_offset += num_images[b] - - if has_mismatch: - # Drop all images: safer than partial trim which is error-prone - forward_kwargs.pop("pixel_values", None) - forward_kwargs.pop("image_grid_thw", None) - mm_ids.zero_() - forward_kwargs["mm_token_type_ids"] = mm_ids - num_images = None - - # When gradient checkpointing is enabled with use_reentrant=True (non default), calling the model inside a - # torch.no_grad() block triggers a harmless PyTorch warning ("None of the inputs have requires_grad=True"). - # Temporarily disable checkpointing to avoid this warning during inference. - with torch.no_grad(), disable_gradient_checkpointing(self.model, self.args.gradient_checkpointing_kwargs): - # If the generation and optimization steps are misaligned—i.e., if generation does not occur at the end of - # a full optimizer step (when gradient_accumulation_steps is not a multiple of generate_every)—then the - # samples may come from an earlier version of the model. In that case, we need to track old_per_token_logps - # for importance sampling. If the steps are aligned, importance sampling isn't necessary and we set - # old_per_token_logps to None. - # When using vLLM, we always compute old_per_token_logps for importance sampling, it was shown that the - # distribution mismatch between vLLM and the training model can be large and harm the training. - generate_every = self.args.steps_per_generation * self.num_iterations # generation frequency - _unsloth_rollout_logps = locals().get('sampling_per_token_logps') - _unsloth_rollout_old = _unsloth_grpo_rollout_logps_as_old( - self, _unsloth_rollout_logps, completion_mask, locals().get('tool_mask'), - ) - if _unsloth_rollout_old: - old_per_token_logps = None - elif self.args.gradient_accumulation_steps % generate_every != 0 or ( - self.use_vllm and self.vllm_importance_sampling_correction - ): - old_per_token_logps, _, _ = self._get_per_token_logps_and_entropies( - self.model, - prompt_completion_ids, - attention_mask, - logits_to_keep, - batch_size=batch_size, - num_images=num_images, - num_tiles=num_tiles, - **forward_kwargs, # may contain pixel_values, image_grid_thw, pixel_attention_mask, spatial_shapes, image_sizes, image_position_ids - ) - else: - old_per_token_logps = None - - # Compute the importance sampling ratio when using vLLM, to correct for potential distribution mismatch - if False and self.use_vllm and self.vllm_importance_sampling_correction: - mask = completion_mask if tool_mask is None else completion_mask * tool_mask - per_token_logps_diff = (old_per_token_logps - sampling_per_token_logps) * mask - # Tokens whose sampling logprob was NaN (unavailable from vLLM) get a zero difference, so their - # importance ratio is exactly 1 (no correction) rather than propagating NaN through the loss. - per_token_logps_diff = torch.nan_to_num(per_token_logps_diff, nan=0.0) - - sequence_level_is = self.vllm_importance_sampling_mode in ["sequence_mask", "sequence_truncate"] - if sequence_level_is: - per_sequence_logps_diff = per_token_logps_diff.sum(dim=-1, keepdim=True) - logps_diff = per_sequence_logps_diff - else: - logps_diff = per_token_logps_diff - - vllm_importance_sampling_ratio = torch.exp(logps_diff) - - # vllm_importance_sampling_ratio.shape: - # token_* modes: (B, T) (per-token ratio) - # sequence_* modes: (B, 1) (per-sequence ratio) - - if self.vllm_importance_sampling_mode in ["sequence_truncate", "token_truncate"]: - vllm_importance_sampling_ratio = torch.clamp( - vllm_importance_sampling_ratio, - min=self.vllm_importance_sampling_clip_min, - max=self.vllm_importance_sampling_clip_max, - ) - elif self.vllm_importance_sampling_mode in ["sequence_mask", "token_mask"]: - min_val = ( - self.vllm_importance_sampling_clip_min - if self.vllm_importance_sampling_clip_min is not None - else -math.inf - ) - max_val = ( - self.vllm_importance_sampling_clip_max - if self.vllm_importance_sampling_clip_max is not None - else math.inf - ) - - invalid_mis_mask = (vllm_importance_sampling_ratio < min_val) | ( - vllm_importance_sampling_ratio > max_val - ) - vllm_importance_sampling_ratio = vllm_importance_sampling_ratio.masked_fill( - invalid_mis_mask, value=0.0 - ) - else: - raise ValueError( - f"Unknown vLLM importance sampling level: {self.vllm_importance_sampling_mode}. Possible values are 'token_truncate', 'token_mask', 'sequence_truncate', and 'sequence_mask'." - ) - - # Compute the per-token log probabilities for the reference model - if self.beta != 0.0: - if self.ref_model is not None: - ref_per_token_logps, _, _ = self._get_per_token_logps_and_entropies( - self.ref_model, - prompt_completion_ids, - attention_mask, - logits_to_keep, - batch_size=batch_size, - num_images=num_images, - num_tiles=num_tiles, - **forward_kwargs, # may contain pixel_values, image_grid_thw, pixel_attention_mask, spatial_shapes, image_sizes, image_position_ids - ) - else: - # When training a PEFT adapter, how we obtain the reference depends on the setup: - # - New adapter: disabling adapters yields the base model. - # - Re-training an existing adapter: an initial copy is loaded under the name "ref". - model = self.accelerator.unwrap_model(self.model) - with use_adapter(model, adapter_name="ref" if "ref" in model.peft_config else None): - ref_per_token_logps, _, _ = self._get_per_token_logps_and_entropies( - self.model, - prompt_completion_ids, - attention_mask, - logits_to_keep, - batch_size=batch_size, - num_images=num_images, - num_tiles=num_tiles, - **forward_kwargs, # may contain pixel_values, image_grid_thw, pixel_attention_mask, spatial_shapes, image_sizes, image_position_ids - ) - else: - ref_per_token_logps = None - - # Merge extra_fields from rollout_func into inputs for reward functions - if extra_fields: - for i, inp in enumerate(inputs): - for key, values in extra_fields.items(): - if isinstance(values, list) and i < len(values): - inp[key] = values[i] - elif not isinstance(values, list): - inp[key] = values - - # Calculate rewards for each reward function. rewards_per_func aggregates rewards across all processes. This is - # important because rewards will be normalized per group, and completions are distributed. We will later slice - # rewards_per_func to extract each process's subset. - _unsloth_reward_locals = locals() - if images is not None and 'prompts_text' in _unsloth_reward_locals: - completions_text = _unsloth_reward_locals.get('completions_text') - if completions_text is None: - completions_text = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) - rewards_per_func = self._calculate_rewards(inputs, prompts_text, completions_text, completion_ids_list) - else: - rewards_per_func = self._calculate_rewards(inputs, prompts, completions, completion_ids_list) - del _unsloth_reward_locals - num_generations = self.num_generations if mode == "train" else self.num_generations_eval - - # A completion for which every reward function returned None is unscorable. nansum would collapse it to 0, - # which both biases the per-group baseline and hands the completion a spurious advantage. Mark these rows NaN - # so they're excluded from the (nan-aware) baseline below; their advantage is forced to 0 afterwards. - unscorable_mask = torch.isnan(rewards_per_func).all(dim=1) - - if self.multi_objective_aggregation == "sum_then_normalize": - # Apply weights to each reward function's output and sum - rewards = (rewards_per_func * self.reward_weights.to(device).unsqueeze(0)).nansum(dim=1) - rewards[unscorable_mask] = torch.nan - mean_grouped_rewards = torch.nanmean(rewards.view(-1, num_generations), dim=1) - mean_grouped_rewards = mean_grouped_rewards.repeat_interleave(num_generations, dim=0) - if self.scale_rewards in ["group", "none"]: - # If self.scale_rewards = "none", we'll only use std_rewards to check for zero std for logging - if num_generations > 1: - std_rewards = nanstd(rewards.view(-1, num_generations), dim=1) - std_rewards = std_rewards.repeat_interleave(num_generations, dim=0) - else: # doesn't occur during training, but could occur in eval when num_generations_eval=1 - std_rewards = torch.zeros_like(rewards) - elif self.scale_rewards == "batch": - # Compute global std - if rewards.numel() > 1: - std_rewards = nanstd(rewards).expand_as(rewards) - else: # doesn't occur during training, but could occur in eval when num_generations_eval=batch_size=1 - std_rewards = torch.zeros_like(rewards) - else: - raise ValueError( - f"Invalid value for scale_rewards: {self.scale_rewards}. Must be one of 'batch', 'group', or 'none'." - ) - - advantages = rewards - mean_grouped_rewards - if self.scale_rewards != "none": - advantages = advantages / (std_rewards + 1e-4) - is_std_zero = torch.isclose(std_rewards, torch.zeros_like(std_rewards)) # for logging - - elif self.multi_objective_aggregation == "normalize_then_sum": - grouped = rewards_per_func.view(-1, num_generations, len(self.reward_funcs)) - mean_k = torch.nanmean(grouped, dim=1, keepdim=True) - std_k = nanstd(grouped, dim=1, keepdim=True) if num_generations > 1 else torch.zeros_like(mean_k) - reward_k = (grouped - mean_k) / (std_k + 1e-4) - reward_k = reward_k.view(-1, len(self.reward_funcs)) - rewards = (reward_k * self.reward_weights.to(device).unsqueeze(0)).nansum(dim=1) - rewards[unscorable_mask] = torch.nan - std_rewards = nanstd(rewards).expand_as(rewards) if rewards.numel() > 1 else torch.zeros_like(rewards) - advantages = (rewards - torch.nanmean(rewards)) / (std_rewards + 1e-4) - is_std_zero = torch.isclose(std_rewards, torch.zeros_like(std_rewards)) # for logging - - else: - raise ValueError( - f"Invalid multi_objective_aggregation: {self.multi_objective_aggregation}. Must be " - "'sum_then_normalize' or 'normalize_then_sum'." - ) - - # Unscorable completions (every reward func returned None) carry no learning signal: their reward is NaN here, - # so zero their advantage to keep them from moving the policy. - advantages = torch.nan_to_num(advantages, nan=0.0) - - # Slice to keep only the local part of the data - process_slice = slice( - self.accelerator.process_index * len(prompts), - (self.accelerator.process_index + 1) * len(prompts), - ) - all_process_advantages = advantages.clone() # keep the aggregated advantages for logging - advantages = advantages[process_slice] - - # Calculate mean reward per function, but only for samples where the function was applied (non-NaN values) - for i, reward_func_name in enumerate(self.reward_func_names): - mean_rewards = torch.nanmean(rewards_per_func[:, i]).item() - self._metrics[mode][f"rewards/{reward_func_name}/mean"].append(mean_rewards) - std_func_rewards = nanstd(rewards_per_func[:, i]).item() - self._metrics[mode][f"rewards/{reward_func_name}/std"].append(std_func_rewards) - rewards = (rewards_per_func * self.reward_weights.to(rewards_per_func.device).unsqueeze(0)).nansum(dim=1) - rewards[unscorable_mask] = torch.nan # exclude unscorable rows from the logged reward stats - self._metrics[mode]["reward"].append(torch.nanmean(rewards).item()) - self._metrics[mode]["reward_std"].append(nanstd(rewards).item()) - self._metrics[mode]["frac_reward_zero_std"].append(is_std_zero.float().mean().item()) - _unsloth_sample_weight = _unsloth_grpo_sample_weight( - self, mode, rewards_per_func, completion_mask, locals().get('tool_mask'), process_slice, - ) - - # Log prompts and completions - self._logs["prompt"].extend(gather_object(prompts)) - self._logs["completion"].extend(gather_object(completions)) - for i, name in enumerate(self.reward_func_names): - self._logs["rewards"][name].extend(rewards_per_func[:, i].tolist()) - self._logs["advantages"].extend(all_process_advantages.tolist()) - - # Flush user-logged extra columns (from log_extra), gathering across processes. - # Keys must be sorted so that all ranks call gather_object in the same order, otherwise values - # get mis-attributed across columns (dict insertion order may differ between processes). - for column in sorted(set(gather_object(list(self._pending_extra_logs)))): - values = self._pending_extra_logs.get(column, [None] * len(prompts)) - self._logs["extra"][column].extend(gather_object(values)) - self._pending_extra_logs.clear() - - # Flush user-logged metrics (from log_metric), averaging across processes. - # Keys must be sorted so that all ranks call accelerator.gather in the same order, otherwise values - # get mis-attributed across metrics (dict insertion order may differ between processes). - for name in sorted(set(gather_object(list(self._pending_metrics)))): - values = self._pending_metrics.get(name, []) - local_stats = torch.tensor([sum(values), len(values)], dtype=torch.float32, device=device) - total, count = self.accelerator.gather(local_stats).view(-1, 2).sum(dim=0).tolist() - self._metrics[mode][name].append(total / count) - self._pending_metrics.clear() - - if images is not None and self.log_multimodal: - self._logs["images"].extend(gather_object(images)) - - if False and self.use_vllm and self.vllm_importance_sampling_correction: - delta = torch.abs(old_per_token_logps - sampling_per_token_logps) - mask = completion_mask.bool() if tool_mask is None else (completion_mask * tool_mask).bool() - # Tokens vLLM could not score carry NaN, so exclude them rather than let them turn the reported - # divergence into NaN. Counting them as zero instead would understate the divergence. - delta = delta[mask & ~torch.isnan(delta)] - mean_delta = torch.mean(delta) if delta.numel() > 0 else torch.tensor(0.0, device=device) - max_delta = torch.max(delta) if delta.numel() > 0 else torch.tensor(0.0, device=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() - ) - if sequence_level_is: - flat_is_ratio = vllm_importance_sampling_ratio.flatten() - else: - flat_is_ratio = vllm_importance_sampling_ratio[mask] - - min_importance_sampling_ratio = ( - torch.min(flat_is_ratio) if flat_is_ratio.numel() > 0 else torch.tensor(0.0, device=device) - ) - mean_importance_sampling_ratio = ( - torch.mean(flat_is_ratio) if flat_is_ratio.numel() > 0 else torch.tensor(0.0, device=device) - ) - max_importance_sampling_ratio = ( - torch.max(flat_is_ratio) if flat_is_ratio.numel() > 0 else torch.tensor(0.0, device=device) - ) - self._metrics[mode]["sampling/importance_sampling_ratio/min"].append( - nanmin(self.accelerator.gather(min_importance_sampling_ratio)).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( - nanmax(self.accelerator.gather(max_importance_sampling_ratio)).item() - ) - - output = { - "prompt_ids": prompt_ids, - "prompt_mask": prompt_mask, - "completion_ids": completion_ids, - "completion_mask": completion_mask, - "advantages": advantages, - "num_items_in_batch": num_items_in_batch, - } - if old_per_token_logps is not None: - output["old_per_token_logps"] = old_per_token_logps - if _unsloth_rollout_old: - output["unsloth_rollout_logps_as_old"] = torch.ones( - completion_ids.size(0), dtype = torch.bool, device = completion_ids.device - ) - output["sampling_per_token_logps"] = _unsloth_rollout_logps - if False and self.use_vllm and self.vllm_importance_sampling_correction: - output["importance_sampling_ratio"] = vllm_importance_sampling_ratio - if sampling_per_token_logps is not None: - output["sampling_per_token_logps"] = sampling_per_token_logps - if ref_per_token_logps is not None: - output["ref_per_token_logps"] = ref_per_token_logps - if "pixel_values" in forward_kwargs: - output["pixel_values"] = forward_kwargs["pixel_values"] - if "image_grid_thw" in forward_kwargs: - output["image_grid_thw"] = forward_kwargs["image_grid_thw"] - if "pixel_attention_mask" in forward_kwargs: - output["pixel_attention_mask"] = forward_kwargs["pixel_attention_mask"] - if "spatial_shapes" in forward_kwargs: - output["spatial_shapes"] = forward_kwargs["spatial_shapes"] - if "image_sizes" in forward_kwargs: - output["image_sizes"] = forward_kwargs["image_sizes"] - if "token_type_ids" in forward_kwargs: - output["token_type_ids"] = forward_kwargs["token_type_ids"] - if "mm_token_type_ids" in forward_kwargs: - output["mm_token_type_ids"] = forward_kwargs["mm_token_type_ids"] - if "image_position_ids" in forward_kwargs: - output["image_position_ids"] = forward_kwargs["image_position_ids"] - if images is not None: - output["num_images"] = num_images - if num_tiles is not None: - output["num_tiles"] = num_tiles - try: - _unsloth_vision_output = _unsloth_grpo_vision_inputs(forward_kwargs) - except NameError: - _unsloth_vision_output = {} - for _vision_key, _vision_value in _unsloth_vision_output.items(): - if _vision_value is not None and _vision_key not in output: - output[_vision_key] = _vision_value - 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 - if tool_mask is not None: - output["tool_mask"] = tool_mask - if _unsloth_sample_weight is not None: - output["unsloth_sample_weight"] = _unsloth_sample_weight - return output - - 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") - - prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"] - completion_ids, completion_mask = ( - inputs["completion_ids"], - inputs["completion_mask"], - ) - _vision_inputs = _unsloth_grpo_vision_inputs(inputs) - pixel_values = _vision_inputs.get("pixel_values", None) - image_grid_thw = _vision_inputs.get("image_grid_thw", None) - pixel_attention_mask = _vision_inputs.get("pixel_attention_mask", None) - image_sizes = _vision_inputs.get("image_sizes", None) - num_images = _vision_inputs.get("num_images", None) - # Transformers 5.x needs token_type_ids/mm_token_type_ids for some vision models. - token_type_ids = _vision_inputs.get("token_type_ids", None) - mm_token_type_ids = _vision_inputs.get("mm_token_type_ids", None) - video_grid_thw = _vision_inputs.get("video_grid_thw", 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) - _rollout_logps_as_old = inputs.get("unsloth_rollout_logps_as_old", None) is not None - _loss_delta = getattr(self.args, "delta", None) - if _rollout_logps_as_old and _loss_delta is None: - # verl's dual-clip default (clip_ratio_c = 3.0). - _loss_delta = 3.0 - current_gradient_accumulation_steps = _unsloth_grpo_accumulation_steps(self) - 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 - or video_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, - ) - _vision_inputs["mm_token_type_ids"] = mm_token_type_ids - # Only the keys the processor produced: an older zoo must not see unknown kwargs. - _vision_inputs = {k: v for k, v in _vision_inputs.items() if v is not 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 - ) - # KL divergence between model and reference: _prepare_inputs no longer returns reference log probs. See trl grpo_trainer.py#L1328. - ref_logps = inputs.get("ref_per_token_logps", None) - # x - x.detach() preserves gradients from x. - advantages = inputs["advantages"] - old_logps = inputs.get("old_per_token_logps", None) - - input_ids = input_ids[:, -logits_to_keep:] - - model_config = _unsloth_get_model_config(model) - # The old and reference logps come from _get_per_token_logps_and_entropies and the gradient logps from here, so both must read the transforms alike or the importance ratio compares two different policies. - if detect_logit_transforms is not None: - # model_config, not model: see _get_per_token_logps_and_entropies. - _transforms = detect_logit_transforms(model_config) - logit_softcapping = _transforms["logit_softcapping"] - logit_scale_multiply = _transforms["logit_scale_multiply"] - logit_scale_divide = _transforms["logit_scale_divide"] - else: - logit_softcapping = _unsloth_get_final_logit_softcapping(model) # Gemma - logit_scale_multiply, logit_scale_divide = _unsloth_resolve_logit_scales(model_config) - - max_left_pad = inputs.get("max_left_pad", 0) - # importance_sampling_level is absent before TRL 0.20.0 (token level). - importance_sampling_level = getattr( - self, - "importance_sampling_level", - getattr(self.args, "importance_sampling_level", "token"), - ) - 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, - ) - if _rollout_logps_as_old: - old_logps, sampling_per_token_logps = sampling_per_token_logps, None - ( - 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 = importance_sampling_level, - epsilon_low = self.epsilon_low, - epsilon_high = self.epsilon_high, - max_completion_length = self.args.max_completion_length, - delta = _loss_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: - # With beta = 0 and num_iterations = 1 the no-grad path (and its zoo gate) never runs, and an older - # grpo_accumulated_loss silently drops pixel_values. Body-local: copied out without imports (#6960). - if pixel_values is not None and not getattr( - self, "_unsloth_grpo_vision_zoo_checked", False - ): - _grpo_vision_chunks = None - try: - from unsloth_zoo.rl_replacements import ( - grpo_vision_chunks as _grpo_vision_chunks, - ) - except Exception: - pass - if _grpo_vision_chunks is None: - raise RuntimeError( - "Unsloth: vision GRPO needs an unsloth_zoo build that exports " - "grpo_vision_chunks, the shared multimodal key tuple and chunker " - "used by both GRPO logprob paths. Please upgrade unsloth_zoo to " - "2026.9.5 or newer: pip install -U unsloth_zoo" - ) - self._unsloth_grpo_vision_zoo_checked = True - - 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: - # Probe by import: the grep below cries "upgrade unsloth_zoo" falsely. - try: - from unsloth_zoo.rl_replacements import grpo_vision_chunks - _supports_num_images = grpo_vision_chunks is not None - except Exception: - pass - 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 = {} - _sample_weight = inputs.get("unsloth_sample_weight", None) - if _sample_weight is not None: - if "sample_weight" not in inspect.getsource(grpo_compute_loss): - raise RuntimeError( - "UNSLOTH_GRPO_ZERO_VARIANCE_EXCLUDE / UNSLOTH_GRPO_LOSS_AGGREGATION need a " - "newer unsloth_zoo. Please upgrade unsloth_zoo." - ) - _grpo_accumulated_loss_kwargs["sample_weight"] = _sample_weight - if tool_mask is not None: - _grpo_accumulated_loss_kwargs["tool_mask"] = tool_mask - if _rollout_logps_as_old: - _grpo_accumulated_loss_kwargs["rollout_logps_as_old"] = True - 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, - 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 = importance_sampling_level, - epsilon_low = self.epsilon_low, - epsilon_high = self.epsilon_high, - max_completion_length = self.args.max_completion_length, - delta = _loss_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, - **_vision_inputs, - **_grpo_accumulated_loss_kwargs, - ) - else: - # For backwards compatibility with trl 0.15.2 and maybe 0.17. - loss, completion_length, mean_kl, coef_1, completion_mask = grpo_accumulated_loss( - trainer = self, - input_ids = _input_ids, - 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, - **_vision_inputs, - **_grpo_accumulated_loss_kwargs, - ) - if "train" in self._metrics: - # model.training like TRL; should_evaluate is False under eval_on_start and evaluate(). - mode = "train" if self.model.training else "eval" - 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 delta is not None and ( - (self.use_vllm and getattr(self, "vllm_importance_sampling_correction", False)) - or _rollout_logps_as_old - ): - delta, flat_is_ratio = _unsloth_grpo_is_metric_values( - delta, - flat_is_ratio, - completion_mask, - (not _rollout_logps_as_old) - and _unsloth_grpo_is_sequence_level( - getattr(self.args, "vllm_importance_sampling_mode", None) - ), - ) - 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): - # a width-1 tensor is sequence-level only when the completion mask spans multiple tokens - if x.shape[1] == 1 and completion_mask.shape[1] != 1: - 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"]: - 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 - - @staticmethod - def get_off_policy_mask( - advantages: torch.Tensor, - per_token_logps: torch.Tensor, - sampling_per_token_logps: torch.Tensor, - mask: torch.Tensor, - off_policy_threshold: float, - ) -> torch.Tensor: - """ - Computes the Off-Policy Sequence Mask from DeepSeek-V3.2 paper. Returns a (B, 1) tensor where 1.0 indicates - "Keep" and 0.0 indicates "Drop". - """ - # forward KL div: log(pi_old) - log(pi_theta) - kl_div = sampling_per_token_logps - per_token_logps.detach() - # Tokens vLLM could not score carry NaN and provide no KL evidence. Left in, the sequence mean is NaN, - # `avg_seq_kl <= off_policy_threshold` is then False, and a negative-advantage sequence would be dropped - # on the strength of a single unscorable token. - kl_div = torch.nan_to_num(kl_div, nan=0.0) - # Sequence-level Mean KL (ignoring prompt+padding) - seq_kl_sum = (kl_div * mask).sum(dim=1, keepdim=True) - avg_seq_kl = seq_kl_sum / mask.sum(dim=1, keepdim=True).clamp(min=1.0) - # Keep if (Advantage >= 0) OR (KL <= delta) - is_pos_adv = advantages >= 0 - is_low_kl = avg_seq_kl <= off_policy_threshold - return (is_pos_adv | is_low_kl).to(dtype=mask.dtype) # (B, 1) - - @staticmethod - @torch.no_grad() - def get_gamma_weights( - advantages: torch.Tensor, - log_ratio_per_token: torch.Tensor, - mask: torch.Tensor, - importance_sampling_ratio: torch.Tensor | None, # (B, T) - k_pos: float = 2.0, - lambda_pos: float = 3.0, - k_neg: float = 3.0, - lambda_neg: float = 2.0, - ) -> torch.Tensor: - """ - Computes the Gamma weights for the VESPO loss. For reference: - φ(w) = e^λ × w^k × e^{-λw} is the gamma weighting (normalized so φ(1)=1) - with w = sequence-level importance sampling ratio - note: we will compute φ(w) in log space - - φ(w) is detached via @torch.no_grad(), only acts as gradient scaling coefficient - - VESPO loss = -φ(w) × A × log_prob, gradient naturally gives φ(w) × A × ∇log π - """ - # reducing clamp range directly to log(1e-8) ~ -18.42, to avoid recomputing log_w=log(w.clamp(min=1e-8)) later - # This is solely for matching truthfully the original implementation, otherwise keeping -20 could be fine. - lower_clamp = math.log(1e-8) - - # Sequence-level log ratio Σ log(π_θ/π_old) (not a mean like for `log_importance_weights`) - log_ratio_clamped = torch.clamp(log_ratio_per_token, -20.0, 20.0) - seq_log_ratio = torch.sum(log_ratio_clamped * mask, dim=-1, keepdim=True) # (B, 1) - - # Apply token-level TIS or MIS correction (in log space) - if importance_sampling_ratio is not None: - log_is_ratio = torch.clamp(torch.log(importance_sampling_ratio), lower_clamp, 20.0) - # log(w) = log(π_θ/π_old) + log(π_old/π_sampler) - seq_log_ratio += torch.sum(log_is_ratio, dim=-1, keepdim=True) - - log_w_seq = torch.clamp(seq_log_ratio, lower_clamp, 20.0) - w_seq = torch.exp(log_w_seq) - - # compute k and lambda based on advantage sign - is_nonneg_adv = advantages >= 0 - k_seq = torch.where(is_nonneg_adv, k_pos, k_neg) - lambda_seq = torch.where(is_nonneg_adv, lambda_pos, lambda_neg).clamp(min=1e-4) - - # log(φ(w)) = λ + k × log(w) - λ × w - log_phi = lambda_seq + k_seq * log_w_seq - lambda_seq * w_seq - phi_seq = torch.exp(log_phi).nan_to_num(nan=0.0, posinf=0.0, neginf=0.0) - - return phi_seq # (B, 1) - - def _compute_loss(self, model, inputs): - # 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"] - input_ids = torch.cat([prompt_ids, completion_ids], dim=1) - attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) - logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens - mask = completion_mask if "tool_mask" not in inputs else completion_mask * inputs["tool_mask"] - - # Compute the per_token_logps and the entropy at each position in the completion - per_token_logps, entropies, aux_loss = self._get_per_token_logps_and_entropies( - model, - input_ids, - attention_mask, - logits_to_keep, - compute_entropy=True, - compute_aux_loss=self.aux_loss_enabled, - pixel_values=inputs.get("pixel_values"), - image_grid_thw=inputs.get("image_grid_thw"), - num_images=inputs.get("num_images"), - pixel_attention_mask=inputs.get("pixel_attention_mask"), - spatial_shapes=inputs.get("spatial_shapes"), - num_tiles=inputs.get("num_tiles"), - image_sizes=inputs.get("image_sizes"), - token_type_ids=inputs.get("token_type_ids"), - mm_token_type_ids=inputs.get("mm_token_type_ids"), - image_position_ids=inputs.get("image_position_ids"), - ) - - if self.top_entropy_quantile < 1.0: - entropy_mask = self.get_high_entropy_mask(entropies, mask, 1 - self.top_entropy_quantile) - else: - entropy_mask = None - - # Compute the loss - advantages = inputs["advantages"] - # In the base GRPO implementation, advantages are expected to have shape (B,). To support subclasses that - # provide advantages with shape (B, T), we *conditionally* unsqueeze the tensor. - if advantages.dim() == 1: - advantages = advantages.unsqueeze(1) - # When num_iterations == 1 and steps_per_generation <= gradient_accumulation_steps, - # old_per_token_logps == per_token_logps. In this case we can skip its computation - # (see _generate_and_score_completions) and instead use per_token_logps.detach(). - # The exception is when using vLLM, where we always compute old_per_token_logps - # for importance sampling - old_per_token_logps = inputs.get("old_per_token_logps") - old_per_token_logps = per_token_logps.detach() if old_per_token_logps is None else old_per_token_logps - - if self.off_policy_mask_threshold is not None: - # OPSM should use inference-time logprobs to detect both sources of off-policyness: - # 1. Drift from gradient updates (always present) - # 2. Drift from training-inference mismatch (when using vLLM) - # When using vLLM, prioritize sampling_per_token_logps, otherwise use old_per_token_logps - sampling_per_token_logps = inputs.get("sampling_per_token_logps", old_per_token_logps) - - off_policy_mask = self.get_off_policy_mask( - advantages=advantages, - per_token_logps=per_token_logps, - sampling_per_token_logps=sampling_per_token_logps, - mask=mask, - off_policy_threshold=self.off_policy_mask_threshold, - ) - - log_ratio = per_token_logps - old_per_token_logps - if self.importance_sampling_level == "token": - log_importance_weights = log_ratio - elif self.importance_sampling_level == "sequence": - log_importance_weights = (log_ratio * mask).sum(-1) / mask.sum(-1).clamp(min=1.0) - log_importance_weights = log_importance_weights.unsqueeze(-1) - else: - raise ValueError( - f"Unknown importance sampling level: {self.importance_sampling_level}. Possible values are 'token' " - "and 'sequence'." - ) - - coef_1 = torch.exp(log_importance_weights) - - # Compute the KL divergence between the model and the reference model - if self.beta != 0.0: - ref_per_token_logps = inputs["ref_per_token_logps"] - per_token_kl = ( - torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1 - ) - # Importance sampling correction for the KL divergence - if self.args.use_bias_correction_kl: - per_token_kl = per_token_kl * coef_1 - - # From here, log_importance_weights (and all subsequent tensors, coef_1, coef_2, etc.) shape depends on - # importance_sampling_level: "token" level: (B, T); "sequence" level: (B, 1) - if self.loss_type == "cispo": - clamped_ratios = torch.clamp(coef_1, max=self.epsilon_high).detach() - per_token_loss = -clamped_ratios * advantages * per_token_logps - elif self.loss_type in ["grpo", "bnpo", "dr_grpo", "dapo", "luspo"]: - coef_2 = torch.clamp(coef_1, 1 - self.epsilon_low, 1 + self.epsilon_high) - # Two-sided clipping - if self.args.delta is not None: - coef_1 = torch.clamp(coef_1, max=self.args.delta) - - per_token_loss1 = coef_1 * advantages - per_token_loss2 = coef_2 * advantages - per_token_loss = -torch.min(per_token_loss1, per_token_loss2) - elif self.loss_type == "sapo": - temperatures = torch.where(advantages > 0, self.args.sapo_temperature_pos, self.args.sapo_temperature_neg) - soft_coef_1 = torch.sigmoid(temperatures * (coef_1 - 1)) * 4 / temperatures - per_token_loss = -soft_coef_1 * advantages - elif self.loss_type == "vespo": - phi_seq = self.get_gamma_weights( - advantages=advantages, - log_ratio_per_token=log_ratio, - mask=mask, - importance_sampling_ratio=inputs.get("importance_sampling_ratio"), - k_pos=self.args.vespo_k_pos, - lambda_pos=self.args.vespo_lambda_pos, - k_neg=self.args.vespo_k_neg, - lambda_neg=self.args.vespo_lambda_neg, - ) - per_token_loss = -phi_seq * advantages * per_token_logps - else: - raise ValueError(f"Unknown loss type: {self.loss_type}") - - if self.off_policy_mask_threshold is not None: - per_token_loss = per_token_loss * off_policy_mask - - if entropy_mask is not None: - per_token_loss = per_token_loss * entropy_mask - - if self.use_vllm and self.vllm_importance_sampling_correction and self.loss_type != "vespo": - per_token_loss = per_token_loss * inputs["importance_sampling_ratio"] - - if self.beta != 0.0: - per_token_loss = per_token_loss + self.beta * per_token_kl - - mode = "train" if self.model.training else "eval" - if self.loss_type in ["grpo", "sapo"]: - loss = ((per_token_loss * mask).sum(-1) / mask.sum(-1).clamp(min=1.0)).mean() - normalizer = self.current_gradient_accumulation_steps if mode == "train" else 1.0 # no accum in eval - policy_loss = loss.detach() - loss = loss / normalizer - elif self.loss_type == "bnpo": - loss = (per_token_loss * mask).sum() / mask.sum().clamp(min=1.0) - normalizer = self.current_gradient_accumulation_steps if mode == "train" else 1.0 # no accum in eval - policy_loss = loss.detach() - loss = loss / normalizer - elif self.loss_type == "dr_grpo": - loss = (per_token_loss * mask).sum() / (per_token_loss.size(0) * self.max_completion_length) - normalizer = self.current_gradient_accumulation_steps if mode == "train" else 1.0 # no accum in eval - policy_loss = loss.detach() - loss = loss / normalizer - elif self.loss_type in ["cispo", "dapo", "vespo"]: - # `num_items_in_batch` spans the generation batch, so rescale it to one accumulation window - normalizer = inputs["num_items_in_batch"].clamp(min=1.0) / self.accelerator.num_processes - if mode == "train": # in eval, the batch is neither split across steps nor accumulated - normalizer = normalizer * self.current_gradient_accumulation_steps / self.args.steps_per_generation - loss = (per_token_loss * mask).sum() / normalizer - policy_loss = loss.detach() - elif self.loss_type == "luspo": - # `per_token_loss` is (B, 1) only in the recommended sequence-level setup; importance_sampling_level= - # "token" (the config default), the KL term, token-level vLLM IS ratios, and the entropy mask all - # broadcast it to (B, T), so mask before aggregating. - loss = (per_token_loss * mask).sum(-1).mean() - normalizer = self.current_gradient_accumulation_steps if mode == "train" else 1.0 - policy_loss = loss.detach() - loss = loss / normalizer - else: - raise ValueError(f"Unknown loss type: {self.loss_type}") - - # Entropy bonus: add entropy regularization to encourage exploration. _entropy_bonus_enabled is set - # whenever a non-zero static coef is set OR adaptive mode is enabled (adaptive stays enabled even when - # entropy_coef has been decremented to entropy_coef_min so it can recover once entropy drops again). - if self._entropy_bonus_enabled: - # When top_entropy_quantile < 1.0, entropy_mask restricts policy gradients to high-entropy - # tokens. Use the same effective mask for the entropy bonus so it acts on the same tokens. - effective_mask = mask if entropy_mask is None else mask * entropy_mask - # Entropy bonus = mean per-token entropy H (the documented objective L = L_policy - coef * H), so - # H does not depend on how each loss type normalizes its policy term. The bonus is a mean over the - # tokens it acts on (effective_mask), scaled only for gradient accumulation, never by a loss-type- - # specific policy normalizer. (The adaptive controller below tracks a window-global token-weighted mean, - # which can differ from this per-micro-batch mean when token counts vary across micro-batches.) - accumulation_factor = self.current_gradient_accumulation_steps if mode == "train" else 1.0 - entropy_loss = ( - (entropies * effective_mask).sum() / effective_mask.sum().clamp(min=1.0) / accumulation_factor - ) - - # Apply the coefficient and gating from the end of the previous optimizer step, so that every - # micro-batch in the current accumulation window applies the same entropy bonus. The adaptive - # update below only takes effect on the next step. - if self.use_adaptive_entropy: - apply_coef = self.entropy_coef if self._last_world_entropy <= self.args.entropy_target else 0.0 - else: - apply_coef = self.entropy_coef - - loss = loss - apply_coef * entropy_loss - - self._metrics[mode]["policy_loss"].append(self.accelerator.gather(policy_loss).nanmean().item()) - - # Adaptive update. Gated on train mode so evaluation cannot mutate the entropy controller state. - if self.use_adaptive_entropy and mode == "train": - # Accumulate the entropy sum and active-token count of every micro-batch into a running window - # buffer, so the controller measures the exact window-global entropy rather than just the last - # micro-batch (which would be a 1 / gradient_accumulation_steps subsample). - stats = torch.stack([(entropies * effective_mask).sum(), effective_mask.sum()]).detach() - if self._entropy_window_stats is None: - self._entropy_window_stats = stats - else: - self._entropy_window_stats = self._entropy_window_stats + stats - # At the optimizer-step boundary, reduce the window totals across ranks (sum and token count - # jointly, for a true global mean unbiased when ranks have different completion lengths), - # update the coefficient for the next step, then reset the window buffer. - if self.accelerator.sync_gradients: - window_stats = self.accelerator.reduce(self._entropy_window_stats, reduction="sum") - world_entropy = (window_stats[0] / window_stats[1].clamp(min=1.0)).item() - if world_entropy <= self.args.entropy_target: - self.entropy_coef = min( - self.entropy_coef + self.args.entropy_coef_delta, self.args.entropy_coef_max - ) - else: - self.entropy_coef = max( - self.entropy_coef - self.args.entropy_coef_delta, self.args.entropy_coef_min - ) - self._last_world_entropy = world_entropy - self._entropy_window_stats = None - - # Log entropy_coef on train optimizer-step boundaries (constant for static control; updated just - # above for adaptive control). sync_gradients is always True in eval (no accumulation context). - if mode == "train" and self.accelerator.sync_gradients: - self._metrics[mode]["entropy_coef"].append(self.entropy_coef) - - # The policy loss above is scaled for gradient accumulation (HF auto-scaling is off here), so scale aux too - if self.aux_loss_enabled: - normalizer = self.current_gradient_accumulation_steps if mode == "train" else 1.0 - loss = loss + self.router_aux_loss_coef * aux_loss / normalizer - self._metrics[mode]["aux_loss"].append(self.accelerator.gather_for_metrics(aux_loss).mean().item()) - - # Log the metrics - def masked_seq_mean(x): - if x.shape[1] == 1: # when importance_sampling_level == "sequence": already one value per sequence - return x.squeeze(1) - return (x * mask).sum(-1) / mask.sum(-1) - - def global_masked_mean(x): - if x.shape[1] == 1: # when importance_sampling_level == "sequence": one value per sequence - local_sum, local_count = x.sum(), torch.tensor(float(x.shape[0]), device=x.device) - else: - local_sum, local_count = (x * mask).sum(), mask.sum().float() - totals = self.accelerator.reduce(torch.stack([local_sum, local_count]), reduction="sum") - return (totals[0] / totals[1].clamp(min=1.0)).item() - - if self.beta != 0.0: - self._metrics[mode]["kl"].append(global_masked_mean(per_token_kl)) - - self._metrics[mode]["entropy"].append(global_masked_mean(entropies)) - - if self.loss_type in ["grpo", "bnpo", "dr_grpo", "dapo", "luspo"]: - # 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 - self._metrics[mode]["clip_ratio/low_mean"].append(global_masked_mean(is_low_clipped.float())) - self._metrics[mode]["clip_ratio/high_mean"].append(global_masked_mean(is_high_clipped.float())) - self._metrics[mode]["clip_ratio/region_mean"].append(global_masked_mean(is_region_clipped.float())) - gathered_low_clip = self.accelerator.gather(masked_seq_mean(is_low_clipped.float())) - self._metrics[mode]["clip_ratio/low_min"].append(nanmin(gathered_low_clip).item()) - gathered_high_clip = self.accelerator.gather(masked_seq_mean(is_high_clipped.float())) - self._metrics[mode]["clip_ratio/high_max"].append(nanmax(gathered_high_clip).item()) - elif self.loss_type == "cispo": - is_cispo_clipped = (coef_1 > self.epsilon_high) & (advantages > 0) - self._metrics[mode]["cispo_clip_ratio"].append(global_masked_mean(is_cispo_clipped.float())) - elif self.loss_type == "vespo": - self._metrics[mode]["vespo/phi_seq_mean"].append(global_masked_mean(phi_seq)) - - return loss - - # During eval, Trainer calls prediction_step. If no labels are present in the inputs, it only runs forward and - # returns logits. We override prediction_step to force compute_loss, because this trainer doesn't involve labels. - def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys: list[str] | None = None): - inputs = self._prepare_inputs(inputs) - with torch.no_grad(): - with self.compute_loss_context_manager(): - loss = self.compute_loss(model, inputs) - loss = loss.mean().detach() - return loss, None, None - - def log(self, logs: dict[str, float], start_time: float | None = None) -> None: - mode = "train" if self.model.training else "eval" - # Average the metrics - metrics = {} - for key, val in self._metrics[mode].items(): - # Filter out NaN values before averaging. A reward function that returns None for all samples - # in a batch produces NaN for that batch's metric. With logging_steps > 1, a naive sum()/len() - # would let a single NaN contaminate valid data from other batches. Only return None when no - # valid values remain (e.g. JSON loggers crash on float NaN). - valid = [v for v in val if not math.isnan(v)] - metrics[key] = sum(valid) / len(valid) if valid else None - - # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` - # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. - if mode == "eval": - metrics = {f"eval_{key}": val for key, val in metrics.items()} - - logs.update(metrics) - super().log(logs, start_time) - self._metrics[mode].clear() - - if self.accelerator.is_main_process and self.log_completions: - if is_rich_available(): - print_prompt_completions_sample( - self._logs["prompt"], - self._logs["completion"], - self._logs["rewards"], - self._logs["advantages"], - self.state.global_step, - self.num_completions_to_print, - extra=dict(self._logs["extra"]), - ) - - logging_backends = [] - if self.args.report_to and "wandb" in self.args.report_to and wandb.run is not None: - logging_backends.append(wandb) - if self.args.report_to and "trackio" in self.args.report_to: - logging_backends.append(trackio) - - table = { - "step": [self.state.global_step] * len(self._logs["prompt"]), - "prompt": [strip_images_from_messages(messages) for messages in self._logs["prompt"]], - "completion": [strip_images_from_messages(messages) for messages in self._logs["completion"]], - **self._logs["rewards"], - **self._logs["extra"], - "advantage": self._logs["advantages"], - } - - df_base = pd.DataFrame(table) - df_base.to_parquet( - os.path.join( - self.args.output_dir, - "completions", - f"completions_{self.state.global_step:05d}.parquet", - ) - ) - - images_raw = self._logs["images"] or [] - - for logging_backend in logging_backends: - if images_raw: - images = [] - for image_list in self._logs["images"]: - if image_list: - images.append([logging_backend.Image(image) for image in image_list]) - else: - images.append([]) - df = pd.concat( - [df_base, pd.Series(images, name="image")], - axis=1, - copy=False, - ) - else: - df = df_base - - if self.log_unique_prompts: - # Prompts are conversations (lists of dicts), so compare them by their text - df = df[~df["prompt"].astype(str).duplicated()] - - logging_backend.log({"completions": logging_backend.Table(dataframe=df)}) - - # Ensure the model card is saved along with the checkpoint - def _save_checkpoint(self, model, trial): - if self.args.hub_model_id is None: - model_name = Path(self.args.output_dir).name - else: - model_name = self.args.hub_model_id.split("/")[-1] - self.create_model_card(model_name=model_name) - super()._save_checkpoint(model, trial) - if self.use_adaptive_entropy and self.args.should_save: - checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}" - output_dir = os.path.join(self._get_output_dir(trial=trial), checkpoint_folder) - with open(os.path.join(output_dir, "entropy_ctrl_state.json"), "w") as f: - json.dump({"entropy_coef": self.entropy_coef, "last_world_entropy": self._last_world_entropy}, f) - - def _load_optimizer_and_scheduler(self, checkpoint): - super()._load_optimizer_and_scheduler(checkpoint) - if self.use_adaptive_entropy and checkpoint is not None: - path = os.path.join(checkpoint, "entropy_ctrl_state.json") - if os.path.exists(path): - with open(path) as f: - state = json.load(f) - self.entropy_coef = state["entropy_coef"] - self._last_world_entropy = state["last_world_entropy"] -class UnslothGRPOTrainer(_UnslothGRPOTrainer): - """ - - Trainer for the Group Relative Policy Optimization (GRPO) method. This algorithm was initially proposed in the - paper [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language - Models](https://huggingface.co/papers/2402.03300). - - Example: - - ```python - >>> from trl import GRPOTrainer - >>> from trl.rewards import accuracy_reward - >>> from datasets import load_dataset - - >>> dataset = load_dataset("trl-lib/DeepMath-103K", split="train") - - >>> trainer = GRPOTrainer( - ... model="Qwen/Qwen2.5-0.5B-Instruct", - ... reward_funcs=accuracy_reward, - ... train_dataset=dataset, - ... ) - >>> trainer.train() - ``` - - Args: - model (`str` or [`~transformers.PreTrainedModel`] or [`~peft.PeftModel`]): - Model to be trained. Can be either: - - - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a - path to a *directory* containing model weights saved using - [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded - using `.from_pretrained` (where `` is derived from the model - config) with the keyword arguments in `args.model_init_kwargs`. If `dtype` is not specified in - `args.model_init_kwargs`, it defaults to `float32`. This differs from - [`~transformers.PreTrainedModel.from_pretrained`], where (since Transformers v5) the dtype is inferred - from the model config. - - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. - - A [`~peft.PeftModel`] object. Only causal language models are supported. - reward_funcs (`RewardFunc | list[RewardFunc]`, *optional*): - Reward functions to be used for computing the rewards. To compute the rewards, we call all the reward - functions with the prompts and completions and sum the rewards. May be omitted when the reward is supplied - by the environment through `environment_factory` (see below). Can be either: - - - A single reward function, such as: - - A string: The *model ID* of a pretrained model hosted inside a model repo on huggingface.co, or a - path to a *directory* containing model weights saved using - [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded - using [`~transformers.AutoModelForSequenceClassification.from_pretrained`] with `num_labels=1` and the - keyword arguments in `args.model_init_kwargs`. - - A [`~transformers.PreTrainedModel`] object: Only sequence classification models are supported. - - A custom reward function: The function is provided with the prompts and the generated completions, - plus any additional columns in the dataset. It should return a list of rewards. Custom reward - functions can be either synchronous or asynchronous and can also return `None` when the reward is - not applicable to those samples. This is useful for multi-task training where different reward - functions apply to different types of samples. When a reward function returns `None` for a sample, - that reward function is excluded from the reward calculation for that sample. For more details, see - [Using a custom reward - function](#using-a-custom-reward-function). - - The trainer's state is also passed to the reward function. The trainer's state is an instance of - [`~transformers.TrainerState`] and can be accessed by accessing the `trainer_state` argument to the - reward function's signature. - - A list of reward functions, where each item can independently be any of the above types. Mixing different - types within the list (e.g., a string model ID and a custom reward function) is allowed. - args ([`GRPOConfig`], *optional*): - Configuration for this trainer. If `None`, a default configuration is used. - train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`], *optional*): - Dataset to use for training. It must include a column `"prompt"`. Any additional columns in the dataset is - ignored. The format of the samples can be either: - - - [Standard](dataset_formats#standard): Each sample contains plain text. - - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role - and content). - - May be omitted only when an `environment_factory` is provided and the environment owns (or procedurally - generates) the data, returning the prompt from its `reset()` method. In that case, `max_steps` must be set - to define the training length. - - When `train_dataset` is an [`~datasets.IterableDataset`] (e.g. a streaming dataset), `max_steps` must be - set in the training arguments, since its length cannot be inferred and the total number of training steps - is required to bound the training loop and configure the learning rate scheduler. - eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`], [`~datasets.DatasetDict`], [`~datasets.IterableDatasetDict`] or `dict[str, Dataset | IterableDataset]`): - Dataset to use for evaluation. It must meet the same requirements as `train_dataset`. - processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.ProcessorMixin`], *optional*): - Processing class used to process the data. The padding side must be set to "left". If `None`, the - processing class is loaded from the model's name with [`~transformers.AutoProcessor.from_pretrained`]. A - padding token, `tokenizer.pad_token`, must be set. If the processing class has not set a padding token, - `tokenizer.eos_token` will be used as the default. - reward_processing_classes ([`~transformers.PreTrainedTokenizerBase`] or `list[PreTrainedTokenizerBase]`, *optional*): - Processing classes corresponding to the reward functions specified in `reward_funcs`. Can be either: - - - A single processing class: Used when `reward_funcs` contains only one reward function. - - A list of processing classes: Must match the order and length of the reward functions in `reward_funcs`. - If set to `None`, or if an element of the list corresponding to a [`~transformers.PreTrainedModel`] is - `None`, the tokenizer for the model is automatically loaded using - [`~transformers.AutoTokenizer.from_pretrained`]. For elements in `reward_funcs` that are custom reward - functions (not [`~transformers.PreTrainedModel`]), the corresponding entries in `reward_processing_classes` - are ignored. - callbacks (list of [`~transformers.TrainerCallback`], *optional*): - List of callbacks to customize the training loop. Will add those to the list of default callbacks detailed - in [here](https://huggingface.co/docs/transformers/main_classes/callback). - - If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`] - method. - optimizers (`tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None]`, *optional*, defaults to `(None, None)`): - A tuple containing the optimizer and the scheduler to use. Will default to an instance of `AdamW` on your - model and a scheduler given by [`~transformers.get_linear_schedule_with_warmup`] controlled by `args`. - quantization_config ([`~transformers.BitsAndBytesConfig`], *optional*): - Quantization configuration used when loading the model from a model identifier. Combine with `peft_config` - for QLoRA training. Ignored if the model is already instantiated. - peft_config ([`~peft.PeftConfig`], *optional*): - PEFT configuration used to wrap the model. If `None`, the model is not wrapped. - tools (list of `Callable`, *optional*): - A list of callable tool functions (sync or async) that the model can invoke during generation. Each tool - should be a standard Python function with properly type-hinted arguments and return values, and a - Google-style docstring describing its purpose, arguments, and return value. For more details, see: - https://huggingface.co/docs/transformers/en/chat_extras#passing-tools. The model uses the function's name, - type hints, and docstring to determine how to call it. Ensure that the model's chat template supports tool - use and that it has been fine-tuned for tool calling. - rollout_func (`RolloutFunc`, *optional*): - Function to use for generating completions. It receives the list of prompts allocated to the current - process and the trainer instance. It must return a dict with `"prompt_ids"`, `"completion_ids"`, and - `"logprobs"` fields, and can optionally return `"logprob_token_ids"` (same shape as `"logprobs"`). Any - other fields are forwarded to the reward functions. The function receives the raw per-process prompt slice - with no duplication; it is responsible for returning the correct number of completions per prompt (see - `num_generations` / `num_generations_eval` on the trainer). This feature is experimental and may change or - be removed at any time without prior notice. - environment_factory (`EnvironmentFactory` or `dict[str, EnvironmentFactory]`, *optional*): - A callable that creates and returns an environment instance, or a dictionary mapping environment names to - such callables. The environment class should define methods that can be invoked as tools during generation. - Each method should comply with the same requirements as the `tools` described above. The environment must - also implement a callable `reset` method that can be used to reset state between generations. The `reset` - method should return either `None` or a string: when it returns a string, that string is appended to the - last user message before generation. The environment may also define a `get_reward` method taking no - argument and returning a `float`: when present, the environment owns the reward, and `get_reward` is called - once per completed rollout to score it from the environment's internal state. It acts as an additional - reward source (with weight 1, logged under the environment's class name) alongside `reward_funcs`, which - then becomes optional. - - With a single callable, every example uses the same environment, with one instance per rollout so their - interactions stay isolated. With a dictionary, each example must carry an `environment` field selecting its - environment by name, and only that environment's tools are exposed in its prompt — letting a single run mix - tasks (e.g. a coding environment and a game). This feature is experimental and may change or be removed at - any time without prior notice. - - """ - def __init__( - self, - model, - reward_funcs = None, - args = None, - train_dataset = None, - eval_dataset = None, - processing_class = None, - reward_processing_classes = None, - callbacks = None, - quantization_config = None, - peft_config = None, - tools = None, - rollout_func = None, - environment_factory = None, - **kwargs - ): - if args is None: args = UnslothGRPOConfig() - _unsloth_precision_requested = tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) - use_bf16 = getattr(args, 'bf16', False) - if type(use_bf16) is not bool: use_bf16 = False - use_fp16 = getattr(args, 'fp16', False) - if type(use_fp16) is not bool: use_fp16 = False - force_float32 = False - try: - from unsloth_zoo.device_type import device_is_bf16_supported as _bf16_supported - except Exception: - _bf16_supported = torch.cuda.is_bf16_supported - try: - from unsloth_zoo.device_type import DEVICE_TYPE as _unsloth_device_type - except Exception: - _unsloth_device_type = None - full_finetuning = getattr(model, '_unsloth_full_finetuning', None) - if full_finetuning is None: full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1' - model_forced_float32 = getattr(model, '_unsloth_forced_float32', None) - if model_forced_float32 is None: model_forced_float32 = os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1' - if model_forced_float32 and not (full_finetuning and _bf16_supported()): - print('Unsloth: Switching to float32 training since model cannot work with float16') - force_float32 = True - mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') - dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None) - if dtype is None: dtype = model.get_input_embeddings().weight.dtype - from unsloth_zoo.utils import _get_dtype - dtype = _get_dtype(dtype) - float16 = dtype == torch.float16 - bfloat16 = dtype == torch.bfloat16 - float32 = dtype == torch.float32 - user_float32 = bool(getattr(model, '_unsloth_user_float32', False)) - if full_finetuning: - if bfloat16 and use_fp16: use_fp16 = False - if float16 and use_bf16: use_bf16 = False - if not force_float32 and bfloat16 and use_fp16 and _unsloth_device_type == 'xpu' and not _bf16_supported(): - print('Unsloth: The model was loaded in bfloat16 on purpose, so training uses bfloat16 mixed precision instead of float16.') - use_fp16 = False - use_bf16 = True - args.fp16 = False - args.bf16 = True - if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`') - if not force_float32 and (bfloat16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`') - if force_float32: - # Forced float32 training - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32' and float32 and user_float32 and not _bf16_supported(): - print('Unsloth: Model is in float32 and this GPU has no bfloat16 support, so training stays in float32. Pass fp16 = True to force float16 mixed precision instead.') - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32': - # Mixed precision training. bf16 only if the GPU supports it; V100/T4 use fp16. - use_bf16_amp = (not float16) and (_bf16_supported() or (bfloat16 and _unsloth_device_type == 'xpu')) - args.fp16 = not use_bf16_amp - args.bf16 = use_bf16_amp - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16_amp else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16_amp else 'fp16' - # args.mixed_precision is a new argument which needs to be set now - elif mixed_precision_dtype == 'bfloat16': - # Both False since bfloat16 full finetuning doesn't do any autocasting. - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif use_bf16 or use_fp16: - # transformers <5 exported this itself from fp16/bf16; 5.x dropped the write, so an - # explicit flag left it unset and GRPO readers defaulted to 'fp16', wrapping a - # bfloat16 model in a float16 autocast. See unslothai/unsloth#4891. - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16 else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16 else 'fp16' - - if getattr(args, 'eval_dataset', None) is not None and getattr(args, 'eval_strategy', 'no') == 'no': - args.eval_strategy = 'steps' - if getattr(args, 'eval_steps', None) is None: args.eval_steps = 0.1 - ga_steps = getattr(args, 'gradient_accumulation_steps', None) - if ga_steps is not None and ga_steps > 1: - from transformers import __version__ as transformers_version - if Version(transformers_version) <= Version('4.45.2'): - print('**** Unsloth: Please use our fixed gradient_accumulation_steps by updating transformers, TRL and Unsloth!\n' - '`pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth transformers trl unsloth_zoo`') - if getattr(args, 'eval_strategy', 'no') != 'no': - eval_bsz = getattr(args, 'per_device_eval_batch_size', 8) - eval_generations = getattr(args, 'num_generations_eval', None) or getattr(args, 'num_generations', None) or 1 - if eval_bsz == 8 and args.per_device_train_batch_size < eval_bsz and (args.per_device_train_batch_size * args.world_size) % eval_generations == 0: args.per_device_eval_batch_size = args.per_device_train_batch_size - if getattr(args, 'eval_accumulation_steps', None) is None and ga_steps is not None: args.eval_accumulation_steps = ga_steps - fp16_full_eval = getattr(args, 'fp16_full_eval', False) - if type(fp16_full_eval) is not bool: fp16_full_eval = False - bf16_full_eval = getattr(args, 'bf16_full_eval', False) - if type(bf16_full_eval) is not bool: bf16_full_eval = False - if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True - if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False - if force_float32: - args.bf16_full_eval = False - args.fp16_full_eval = False - elif os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16': - args.bf16_full_eval = True - args.fp16_full_eval = False - elif not bf16_full_eval and not fp16_full_eval: - args.bf16_full_eval = args.bf16 - args.fp16_full_eval = args.fp16 - if tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) != _unsloth_precision_requested: - from unsloth.models.loader_utils import sync_deepspeed_mixed_precision - sync_deepspeed_mixed_precision(args) - if model is not None: - from unsloth.models.loader_utils import apply_deepspeed_zero3_model_setup - apply_deepspeed_zero3_model_setup(model) - if model is not None: - _warnings_issued = getattr(model, 'warnings_issued', None) - if _warnings_issued is None: - model.warnings_issued = {} - elif not isinstance(_warnings_issued, dict): - try: - model.warnings_issued = dict(_warnings_issued) - except Exception: - model.warnings_issued = {} - if 'max_seq_length' not in locals() and not hasattr(args, 'max_seq_length'): - pass - else: - model_max_seq_length = getattr(model, 'max_seq_length', None) - args_max_seq_length = getattr(args, 'max_seq_length', None) - if args_max_seq_length is None and model_max_seq_length is not None: - max_seq_length = model.max_seq_length - if hasattr(args, 'max_seq_length'): args.max_seq_length = max_seq_length - elif args_max_seq_length is not None and model_max_seq_length is not None: - if args_max_seq_length > model_max_seq_length: - print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but ' - 'the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.') - args.max_seq_length = model_max_seq_length - if model is not None and hasattr(model, 'for_training'): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right' - if 'processing_class' in locals(): - if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right' - if hasattr(processing_class, 'tokenizer') and hasattr(processing_class.tokenizer, 'padding_side'): processing_class.tokenizer.padding_side = 'right' - other_metrics = [] - if not isinstance(reward_funcs, list): _reward_funcs = [reward_funcs] - else: _reward_funcs = reward_funcs - for reward_func in _reward_funcs: - try: - reward_func_name = reward_func.__name__ - if True: - other_metrics.append(f'rewards/{reward_func_name}/mean') - if True: - other_metrics.append(f'rewards/{reward_func_name}/std') - if False: - other_metrics.append(f'rewards/{reward_func_name}') - except: pass - - from unsloth_zoo.logging_utils import PatchRLStatistics - PatchRLStatistics('grpo_trainer', other_metrics) - - # [TODO] Fix up DataParallel multiplying batch sizes - # [TODO] DDP works, but DP seems to not work? [TODO] - if getattr(args, "parallel_mode", None) == ParallelMode.NOT_DISTRIBUTED and args.n_gpu > 1: - if getattr(args, "_n_gpu", 1) != 1: - args._n_gpu = 1 - if "model" in locals(): - try: - from unsloth.models._gc_config_default import honour_gradient_checkpointing_off - honour_gradient_checkpointing_off(args, model) - except Exception: - pass - if "model" in locals() and hasattr(model, "for_training"): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if "model" in locals(): - try: - from unsloth.models._ddp_defaults import unsloth_default_ddp_find_unused_parameters - unsloth_default_ddp_find_unused_parameters(args, model) - except Exception: - pass - super().__init__( - model = model, - reward_funcs = reward_funcs, - args = args, - train_dataset = train_dataset, - eval_dataset = eval_dataset, - processing_class = processing_class, - reward_processing_classes = reward_processing_classes, - callbacks = callbacks, - quantization_config = quantization_config, - peft_config = peft_config, - tools = tools, - rollout_func = rollout_func, - environment_factory = environment_factory,**kwargs) - # A tensor parallel model (tp_native.tp_trainer_setup) stays on its training path: for_inference would - # route the decoder layers onto decode paths outside the tensor parallel regions. - if "model" in locals() and hasattr(model, "for_inference") and not getattr(self, "_unsloth_tensor_parallel", False): - model.for_inference() - if hasattr(self, 'neftune_hook_handle'): - self.neftune_hook_handle.remove() - if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle - if getattr(args, 'neftune_noise_alpha', None) is not None: - model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha - pass - if hasattr(self, 'accelerator'): - scaler = self.accelerator.scaler - current_model = model - while hasattr(current_model, 'model'): - current_model.accelerator_scaler = scaler - current_model = current_model.model - current_model.accelerator_scaler = scaler - pass - if hasattr(self, 'train'): - self.train = MethodType(prepare_for_training_mode(self.__class__.train), self) - pass - if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'): - _vllm_tok = self.llm.get_tokenizer() - _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None) - if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None: - _vllm_tok.chat_template = _pc.chat_template - pass - if getattr(self, 'aux_loss_enabled', False) and hasattr(getattr(self, 'model', None), 'config'): - _text_config = self.model.config - if hasattr(_text_config, 'get_text_config'): _text_config = _text_config.get_text_config() - _n_experts = [getattr(_text_config, _k) for _k in ('num_local_experts', 'num_experts', 'n_routed_experts', 'moe_num_experts') if isinstance(getattr(_text_config, _k, None), int)] - if _n_experts and all(_n <= 1 for _n in _n_experts): - self.aux_loss_enabled = False - _text_config.output_router_logits = False - pass - if getattr(self, 'vllm_generation', None) is not None: - self.vllm_generation._unsloth_vllm_sampling_params = getattr(getattr(self, 'args', None), 'vllm_sampling_params', None) - pass - if os.environ.get('UNSLOTH_GRPO_ROUTING_REPLAY', '0') == '1' and hasattr(model, 'vllm_engine'): - try: - from unsloth_zoo.routing_replay import install_routing_replay - except ImportError: - print('Unsloth: UNSLOTH_GRPO_ROUTING_REPLAY=1 needs a newer unsloth_zoo: pip install -U unsloth_zoo') - else: - from unsloth.models._routing_replay import routing_replay_hook_snapshot, prepare_routing_replay - # TRL scores the PEFT reference through self.model.disable_adapter(), so prepare that object. - _unsloth_replay_model = getattr(self, 'model', None) - _unsloth_replay_model = model if _unsloth_replay_model is None else _unsloth_replay_model - _unsloth_replay_hooks_before = routing_replay_hook_snapshot(_unsloth_replay_model) - install_routing_replay(model, model.vllm_engine, getattr(self, 'processing_class', None)) - # Pick the router weight formulas outside compilation; PEFT reference scoring routes naturally. - prepare_routing_replay(_unsloth_replay_model, _unsloth_replay_hooks_before) - pass - -pass - - -if hasattr(logger, "addFilter"): - import logging - class HideLoggingMessage(logging.Filter): - def __init__(self, text): self.text = text - def filter(self, x): return not (self.text in x.getMessage()) - pass - logger.addFilter(HideLoggingMessage("`use_cache=True`")) - diff --git a/unsloth_compiled_cache/UnslothKTOTrainer.py b/unsloth_compiled_cache/UnslothKTOTrainer.py deleted file mode 100644 index bb09abf50358b89e5a6d057e30c2ecca9c95b94c..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/UnslothKTOTrainer.py +++ /dev/null @@ -1,2947 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from unsloth_zoo.temporary_patches.common import torch_compile -from unsloth_zoo.temporary_patches.common import _maybe_compile -import functools -import transformers -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from trl.trainer.kto_trainer import (Any, AutoProcessor, BaseTunerLayer, BitsAndBytesConfig, Callable, DataCollator, DataCollatorForUnpairedPreference, DataCollatorForVisionUnpairedPreference, DataLoader, Dataset, DatasetDict, EvalPrediction, F, Hasher, IterableDataset, IterableDatasetDict, KTOConfig, KTOTrainer, LoraConfig, Path, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, Sampler, SequentialSampler, SyncRefModelCallback, TrainerCallback, Version, _BaseTrainer, _tokenize, accelerate, add_fused_lm_head, apply_chat_template, broadcast_object_list, common_prefix_length, concatenate_datasets, contextlib, create_model_from_path, dataclass, defaultdict, disable_dropout_in_model, disable_gradient_checkpointing, extract_prompt, flush_left, get_act_offloading_ctx_manager, get_config_model_id, get_peft_model, global_then_local_main_first, has_length, hash_module, is_conversational, is_peft_available, is_peft_model, json, logger, os, pad, peft, prepare_deepspeed, prepare_fsdp, set_seed, textwrap, torch, tqdm, transformers, unpair_preference_dataset, use_adapter, warnings, AutoProcessor, BaseTunerLayer, BitsAndBytesConfig, Callable, DataCollator, DataCollatorForUnpairedPreference, DataCollatorForVisionUnpairedPreference, Dataset, DatasetDict, EvalPrediction, F, IterableDataset, IterableDatasetDict, KTOConfig, KTOTrainer, LoraConfig, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, SyncRefModelCallback, TrainerCallback, Version, accelerate, add_fused_lm_head, contextlib, create_model_from_path, defaultdict, disable_dropout_in_model, get_act_offloading_ctx_manager, get_config_model_id, get_peft_model, hash_module, is_peft_available, is_peft_model, logger, os, pad, peft, prepare_deepspeed, prepare_fsdp, set_seed, torch, transformers, unpair_preference_dataset, warnings, F, PeftModel, PreTrainedModel, is_peft_available, logger, os, peft, torch) - - -import os -import math -import logging -from typing import * -from dataclasses import dataclass, field -from packaging.version import Version -import torch -import numpy as np -from contextlib import nullcontext -from torch.nn import functional as F -import inspect -from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling -from transformers.training_args import ParallelMode -from unsloth_zoo.device_type import DEVICE_TYPE, DEVICE_TYPE_TORCH, device_synchronize - -# Wrap trainer with padding to right and enable training mode -import functools -from types import MethodType -try: - from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers -except: - def reset_unsloth_gradient_checkpointing_buffers(): pass -# Canonical reset lives in unsloth.models._utils so the SFT auto-packing wrapper and the plain -# Trainer loop can import the same helper; fall back to a no-op only if it can't be imported. -try: - from unsloth.models._utils import _unsloth_reset_stray_compile_cache -except Exception: - def _unsloth_reset_stray_compile_cache(self): pass -try: - from unsloth.models._utils import _unsloth_dataset_column_names -except Exception: - def _unsloth_dataset_column_names(dataset): return dataset.column_names -# Drops/renames config arguments the installed TRL no longer accepts, so a -# script pinned to an older TRL keeps working after an upgrade. Falls back to -# the historical raw passthrough so this can never break trainer construction. -try: - from unsloth.models.rl_config_compat import filter_config_init_kwargs as _unsloth_filter_config_init_kwargs - # A cache file generated here can be imported by an older Unsloth whose filter - # predates `mirrored_from`, so drop the argument rather than raise TypeError. - if "mirrored_from" not in inspect.signature(_unsloth_filter_config_init_kwargs).parameters: - _unsloth_filter_config_init_kwargs_old = _unsloth_filter_config_init_kwargs - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): - return _unsloth_filter_config_init_kwargs_old(config_class, kwargs) -except Exception: - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): return kwargs -def prepare_for_training_mode(f): - @functools.wraps(f) - def wrapper(self, *args, **kwargs): - # Drop any torch.compile graph cache poisoned by a stray pre-train forward. - try: - _unsloth_reset_stray_compile_cache(self) - except Exception: - pass - # Finish the previous W&B run if this is a subsequent train() call. - # We do this at the START of train() (not the end) so that - # evaluate() / log() still work after train() completes. - # HF's WandbCallback.setup() will call wandb.init() for the new run. - # See: https://github.com/unslothai/unsloth/issues/3954 - if getattr(self, '_unsloth_training_completed', False): - try: - import wandb - if wandb.run is not None: - wandb.finish() - # Reset HF's WandbCallback so it calls wandb.init() for the new run - for cb in self.callback_handler.callbacks: - if type(cb).__name__ == 'WandbCallback': - cb._initialized = False - break - except: - pass - # Enable training mode - _was_training = None - # Restore the GC mode the model was configured with at setup; fall back to - # the training args only when it wasn't recorded (issue #4735). Use hasattr, - # not a None sentinel, so a deliberately-recorded None is restored verbatim. - _model = getattr(self, 'model', None) - if hasattr(_model, '_unsloth_gradient_checkpointing'): - use_gc = _model._unsloth_gradient_checkpointing - else: - use_gc = getattr(self.args, 'gradient_checkpointing', True) - if hasattr(self, 'model') and hasattr(self.model, "training"): - _was_training = self.model.training - if hasattr(self, 'model') and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - output = f(self, *args, **kwargs) - # Restore previous mode when possible - if hasattr(self, 'model') and hasattr(self.model, "for_inference"): - if _was_training is False: - self.model.for_inference() - elif _was_training is True and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - # Reset gradient checkpointing buffers to free memory while staying ready for next run - try: - reset_unsloth_gradient_checkpointing_buffers() - except: - pass - # Mark that training completed so the next train() call can - # finish this W&B run before starting a new one - self._unsloth_training_completed = True - return output - return wrapper -pass - -torch_compile_options = { - "epilogue_fusion" : True, - "max_autotune" : False, - "shape_padding" : True, - "trace.enabled" : False, - "triton.cudagraphs" : False, -} - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_hidden_states_selective_log_softmax( - hidden_states: torch.Tensor, - lm_head: torch.Tensor, - index: torch.Tensor, - chunks: int = 4, - logit_scale_multiply: float = 0.0, - logit_scale_divide: float = 0.0, - logit_softcapping: float = 0.0, - temperature: float = 1.0, - # Rows per chunk cap. Read HERE, in the default, not in the body: the body - # is traced with fullgraph = True, and `os.environ` is an unsupported op - # there on torch 2.4 -- Dynamo raises - # torch._dynamo.exc.Unsupported: const method call bytes.decode - # from os._Environ.__getitem__, which the eager fallback does not catch - # (it only catches recompile-limit and disabled-hook breaks), so the very - # first call would die even with the variable unset. A default is evaluated - # once when the def runs, which is import time, outside any traced region. - # It is also a plain int argument, so Dynamo guards on it instead of - # constant-folding an unguarded read (2.7+ never notice a later change). - # A non-numeric value is ignored rather than raised on, so a typo cannot - # break the import. 0 keeps the previous chunk boundaries exactly. - max_rows_per_chunk: int = ( - int(os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip()) - if os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip().isdigit() - else 0 - ), - lm_head_bias: Optional[torch.Tensor] = None, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, - # Read in the default for the same Dynamo reason as max_rows_per_chunk. Pairs with vLLM's - # head_dtype=float32 that load_vllm requests under the same flag. - fp32_lm_head: bool = ( - os.environ.get("UNSLOTH_VLLM_FP32_LM_HEAD", "0").strip().lower() in ("1", "true", "on", "yes") - ), - # Older Unsloth scores old / ref log-probs here and the policy with the fused kernel: one scorer keeps ratio at 1. - fused_linear: bool = ( - os.environ.get("UNSLOTH_FUSED_LINEAR", "1") != "0" - and os.environ.get("UNSLOTH_GRPO_CHUNKED_LOGPROBS", "1") != "0" - ), -) -> torch.Tensor: - # All Unsloth Zoo code licensed under AGPL3 - from unsloth_zoo.fused_losses import chunked_logprob as _chunked_logprob - if (fused_linear and _chunked_logprob._FUSED_LINEAR and not fp32_lm_head - and support_top_k == 0 and support_top_p == 1.0 and support_min_p == 0.0 - and support_excluded_ids is None): - return _chunked_logprob.unsloth_chunked_logprobs( - hidden_states, lm_head, index, - temperature = temperature, - logit_scale_multiply = logit_scale_multiply, - logit_scale_divide = logit_scale_divide, - logit_softcapping = logit_softcapping, - lm_head_bias = lm_head_bias, - max_rows_per_chunk = max_rows_per_chunk if max_rows_per_chunk > 0 else None, - ).reshape((hidden_states.shape[0], hidden_states.shape[1])) - # Reshape on this tensor's own last dim: a no-op, so a wrong-width caller - # cannot have its row count silently rewritten and instead fails at the - # matmul below, which prints both operands. Do not swap in a bare - # torch._check: it reports only "Expected cond to be True", naming neither - # operand, and Dynamo rejects a message-carrying one. Callers dispatch on - # the width first -- see `compute_logprobs_chunk`, the packed path and - # `_pg_grad_forward`. - flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) - flat_index = index.reshape(-1) - - # Each chunk materialises rows x vocab logits and then a float32 copy of - # them, all on the device holding the output head. With a large vocabulary - # and a fixed chunk count that grows with the batch, so the peak scales with - # the batch rather than staying bounded. max_rows_per_chunk caps the rows - # per chunk instead, which is pure loop splitting: more, smaller chunks, - # same concatenated result. 0 (the default) keeps the previous chunk - # boundaries exactly. - if max_rows_per_chunk > 0: - n_rows = flat_hidden_states.shape[0] - chunks = max(chunks, -(-n_rows // max_rows_per_chunk)) - chunks = min(chunks, max(n_rows, 1)) - - # A one-row chunk's head gradient is a K=1 matmul Inductor mis-lowers under dynamic shapes (illegal - # memory access or NaN head gradients, torch 2.13), so keep chunks at 2+ rows when the head trains. - # Exactly `chunks` row ranges, so the unrolled loop length never depends on T: torch.chunk - # returns ceil(T / ceil(T / chunks)) pieces and recompiled per new count; tensor_split guards - # on T % chunks. At most rows // 2 chunks when the head trains keeps every piece at 2+ rows. - lm_head_grad = torch.is_grad_enabled() and lm_head.requires_grad - if lm_head_grad and flat_hidden_states.shape[0] < 2 * chunks: - chunks = max(flat_hidden_states.shape[0] // 2, 1) - - n_rows = flat_hidden_states.shape[0] - chunks = max(1, min(chunks, n_rows)) - bounds = [(i * n_rows) // chunks for i in range(chunks + 1)] - chunked_hidden_states = [flat_hidden_states[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - - # Cast once, outside the loop, so every chunk shares one fp32 head copy (also for autograd). - if fp32_lm_head: - lm_head = lm_head.to(torch.float32) - if lm_head_bias is not None: - lm_head_bias = lm_head_bias.to(torch.float32) - - all_per_token_logps = [] - - for chunk_hidden_states, chunk_index in zip(chunked_hidden_states, chunked_index): - # When the model is dispatched over several devices, the output head can - # sit on a different one from the hidden states, because accelerate - # places the tail of the model on the last device it fills. Co-locate on - # the head's device before the matmul, otherwise this raises - # Unhandled FakeTensor Device Propagation for aten.mm.default, - # found two different devices cuda:0, cuda:1 - # On a single device every .to() here is a no-op and the result is - # bit-identical to before. - chunk_hidden_states = chunk_hidden_states.to(device = lm_head.device, dtype = lm_head.dtype) - chunk_index = chunk_index.to(lm_head.device) - if fp32_lm_head: - # The GRPO loss runs this under bf16 / fp16 autocast, which would downcast the matmul. - with torch.autocast(device_type = chunk_hidden_states.device.type, enabled = False): - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - elif lm_head_bias is None: - chunk_logits = chunk_hidden_states @ lm_head.t() - else: - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - - if logit_scale_multiply != 0.0: - chunk_logits = chunk_logits * logit_scale_multiply - if logit_scale_divide != 0.0: - chunk_logits = chunk_logits / logit_scale_divide - if logit_softcapping != 0.0: - chunk_logits = logit_softcapping * torch.tanh(chunk_logits / logit_softcapping) - - chunk_logits = chunk_logits.to(torch.float32) - - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - per_token_logps = _truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - ) - else: - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - # Return to the caller's device so the concatenation below and every - # downstream consumer see the device they started on. - all_per_token_logps.append(per_token_logps.to(hidden_states.device)) - - all_per_token_logps = torch.concat(all_per_token_logps) - - all_per_token_logps = all_per_token_logps.reshape((hidden_states.shape[0], hidden_states.shape[1])) - return all_per_token_logps - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_selective_log_softmax( - logits, - index, - temperature: float = 1.0, - chunks: int = 4, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, -): - # Fixed `chunks` row ranges; see chunked_hidden_states_selective_log_softmax. - flat_logits = logits.reshape(-1, logits.shape[-1]) - flat_index = index.reshape(-1) - # Fewer rows than chunks: one row per range, as torch.chunk did, never empty ranges. - chunks = max(1, min(chunks, flat_logits.shape[0])) - bounds = [(i * flat_logits.shape[0]) // chunks for i in range(chunks + 1)] - chunked_logits = [flat_logits[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - all_per_token_logps = [] - # Per-chunk selective_log_softmax. - for chunk_logits, chunk_index in zip(chunked_logits, chunked_index): - chunk_logits = chunk_logits.to(torch.float32) - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - all_per_token_logps.append(_truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - )) - continue - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - all_per_token_logps.append(per_token_logps) - pass - all_per_token_logps = torch.concat(all_per_token_logps) - all_per_token_logps = all_per_token_logps.reshape((logits.shape[0], logits.shape[1])) - return all_per_token_logps - -def calculate_pad_tokens_in_prompt( - input_ids: torch.Tensor, - logits_to_keep: int, - pad_token_id: int -) -> torch.Tensor: - """Count left-padded tokens per sequence, e.g. [pad, pad, pad, cat] -> 3.""" - if logits_to_keep >= input_ids.shape[1]: - raise ValueError("logits_to_keep must be smaller than the sequence length.") - - prompt_section = input_ids[:, :-logits_to_keep] - - padding_mask = (prompt_section == pad_token_id) - - pad_token_counts = padding_mask.sum(dim=1) - - return pad_token_counts - -def create_completion_attention_mask( - completion_input_ids: torch.Tensor, - left_pad_tokens_per_prompt: torch.Tensor, - max_left_pad: int, - pad_token_id: int -) -> torch.Tensor: - """Build a completion mask that zeros leading prompt and trailing pad tokens. - - For [p,p,p,c,c,c,pad,pad,pad] (p=sliced prompt, c=completion, pad=padding) - this returns [0,0,0,1,1,1,0,0,0]. - """ - batch_size, completion_len = completion_input_ids.shape - device = completion_input_ids.device - - num_tokens_to_mask = max_left_pad - left_pad_tokens_per_prompt - - indices = torch.arange(completion_len, device=device).unsqueeze(0) - shift_mask = indices >= num_tokens_to_mask.unsqueeze(1) - - non_padding_mask = (completion_input_ids != pad_token_id) - - final_mask = shift_mask & non_padding_mask - - return final_mask - -def left_pack_padding(tensor: torch.Tensor, pad_id: int) -> torch.Tensor: - """Move all padding tokens in each sequence to the right.""" - mask = (tensor != pad_id) - # stable=True since the binary mask is unordered. - sorted_indices = torch.argsort(mask, dim=1, descending=True, stable=True) - packed_tensor = torch.gather(tensor, 1, sorted_indices) - return packed_tensor - -def align_logprobs_with_mask( - logprob_tensor: torch.Tensor, - attention_mask: torch.Tensor, - pad_value: float = 0.0 -) -> torch.Tensor: - """Align a log probability tensor with a given attention mask.""" - - device = logprob_tensor.device - batch_size, logprob_seq_len = logprob_tensor.shape - mask_seq_len = attention_mask.shape[1] - - padded_logprobs = torch.full( - attention_mask.shape, - fill_value=pad_value, - dtype=logprob_tensor.dtype, - device=device - ) - - left_pad_counts = torch.argmax(attention_mask, dim=1) - - cols = torch.arange(logprob_seq_len, device=device) - dest_indices = left_pad_counts.unsqueeze(1) + cols - - # Destination row indices, shape [batch_size, logprob_seq_len]. - row_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand_as(dest_indices) - - # Keep only in-bounds destinations, then scatter via advanced indexing. - valid_mask = dest_indices < mask_seq_len - valid_rows = row_indices[valid_mask] - valid_cols = dest_indices[valid_mask] - valid_vals = logprob_tensor[valid_mask] - padded_logprobs[valid_rows, valid_cols] = valid_vals - - return padded_logprobs - -def align_completion_tool_mask( - tool_mask: torch.Tensor, - completion_mask: torch.Tensor, -) -> torch.Tensor: - """Align a raw completion-length tool/env mask with Unsloth's repacked loss mask.""" - if tool_mask is None: - return completion_mask - if tool_mask.shape[0] != completion_mask.shape[0]: - raise ValueError("tool_mask batch size must match completion_mask batch size.") - - tool_mask = tool_mask.to(device=completion_mask.device) - if tool_mask.shape == completion_mask.shape: - aligned_tool_mask = tool_mask - else: - aligned_tool_mask = align_logprobs_with_mask( - tool_mask, - completion_mask, - pad_value=0, - ) - return completion_mask * aligned_tool_mask.to(dtype=completion_mask.dtype) - -def autotune_batch_and_chunks( - total_input_rows, - seq_len, - hidden_size, - vocab_size, - dtype_bytes=16, - multiplier=None -): - if multiplier is None: - final_m = max(4, seq_len // 4096) - else: - final_m = multiplier - - if torch.cuda.is_available(): - free_bytes, _ = torch.cuda.mem_get_info() - limit_gb = (free_bytes / (1024**3))*.80 - elif hasattr(torch, "xpu") and torch.xpu.is_available(): - # XPU: estimate free memory as total - reserved. - total_mem = torch.xpu.get_device_properties(0).total_memory - reserved_mem = torch.xpu.memory_reserved() - free_bytes = total_mem - reserved_mem - limit_gb = (free_bytes / (1024**3)) * 0.80 - else: - # Fallback: assume 8GB available. - limit_gb = 8.0 - - bytes_to_gb = 1024**3 - - b_vals = torch.arange(total_input_rows, 0, -1, device='cpu', dtype=torch.float32) - - hidden_gb = (b_vals * seq_len * hidden_size * dtype_bytes) / bytes_to_gb - - base_logits = ((b_vals/total_input_rows) * b_vals * seq_len * vocab_size * dtype_bytes) / bytes_to_gb - logits_gb = base_logits / final_m - - total_mem_gb = hidden_gb + logits_gb - - valid_mask = total_mem_gb <= limit_gb - valid_indices = torch.nonzero(valid_mask, as_tuple=False) - - if valid_indices.shape[0] == 0: - #This means your GPU will OOM - # Capped at the row count: unsloth's no-grad pass divides rows by this without max(1, ...). - return max(1, min(4, total_input_rows)), final_m - - best_idx = valid_indices[0].item() - final_b = int(b_vals[best_idx].item()) - - return final_b, final_m - -def sanitize_logprob(logprob): - """Local port of trl.scripts.vllm_serve.sanitize_logprob. - Filters NaN logprobs from vLLM outputs.""" - value = logprob.logprob - if math.isnan(value): - logging.getLogger(__name__).warning( - f"Generated NaN logprob, token logprob '{logprob}' will be ignored" - ) - return None - return value -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_text_configs(config): - """``config`` and its text sub-config, the two places a transform can live: Gemma-4-style configs keep it on ``config.text_config``, T5Gemma-style ones only reach it via ``config.get_text_config()``.""" - if config is None: - return [] - holders = [config] - 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: - holders.append(text_cfg) - return holders - -def _unsloth_get_final_logit_softcapping(model): - """The soft cap the loss applies, under any of its three spellings: ``final_logit_softcapping`` (Gemma), ``logits_soft_cap`` (RecurrentGemma), ``output_logit_soft_cap`` (xLSTM). Returns 0 if unset.""" - config = _unsloth_get_model_config(model) - if config is None: - return 0 - for holder in _unsloth_text_configs(config): - for name in ("final_logit_softcapping", "logits_soft_cap", "output_logit_soft_cap"): - softcap = getattr(holder, name, None) - if softcap: - return softcap - return 0 - -def _unsloth_resolve_logit_scales(model_config): - """``(multiply, divide)`` for the logits, read with the same field table and per-family overrides ``detect_logit_transforms`` uses. - - Only reached when the installed unsloth_zoo predates that helper. Must stay in step with ``resolve_logit_transforms`` in unsloth/models/llama.py: if the forward applies a transform GRPO does not, the policy log-probabilities come from different logits than the ones generated, which silently shifts every importance ratio. - """ - # Granite divides, HyperCLOVA X multiplies, MiniCPM3 scales hidden states (not a logit transform). - overrides = { - ("logits_scaling", "hyperclovax"): "multiply", - ("logits_scaling", "minicpm3"): None, - } - found = {"multiply": 0, "divide": 0} - for holder in _unsloth_text_configs(model_config): - model_type = getattr(holder, "model_type", "") or "" - for bucket, names in ( - ("multiply", ("logit_scale", "lm_head_multiplier", "output_multiplier")), - ("divide", ("logits_scaling",)), - ): - for name in names: - target = overrides.get((name, model_type), bucket) - if target is None or found[target]: - continue - value = getattr(holder, name, 0) or 0 - if value: - found[target] = value - break - return found["multiply"], found["divide"] - -def _unsloth_grpo_returns_hidden_states(model, tensor, lm_head): - """Does ``tensor`` (a forward's ``.logits``) carry hidden states or real logits? - - ``_get_per_token_logps_and_entropies`` sets ``UNSLOTH_RETURN_HIDDEN_STATES=1``, but only a forward that honours the name hands hidden states back as ``.logits``; any other forward returns a real ``[.., vocab]`` tensor that must not reach the ``lm_head`` matmul. - - The primary test is an explicit signal that the forward honours the flag, both set outside this file: ``__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__`` on the generated class, written by ``unsloth_zoo.compiler.create_standalone_class`` exactly when ``apply_fused_lm_head`` gave that forward its own branch; and ``_unsloth_grpo_hidden_states_forward_wrapped``, set by ``_install_grpo_hidden_states_forward_wrapper`` in ``unsloth/models/rl.py`` for models the compiler did not rewrite. That wrapper degrades to real logits when the model cannot produce hidden states and records whether it did so in ``_unsloth_grpo_hidden_states_degraded`` before returning, so reading the pair after a forward describes the call that just finished; degradation is per call, not per model. - - The width comparison stays as the fallback, for an ``unsloth_zoo`` old enough that it never writes the marker. It is decisive whenever ``vocab_size != hidden_size``, and the signal may only overrule it when it is not, which is the one case the shape cannot answer. - """ - if tensor.shape[-1] != lm_head.shape[1]: - return False # vocab-wide: real logits, whatever any signal claims - if lm_head.shape[0] != lm_head.shape[1]: - return True # hidden-wide and vocab_size != hidden_size: hidden states - return _unsloth_grpo_hidden_states_signal(model) is not False - -def _unsloth_grpo_hidden_states_signal(model): - """``True``/``False`` if the forward honours ``UNSLOTH_RETURN_HIDDEN_STATES``, ``None`` when neither marker is present. See ``_unsloth_grpo_returns_hidden_states`` for where each marker is set. Walks the wrapper chain because the markers are set on whichever object the trainer saw, which may be the DDP module or the PEFT base model rather than the object handed to the logprob loop.""" - candidates = [] - pending = [model] - while pending and len(candidates) < 8: - candidate = pending.pop(0) - if candidate is None or any(candidate is seen for seen in candidates): - continue - candidates.append(candidate) - get_base_model = getattr(candidate, "get_base_model", None) - if callable(get_base_model): - try: - pending.append(get_base_model()) - except Exception: - pass - for _attr in ("module", "base_model", "model"): - child = getattr(candidate, _attr, None) - if child is not None and hasattr(child, "forward"): - pending.append(child) - for candidate in candidates: - if getattr(candidate, "__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__", False): - return True - if getattr(type(candidate), "__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__", False): - return True - if any( - getattr(candidate, "_unsloth_grpo_hidden_states_forward_wrapped", False) - for candidate in candidates - ): - # The wrapper honours the flag unless it recorded that this call could not. - if any( - hasattr(candidate, "_unsloth_grpo_hidden_states_degraded") for candidate in candidates - ): - return not any( - getattr(candidate, "_unsloth_grpo_hidden_states_degraded", False) - for candidate in candidates - ) - # An unsloth/models/rl.py predating the per-call attribute set only the warn-once flag; it is the best signal such a wrapper offers. - return not any( - getattr(candidate, "_unsloth_grpo_hidden_states_warning_issued", False) - for candidate in candidates - ) - return None - -def _unsloth_pref_returns_hidden_states(model, tensor, lm_head): - """``_unsloth_grpo_returns_hidden_states`` for the preference routes, which need a zoo that writes the hidden-state - markers: a square head with no signal is a forward that ignored the flag, so its output is real logits.""" - if tensor.shape[-1] != lm_head.shape[1]: - return False - if lm_head.shape[0] != lm_head.shape[1]: - return True - return _unsloth_grpo_hidden_states_signal(model) is True - -def _unsloth_gkd_dense_head(model): - """The output head when it is a plain dense ``[vocab, hidden]`` projection the chunked loss can read directly. - Rejects bnb heads (``nn.Linear`` subclasses with packed weights), PEFT heads and ZeRO-3 partitioned weights. - """ - get_output_embeddings = getattr(model, "get_output_embeddings", None) - if not callable(get_output_embeddings): - return None - try: - head = get_output_embeddings() - except Exception: - return None - if type(head) is not torch.nn.Linear: - return None - # An lm_head that is not the output embeddings runs before them (ModernBERT decoder, RoBERTa-style heads). - lm_head = getattr(model, "lm_head", None) - if isinstance(lm_head, torch.nn.Module) and lm_head is not head: - return None - weight = getattr(head, "weight", None) - if not isinstance(weight, torch.Tensor) or weight.dim() != 2 or weight.numel() == 0: - return None - if weight.dtype not in (torch.float16, torch.bfloat16, torch.float32, torch.float64): - return None - if type(weight) not in (torch.Tensor, torch.nn.Parameter): - return None - bias = getattr(head, "bias", None) - if bias is not None and (bias.dim() != 1 or bias.shape[0] != weight.shape[0]): - return None - if weight.is_meta or (bias is not None and bias.is_meta): - return None - return head - -def _unsloth_gkd_logit_transforms(model): - """``(scale, softcap)`` exactly as the model's own forward applies them to its logits.""" - model_config = _unsloth_get_model_config(model) - if detect_logit_transforms is not None: - _transforms = detect_logit_transforms(model_config) - logit_softcapping = _transforms["logit_softcapping"] - logit_scale_multiply = _transforms["logit_scale_multiply"] - logit_scale_divide = _transforms["logit_scale_divide"] - else: - logit_softcapping = _unsloth_get_final_logit_softcapping(model) - logit_scale_multiply, logit_scale_divide = _unsloth_resolve_logit_scales(model_config) - scale = 1.0 - if logit_scale_multiply: - scale = scale * float(logit_scale_multiply) - if logit_scale_divide: - scale = scale / float(logit_scale_divide) - return scale, float(logit_softcapping or 0.0) - -def _unsloth_pref_logit_transforms(model): - """``(multiply, divide, softcap)`` for the chunked kernel, from the GKD resolver (one scale, 0.0 = off).""" - scale, softcap = _unsloth_gkd_logit_transforms(model) - return (0.0 if scale == 1.0 else scale), 0.0, softcap - -def _unsloth_stock_logps_hook( - self, - method, - hooks = ("get_batch_logps",), -): - """False when a subclass overrides one of ``hooks`` below the class whose ``method`` is the rewritten one: TRL's - body would call the override (or it reads the full logits the chunked path replaces with stand-ins).""" - owner = None - for klass in type(self).__mro__: - # The rewritten body calls a chunked or fused-pass helper; a subclass method wrapping it via super() does not. - names = getattr(getattr(vars(klass).get(method), "__code__", None), "co_names", ()) - if any( - name.startswith("_unsloth_") and ("chunked" in name or "fused_pass" in name) - for name in names - ): - owner = klass - break - if owner is None: - return True - for name in hooks: - hook = next((klass for klass in type(self).__mro__ if name in vars(klass)), None) - if hook is not None and hook not in owner.__mro__: - return False - return True - -def _unsloth_dpo_chunked_head( - self, - model, - concatenated_batch, - method = "concatenated_forward", -): - """``(unwrapped_model, head, (multiply, divide, softcap))`` when concatenated_forward can score from hidden - states, else ``None`` (TRL's own logits path runs). Decided before the forward, which then skips the logits.""" - reason = None - if _unsloth_masked_logprobs is None: - reason = "unsloth_zoo has no chunked_logprob" - elif os.environ.get("UNSLOTH_FUSED_LINEAR", "1") == "0": - reason = "UNSLOTH_FUSED_LINEAR=0" - elif os.environ.get("UNSLOTH_DPO_CHUNKED", "1").lower() in ("0", "false", "no", "off"): - reason = "UNSLOTH_DPO_CHUNKED=0" - elif getattr(self, "is_encoder_decoder", False): - reason = "encoder-decoder" - elif getattr(self, "aux_loss_enabled", False): - reason = "router aux loss needs the stock forward" - elif not _unsloth_stock_logps_hook( - self, - method, - ("get_batch_logps",) - if method == "concatenated_forward" - else ("get_batch_logps", "get_batch_loss_metrics", method), - ): - # KTO hands its callers logits stand-ins; DPO's outputs are exact, so a DPO wrapper is fine. - reason = "get_batch_logps / get_batch_loss_metrics overridden" - elif getattr(self, "padding_free", False): - reason = "padding_free" - elif getattr(self, "use_weighting", False): - reason = "use_weighting needs full-vocabulary log-probs" - elif getattr(self, "is_fsdp_enabled", False) or getattr(self, "is_deepspeed_enabled", False): - reason = "FSDP / DeepSpeed" - elif any( - k in concatenated_batch - for k in ("pixel_values", "pixel_attention_mask", "image_sizes", "token_type_ids") - + tuple(_DPO_VISION_KEYS) - ): - reason = "vision inputs" - if reason is None: - try: - unwrapped = self.accelerator.unwrap_model(model) - except Exception: - unwrapped = None - reason = "unwrap_model" - if reason is None: - head = _unsloth_gkd_dense_head(unwrapped) - if head is None: - reason = "output head is not a dense nn.Linear" - # The kernel reads weight / bias directly: a hooked head (accelerate, forward / backward hooks) is not - # F.linear(h, W, b), and its backward hooks would never see the gradient. - elif ( - hasattr(head, "_hf_hook") - or "forward" in vars(head) - or head._forward_hooks - or head._forward_pre_hooks - or head._backward_hooks - or getattr(head, "_backward_pre_hooks", None) - ): - reason = "output head has hooks or a patched forward" - if reason is None: - for config in _unsloth_text_configs(_unsloth_get_model_config(unwrapped)): - if getattr(config, "model_type", None) == "minicpm3": - reason = "minicpm3 scales hidden states before the head" - if reason is None and getattr(model, "find_unused_parameters", False): - if any(p is not None and p.requires_grad for p in (head.weight, head.bias)): - reason = "DDP find_unused_parameters with a trainable head" - if reason is not None: - try: - fallbacks = self.__dict__.setdefault("_unsloth_dpo_chunked_fallbacks", {}) - fallbacks[reason] = fallbacks.get(reason, 0) + 1 - except Exception: - pass - return None - return unwrapped, head, _unsloth_pref_logit_transforms(unwrapped) - -def _unsloth_dpo_model_forward(model, input_ids, model_kwargs, want_hidden_states): - if not want_hidden_states: - return model(input_ids, **model_kwargs) - prior = os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES") - os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" - try: - return model(input_ids, **model_kwargs) - finally: - if prior is None: - os.environ.pop("UNSLOTH_RETURN_HIDDEN_STATES", None) - else: - os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = prior - -def _unsloth_dpo_returns_hidden_states(self, chunked, logits): - if chunked is None: - return False - unwrapped, head, _ = chunked - hidden = _unsloth_pref_returns_hidden_states(unwrapped, logits, head.weight) - if not hidden: - fallbacks = self.__dict__.setdefault("_unsloth_dpo_chunked_fallbacks", {}) - reason = "forward returned logits, not hidden states" - fallbacks[reason] = fallbacks.get(reason, 0) + 1 - return hidden - -def _unsloth_kto_hidden_logps(self, chunked, hidden_states, labels, label_pad_token_id): - """``(completion_logps [B], logits stand-in [B, S, 1])`` for TRL's ``get_batch_logps(logits, labels, - average_log_prob=False)``. The stand-in holds each position's vocab sum of the logits, so the - ``logits/chosen_sum`` metrics (a sum over every position) stay exact.""" - _, head, (multiply, divide, softcap) = chunked - # The model may truncate input_ids (see kto_trainer_get_batch_logps): score the common length. - n = min(hidden_states.shape[1], labels.shape[1]) - out_device = hidden_states.device - hidden_states = hidden_states[:, :n] - labels = labels[:, :n].to(head.weight.device) - shifted = labels[:, 1:] - mask = shifted != label_pad_token_id - if _unsloth_masked_logprobs is not None and _unsloth_linear_mean_logits is not None: - # Only completion tokens reach the head; the logits/* sums come from the head's affine mean (no backward). - hidden_states = hidden_states.to(head.weight.device) - per_token_logps = _unsloth_masked_logprobs( - hidden_states[:, :-1], - head.weight, - shifted, - mask = mask, - logit_scale_multiply = multiply, - logit_scale_divide = divide, - logit_softcapping = softcap, - lm_head_bias = head.bias, - ) - mean_logits = _unsloth_linear_mean_logits( - hidden_states, head.weight, head.bias, multiply, divide, softcap - ) - else: - index = torch.cat( - (shifted.masked_fill(~mask, 0), shifted.new_zeros(shifted.shape[0], 1)), dim = 1 - ) - per_token_logps, mean_logits = _unsloth_chunked_logprobs( - hidden_states.to(head.weight.device), - head.weight, - index, - logit_scale_multiply = multiply, - logit_scale_divide = divide, - logit_softcapping = softcap, - lm_head_bias = head.bias, - return_mean_logits = True, - ) - per_token_logps = per_token_logps[:, :-1] - try: - self._unsloth_dpo_chunked_calls = getattr(self, "_unsloth_dpo_chunked_calls", 0) + 1 - except Exception: - pass - completion_logps = (per_token_logps * mask).sum(-1) - logits_sum = (mean_logits * head.weight.shape[0]).unsqueeze(-1) - return completion_logps.to(out_device), logits_sum.to(out_device) - -try: - from unsloth_zoo.device_map_planner import detect_logit_transforms -except Exception: - detect_logit_transforms = None -try: - from unsloth_zoo.fused_losses.chunked_logprob import ( - unsloth_chunked_logprobs as _unsloth_chunked_logprobs, - unsloth_masked_logprobs as _unsloth_masked_logprobs, - ) -except Exception: - _unsloth_chunked_logprobs = _unsloth_masked_logprobs = None -try: - from unsloth_zoo.fused_losses.preference_loss import linear_mean_logits as _unsloth_linear_mean_logits -except Exception: - _unsloth_linear_mean_logits = None -_DPO_VISION_KEYS = ('pixel_position_ids', 'image_position_ids', 'mm_token_type_ids') - -def _unsloth_kto_ref_chunked_forward(model, want_hidden_states, *args, **kwargs): - """``model(*args, **kwargs)``, returning hidden states as ``.logits`` when ``want_hidden_states``.""" - if not want_hidden_states: - return model(*args, **kwargs) - prior = os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES") - os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" - try: - return model(*args, **kwargs) - finally: - if prior is None: - os.environ.pop("UNSLOTH_RETURN_HIDDEN_STATES", None) - else: - os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = prior - -def _unsloth_kto_ref_batch_logps(self, chunked, hidden_states, labels, label_pad_token_id): - """TRL < 1.6's ``get_batch_logps(logits, labels, average_log_prob=False)`` (decoder-only) from hidden states.""" - shifted = labels[:, 1:] - mask = shifted != label_pad_token_id - return _unsloth_fused_pass_logps( - self, chunked, hidden_states[:, :-1], shifted.masked_fill(~mask, 0), mask - ).sum(-1) - -def _unsloth_dpo_forward(model, model_kwargs, want_hidden_states): - if not want_hidden_states: - return model(**model_kwargs) - prior = os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES") - os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" - try: - return model(**model_kwargs) - finally: - if prior is None: - os.environ.pop("UNSLOTH_RETURN_HIDDEN_STATES", None) - else: - os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = prior - -def _unsloth_fused_pass_head(self, model, inputs, return_outputs, method): - """``_unsloth_dpo_chunked_head`` for a rewritten pass that also needs the fused linear op, keeps TRL's path for - the kill switch, ``return_outputs`` (the caller reads ``outputs.logits``) and TRL's own liger / chunked head.""" - reason = None - if _unsloth_fused_linear_logprobs is None: - reason = "unsloth_zoo has no fused_linear_logprobs" - elif os.environ.get("UNSLOTH_FUSED_LINEAR", "1") == "0": - reason = "UNSLOTH_FUSED_LINEAR=0" - elif return_outputs: - reason = "return_outputs needs the logits" - elif getattr(self, "use_liger_kernel", False): - reason = "use_liger_kernel runs TRL's own chunked head" - if reason is not None: - try: - fallbacks = self.__dict__.setdefault("_unsloth_dpo_chunked_fallbacks", {}) - fallbacks[reason] = fallbacks.get(reason, 0) + 1 - except Exception: - pass - return None - return _unsloth_dpo_chunked_head(self, model, inputs, method) - -def _unsloth_fused_pass_logps( - self, - chunked, - hidden_states, - labels, - mask, - stats = False, - grad = False, -): - """TRL's ``selective_log_softmax(head(hidden_states), labels)`` with rows outside ``mask`` zeroed (TRL's - ``row_mask`` or its ``logps[mask == 0] = 0.0``). ``stats`` also returns the entropy (no grad, as TRL computes it - from detached logits) and the per-position vocab mean ``[.., 1]``: ``x[mask].mean(-1)`` on it is TRL's mean of - the logits row. A shorter ``hidden_states`` (the model truncated input_ids) is scored over the common length and - zero padded. Only the policy pass (``grad``) keeps a graph: TRL scores the reference / KL logits it built under - no_grad, so a trainable head must get no gradient from them.""" - _, head, (multiply, divide, softcap) = chunked - width = labels.shape[1] - n = min(hidden_states.shape[1], width) - with torch.set_grad_enabled(grad and torch.is_grad_enabled()): - out = _unsloth_fused_linear_logprobs( - hidden_states[:, :n], - head.weight, - labels[:, :n], - bias = head.bias, - mask = mask[:, :n].bool(), - logit_scale_multiply = multiply, - logit_scale_divide = divide, - logit_softcapping = softcap, - return_entropy = stats, - return_mean_logits = stats, - entropy_requires_grad = False, - ) - try: - self._unsloth_dpo_chunked_calls = getattr(self, "_unsloth_dpo_chunked_calls", 0) + 1 - except Exception: - pass - out = out if isinstance(out, tuple) else (out,) - out = tuple(torch.nn.functional.pad(x, (0, width - n)) if n != width else x for x in out) - if not stats: - return out[0].to(labels.device) - logps, entropy, mean_logits = out - return ( - logps.to(labels.device), - entropy.to(labels.device), - mean_logits.unsqueeze(-1).to(labels.device), - ) - -def _unsloth_fused_pass_vocab_sums(self, chunked, hidden_states): - """``[B, S, 1]`` per-position vocab sums of ``head(hidden_states)`` (no grad): TRL 1.6 - 1.8 KTO logs - ``logits.index_select(0, idx).nansum()``, a sum over every position and vocab entry, which this keeps exact.""" - _, head, (multiply, divide, softcap) = chunked - with torch.no_grad(): - hidden_states = hidden_states.detach() - _, mean_logits = _unsloth_fused_linear_logprobs( - hidden_states, - head.weight, - torch.zeros(hidden_states.shape[:-1], dtype = torch.long, device = hidden_states.device), - bias = head.bias, - logit_scale_multiply = multiply, - logit_scale_divide = divide, - logit_softcapping = softcap, - return_mean_logits = True, - ) - return (mean_logits * head.weight.shape[0]).unsqueeze(-1).to(hidden_states.device) - -try: - from unsloth_zoo.fused_losses.fused_linear import ( - fused_linear_logprobs as _unsloth_fused_linear_logprobs, - ) -except Exception: - _unsloth_fused_linear_logprobs = None -@dataclass -class UnslothKTOConfig(KTOConfig): - """ - - Configuration class for the [`KTOTrainer`]. - - This class includes only the parameters that are specific to KTO training. For a full list of training arguments, - please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may - differ from those in [`~transformers.TrainingArguments`]. - - Using [`~transformers.HfArgumentParser`] we can turn this class into - [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the - command line. - - Parameters: - > Parameters that control the model - - model_init_kwargs (`dict[str, Any]`, *optional*): - Keyword arguments for [`~transformers.AutoModelForCausalLM.from_pretrained`], used when the `model` - argument of the [`KTOTrainer`] is provided as a string. The `revision` value is also used when loading the - processing class. - trust_remote_code (`bool`, *optional*, defaults to `False`): - Whether to allow loading models and tokenizers that ship custom Python code from the Hub. Forwarded to - [`~transformers.AutoModelForCausalLM.from_pretrained`] and [`~transformers.AutoProcessor.from_pretrained`]. - router_aux_loss_coef (`float`, *optional*): - Coefficient of the load-balancing auxiliary loss for Mixture-of-Experts (MoE) models, added to the training - loss with this weight. When not set, the value declared by the model config is used. Set to `0.0` to - disable it. Fails when used with a non-MoE model. - disable_dropout (`bool`, *optional*, defaults to `True`): - Whether to disable dropout in the model and reference model. - - > Parameters that control the data preprocessing - - dataset_num_proc (`int`, *optional*): - Number of processes to use for processing the dataset. - max_length (`int` or `None`, *optional*, defaults to `1024`): - Maximum length of the tokenized sequence. Sequences longer than `max_length` are truncated from the right. - If `None`, no truncation is applied. - pad_to_multiple_of (`int`, *optional*): - If set, the sequences will be padded to a multiple of this value. - precompute_ref_log_probs (`bool`, *optional*, defaults to `False`): - Whether to precompute the reference model log probabilities for the entire training dataset before - training. This allows to save memory during training, as the reference model does not need to be kept in - memory. - precompute_ref_batch_size (`int`, *optional*): - Batch size to use when precomputing reference model log probabilities. This can be set higher than the - training batch size to speed up preprocessing. If `None`, defaults to `per_device_train_batch_size` for - training and `per_device_eval_batch_size` for evaluation. - - > Parameters that control the training - - loss_type (`str`, *optional*, defaults to `"kto"`): - Type of loss to use. Possible values are: - - - `"kto"`: KTO loss from the [KTO](https://huggingface.co/papers/2402.01306) paper. - - `"apo_zero_unpaired"`: Unpaired variant of APO-zero loss from the - [APO](https://huggingface.co/papers/2408.06266) paper. - - beta (`float`, *optional*, defaults to `0.1`): - Parameter controlling the deviation from the reference model. Higher β means less deviation from the - reference model. - desirable_weight (`float`, *optional*, defaults to `1.0`): - Desirable losses are weighed by this factor to counter unequal number of desirable and undesirable pairs. - undesirable_weight (`float`, *optional*, defaults to `1.0`): - Undesirable losses are weighed by this factor to counter unequal number of desirable and undesirable pairs. - activation_offloading (`bool`, *optional*, defaults to `False`): - Whether to offload the activations to the CPU. - sync_ref_model (`bool`, *optional*, defaults to `False`): - Whether to synchronize the reference model with the active model every `ref_model_sync_steps` steps, using - the `ref_model_mixup_alpha` parameter. This synchronization originates from the - [TR-DPO](https://huggingface.co/papers/2404.09656) paper. `sync_ref_model=True` is not yet compatible with - PEFT or `precompute_ref_log_probs=True`. - ref_model_mixup_alpha (`float`, *optional*, defaults to `0.6`): - α parameter from the TR-DPO paper, which controls the mix between the current policy and the previous - reference policy during updates. The reference policy is updated according to the equation: `π_ref = α * - π_θ + (1 - α) * π_ref_prev`. To use this parameter, you must set `sync_ref_model=True`. - ref_model_sync_steps (`int`, *optional*, defaults to `512`): - τ parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which determines how - frequently the current policy is synchronized with the reference policy. To use this parameter, you must - set `sync_ref_model=True`. - - > [!NOTE] - > These parameters have default values different from [`~transformers.TrainingArguments`]: - > - `logging_steps`: Defaults to `10` instead of `500`. - > - `gradient_checkpointing`: Defaults to `True` instead of `False`. - > - `bf16`: Defaults to `True` if `fp16` is not set, instead of `False`. - > - `learning_rate`: Defaults to `1e-6` instead of `5e-5`. - > - `train_sampling_strategy`: Defaults to `"sequential"` instead of `"random"`. Loss types - > that estimate the KL divergence term (all except `"apo_zero_unpaired"`) require sequential - > sampling because the KL completion for each example is precomputed against its neighbors in - > a fixed-order batch; any other strategy breaks that pairing. - - """ - _unsloth_accepts_adamw_flash = True - vllm_sampling_params: Optional[Any] = field( - default = None, - metadata = {'help': 'vLLM SamplingParams'}, - ) - unsloth_num_chunks : Optional[int] = field( - default = -1, - metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'}, - ) - unsloth_logit_chunk_multiplier : Optional[int] = field( - default = None, - metadata = {'help': 'Multiplier for chunked logit computations.'}, - ) - unsloth_grpo_mini_batch : Optional[int] = field( - default = None, - metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'}, - ) - max_seq_length : Optional[int] = field( - default = None, - metadata = {'help': 'Maximum sequence length to truncate to.'}, - ) - def __init__( - self, - output_dir = None, - per_device_train_batch_size = 4, - num_train_epochs = 3.0, - max_steps = -1, - learning_rate = 5e-05, - lr_scheduler_type = 'linear', - lr_scheduler_kwargs = None, - warmup_steps = 0.1, - optim = 'adamw_8bit', - optim_args = None, - weight_decay = 0.001, - adam_beta1 = 0.9, - adam_beta2 = 0.999, - adam_epsilon = 1e-08, - optim_target_modules = None, - gradient_accumulation_steps = 2, - average_tokens_across_devices = True, - max_grad_norm = 1.0, - label_smoothing_factor = 0.0, - bf16 = False, - fp16 = False, - bf16_full_eval = False, - fp16_full_eval = False, - tf32 = None, - gradient_checkpointing = True, - gradient_checkpointing_kwargs = None, - torch_compile = False, - torch_compile_backend = None, - torch_compile_mode = None, - use_liger_kernel = False, - liger_kernel_config = None, - use_cache = False, - neftune_noise_alpha = None, - torch_empty_cache_steps = 250, - auto_find_batch_size = False, - logging_strategy = 'steps', - logging_steps = 1, - logging_first_step = False, - log_on_each_node = True, - logging_nan_inf_filter = False, - include_num_input_tokens_seen = False, - log_level = 'passive', - log_level_replica = 'warning', - disable_tqdm = None, - report_to = 'none', - run_name = None, - project = 'huggingface', - trackio_space_id = 'trackio', - eval_strategy = 'no', - eval_steps = None, - eval_delay = 0, - per_device_eval_batch_size = 4, - prediction_loss_only = False, - eval_on_start = False, - eval_do_concat_batches = True, - eval_use_gather_object = False, - eval_accumulation_steps = 2, - batch_eval_metrics = False, - save_only_model = False, - save_strategy = 'steps', - save_steps = 500, - save_on_each_node = False, - save_total_limit = None, - enable_jit_checkpoint = False, - push_to_hub = False, - hub_token = None, - hub_private_repo = None, - hub_model_id = None, - hub_strategy = 'every_save', - hub_always_push = False, - hub_revision = None, - load_best_model_at_end = False, - metric_for_best_model = None, - greater_is_better = None, - ignore_data_skip = False, - restore_callback_states_from_checkpoint = False, - full_determinism = False, - seed = 3407, - data_seed = None, - use_cpu = False, - accelerator_config = None, - parallelism_config = None, - dataloader_drop_last = False, - dataloader_num_workers = 0, - dataloader_pin_memory = True, - dataloader_persistent_workers = False, - dataloader_prefetch_factor = None, - remove_unused_columns = True, - label_names = None, - train_sampling_strategy = 'sequential', - length_column_name = 'length', - ddp_find_unused_parameters = None, - ddp_bucket_cap_mb = None, - ddp_broadcast_buffers = None, - ddp_backend = None, - ddp_timeout = 1800, - fsdp = None, - fsdp_config = None, - deepspeed = None, - debug = '', - skip_memory_metrics = True, - do_train = False, - do_eval = False, - do_predict = False, - resume_from_checkpoint = None, - warmup_ratio = None, - logging_dir = None, - local_rank = -1, - model_init_kwargs = None, - trust_remote_code = False, - router_aux_loss_coef = None, - disable_dropout = True, - dataset_num_proc = None, - max_length = 1024, - pad_to_multiple_of = None, - precompute_ref_log_probs = False, - precompute_ref_batch_size = None, - loss_type = 'kto', - beta = 0.1, - desirable_weight = 1.0, - undesirable_weight = 1.0, - activation_offloading = False, - sync_ref_model = False, - ref_model_mixup_alpha = 0.6, - ref_model_sync_steps = 512, - vllm_sampling_params = None, - unsloth_num_chunks = -1, - unsloth_logit_chunk_multiplier = None, - unsloth_grpo_mini_batch = None, - max_seq_length = None, - **kwargs, - ): - if learning_rate < 1e-7: print(f'Unsloth: Your learning rate of `{learning_rate}` is too small and less than 1e-7! Consider increasing it, otherwise gradient updates will be close to 0!') - if learning_rate > 1: print(f'Unsloth: Your learning rate of `{learning_rate}` is way too larger > 1! Consider decreasing it to 1e-1, otherwise gradient updates will explode!') - _unsloth_fsdp2 = os.environ.get('FSDP_VERSION', '').strip() == '2' - _unsloth_fsdp_cfg = fsdp_config - if fsdp and isinstance(_unsloth_fsdp_cfg, str) and os.path.isfile(_unsloth_fsdp_cfg): - try: - import json as _unsloth_json - with open(_unsloth_fsdp_cfg, encoding = 'utf-8') as _unsloth_f: - _unsloth_fsdp_cfg = _unsloth_json.load(_unsloth_f) - except Exception: - _unsloth_fsdp_cfg = None - if fsdp and _unsloth_fsdp_cfg is None: - _unsloth_fsdp_cfg = {} - if fsdp and isinstance(_unsloth_fsdp_cfg, dict): - _unsloth_fsdp2 = _unsloth_fsdp2 or str(_unsloth_fsdp_cfg.get('fsdp_version', _unsloth_fsdp_cfg.get('version', 1))) == '2' - _unsloth_flash_adamw = str(getattr(optim, 'value', optim)).strip().lower() == 'adamw_flash' - if _unsloth_flash_adamw: optim = 'adamw_torch' - if _unsloth_fsdp2 and str(getattr(optim, 'value', optim)) == 'adamw_8bit': - print("Unsloth: FSDP2 shards parameters as DTensors, which bitsandbytes' 8-bit optimizers cannot step. Switching optim = 'adamw_8bit' to 'adamw_torch_fused'.") - optim = 'adamw_torch_fused' - if num_train_epochs is None: - num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override - if output_dir is None and save_strategy == 'steps' and save_steps == 500: - output_dir = 'unsloth_training_checkpoints' - save_strategy = 'no' - try: - from unsloth_zoo.dataset_num_proc import get_dataset_num_proc as _unsloth_get_dataset_num_proc - except Exception: - try: - from unsloth.dataset_num_proc import get_dataset_num_proc as _unsloth_get_dataset_num_proc - except Exception: - _unsloth_get_dataset_num_proc = None - if _unsloth_get_dataset_num_proc is not None: - dataset_num_proc = _unsloth_get_dataset_num_proc(dataset_num_proc, serial_as_none = True) - if os.environ.get('UNSLOTH_ENABLE_FLEX_ATTENTION', '0') == '1': - from unsloth_zoo.flex_attention import HAS_FLEX_ATTENTION - if HAS_FLEX_ATTENTION and pad_to_multiple_of is None: - from unsloth_zoo.flex_attention import FLEX_ATTENTION_BLOCK_SIZE - pad_to_multiple_of = FLEX_ATTENTION_BLOCK_SIZE - - if eval_steps is not None and eval_strategy != 'steps': - print(f'Unsloth: `eval_steps = {eval_steps}` is ignored because `eval_strategy` is {getattr(eval_strategy, "value", eval_strategy)!r}. Set `eval_strategy = "steps"` to evaluate every `eval_steps` steps.') - - - # One dict so the filter sees the mirrored parameters AND `**kwargs`: - # filtering kwargs alone would double-bind any argument TRL renamed, - # since the new name is itself a mirrored parameter. - _unsloth_config_arguments = dict( - output_dir = output_dir, - per_device_train_batch_size = per_device_train_batch_size, - num_train_epochs = num_train_epochs, - max_steps = max_steps, - learning_rate = learning_rate, - lr_scheduler_type = lr_scheduler_type, - lr_scheduler_kwargs = lr_scheduler_kwargs, - warmup_steps = warmup_steps, - optim = optim, - optim_args = optim_args, - weight_decay = weight_decay, - adam_beta1 = adam_beta1, - adam_beta2 = adam_beta2, - adam_epsilon = adam_epsilon, - optim_target_modules = optim_target_modules, - gradient_accumulation_steps = gradient_accumulation_steps, - average_tokens_across_devices = average_tokens_across_devices, - max_grad_norm = max_grad_norm, - label_smoothing_factor = label_smoothing_factor, - bf16 = bf16, - fp16 = fp16, - bf16_full_eval = bf16_full_eval, - fp16_full_eval = fp16_full_eval, - tf32 = tf32, - gradient_checkpointing = gradient_checkpointing, - gradient_checkpointing_kwargs = gradient_checkpointing_kwargs, - torch_compile = torch_compile, - torch_compile_backend = torch_compile_backend, - torch_compile_mode = torch_compile_mode, - use_liger_kernel = use_liger_kernel, - liger_kernel_config = liger_kernel_config, - use_cache = use_cache, - neftune_noise_alpha = neftune_noise_alpha, - torch_empty_cache_steps = torch_empty_cache_steps, - auto_find_batch_size = auto_find_batch_size, - logging_strategy = logging_strategy, - logging_steps = logging_steps, - logging_first_step = logging_first_step, - log_on_each_node = log_on_each_node, - logging_nan_inf_filter = logging_nan_inf_filter, - include_num_input_tokens_seen = include_num_input_tokens_seen, - log_level = log_level, - log_level_replica = log_level_replica, - disable_tqdm = disable_tqdm, - report_to = report_to, - run_name = run_name, - project = project, - trackio_space_id = trackio_space_id, - eval_strategy = eval_strategy, - eval_steps = eval_steps, - eval_delay = eval_delay, - per_device_eval_batch_size = per_device_eval_batch_size, - prediction_loss_only = prediction_loss_only, - eval_on_start = eval_on_start, - eval_do_concat_batches = eval_do_concat_batches, - eval_use_gather_object = eval_use_gather_object, - eval_accumulation_steps = eval_accumulation_steps, - batch_eval_metrics = batch_eval_metrics, - save_only_model = save_only_model, - save_strategy = save_strategy, - save_steps = save_steps, - save_on_each_node = save_on_each_node, - save_total_limit = save_total_limit, - enable_jit_checkpoint = enable_jit_checkpoint, - push_to_hub = push_to_hub, - hub_token = hub_token, - hub_private_repo = hub_private_repo, - hub_model_id = hub_model_id, - hub_strategy = hub_strategy, - hub_always_push = hub_always_push, - hub_revision = hub_revision, - load_best_model_at_end = load_best_model_at_end, - metric_for_best_model = metric_for_best_model, - greater_is_better = greater_is_better, - ignore_data_skip = ignore_data_skip, - restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint, - full_determinism = full_determinism, - seed = seed, - data_seed = data_seed, - use_cpu = use_cpu, - accelerator_config = accelerator_config, - parallelism_config = parallelism_config, - dataloader_drop_last = dataloader_drop_last, - dataloader_num_workers = dataloader_num_workers, - dataloader_pin_memory = dataloader_pin_memory, - dataloader_persistent_workers = dataloader_persistent_workers, - dataloader_prefetch_factor = dataloader_prefetch_factor, - remove_unused_columns = remove_unused_columns, - label_names = label_names, - train_sampling_strategy = train_sampling_strategy, - length_column_name = length_column_name, - ddp_find_unused_parameters = ddp_find_unused_parameters, - ddp_bucket_cap_mb = ddp_bucket_cap_mb, - ddp_broadcast_buffers = ddp_broadcast_buffers, - ddp_backend = ddp_backend, - ddp_timeout = ddp_timeout, - fsdp = fsdp, - fsdp_config = fsdp_config, - deepspeed = deepspeed, - debug = debug, - skip_memory_metrics = skip_memory_metrics, - do_train = do_train, - do_eval = do_eval, - do_predict = do_predict, - resume_from_checkpoint = resume_from_checkpoint, - warmup_ratio = warmup_ratio, - logging_dir = logging_dir, - local_rank = local_rank, - model_init_kwargs = model_init_kwargs, - trust_remote_code = trust_remote_code, - router_aux_loss_coef = router_aux_loss_coef, - disable_dropout = disable_dropout, - dataset_num_proc = dataset_num_proc, - max_length = max_length, - pad_to_multiple_of = pad_to_multiple_of, - precompute_ref_log_probs = precompute_ref_log_probs, - precompute_ref_batch_size = precompute_ref_batch_size, - loss_type = loss_type, - beta = beta, - desirable_weight = desirable_weight, - undesirable_weight = undesirable_weight, - activation_offloading = activation_offloading, - sync_ref_model = sync_ref_model, - ref_model_mixup_alpha = ref_model_mixup_alpha, - ref_model_sync_steps = ref_model_sync_steps,**kwargs) - super().__init__(**_unsloth_filter_config_init_kwargs(KTOConfig, _unsloth_config_arguments, mirrored_from = __class__)) - self.vllm_sampling_params = vllm_sampling_params - self.unsloth_num_chunks = unsloth_num_chunks - if unsloth_grpo_mini_batch is not None: - if self.generation_batch_size >= unsloth_grpo_mini_batch: - self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch - else: - raise ValueError( - f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, " - f"which is self.per_device_train_batch_size * gradient_accumulation_steps." - ) - self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier - self.max_seq_length = max_seq_length - # Unsloth: keep the reentrant checkpoint path - if getattr(self, 'gradient_checkpointing', False): - _gc_kwargs = getattr(self, 'gradient_checkpointing_kwargs', None) or {} - if _gc_kwargs.get('context_fn') is None and not _gc_kwargs.get('debug', False): - _gc_kwargs['use_reentrant'] = True - self.gradient_checkpointing_kwargs = _gc_kwargs - from unsloth.optimizers.flash_adamw import is_flash_adamw_marked, mark_flash_adamw, unmark_flash_adamw - if _unsloth_flash_adamw or is_flash_adamw_marked(self): - mark_flash_adamw(self) - else: - unmark_flash_adamw(self) - -pass - -class _UnslothKTOTrainer(_BaseTrainer): - """""" - - loss_is_scaled_for_ga = False - - _tag_names = ["trl", "kto"] - _name = "KTO" - _paper = { - "title": "KTO: Model Alignment as Prospect Theoretic Optimization", - "id": "2402.01306", - # docstyle-ignore - "citation": textwrap.dedent("""\ - @article{ethayarajh2024kto, - title = {{KTO: Model Alignment as Prospect Theoretic Optimization}}, - author = {Kawin Ethayarajh and Winnie Xu and Niklas Muennighoff and Dan Jurafsky and Douwe Kiela}, - year = 2024, - eprint = {arXiv:2402.01306}, - }"""), - } - - def __init__( - self, - model: "str | PreTrainedModel | PeftModel", - ref_model: PreTrainedModel | None = None, - args: KTOConfig | None = None, - data_collator: DataCollator | None = None, - train_dataset: Dataset | IterableDataset | None = None, - eval_dataset: Dataset - | IterableDataset - | DatasetDict - | IterableDatasetDict - | dict[str, Dataset | IterableDataset] - | None = None, - processing_class: PreTrainedTokenizerBase | ProcessorMixin | None = None, - compute_metrics: Callable[[EvalPrediction], dict] | None = None, - callbacks: list[TrainerCallback] | None = None, - optimizers: tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None] = (None, None), - quantization_config: "BitsAndBytesConfig | None" = None, - peft_config: "PeftConfig | None" = None, - ): - # Args - if args is None: - model_name = model if isinstance(model, str) else get_config_model_id(model.config) - model_name = model_name.split("/")[-1] - args = KTOConfig(f"{model_name}-KTO") - - if train_dataset is None: - raise ValueError("`train_dataset` is required") - elif isinstance(train_dataset, IterableDataset): - # IterableDataset requires dispatch_batches=False because Accelerate's dispatch mode may try to concatenate - # batches from multiple processes, leading to mismatch errors. - if args.accelerator_config.dispatch_batches is True: - logger.warning( - "You are using an `IterableDataset` for training with `dispatch_batches=True`. `dispatch_batches` " - "is forced to `False` when using an `IterableDataset`. To remove this warning, unset " - "`dispatch_batches` in `KTOConfig` or set it to `False`." - ) - args.accelerator_config.dispatch_batches = False - elif not isinstance(train_dataset, Dataset): - raise TypeError( - f"`train_dataset` must be a `Dataset` or `IterableDataset`, got `{type(train_dataset).__name__}`." - ) - - # Model - # PEFT initializes the adapter weights randomly, so set_seed must be done before creating the model to ensure - # reproducibility. - set_seed(args.seed) - if isinstance(model, str): - model_init_kwargs = dict(args.model_init_kwargs or {}) # copy to avoid mutating model_init_kwargs - if quantization_config is not None: - if "quantization_config" in model_init_kwargs: - raise ValueError( - "You set `quantization_config` both as a trainer argument and in `args.model_init_kwargs`. " - "Please set it in only one place, preferably as a trainer argument." - ) - model_init_kwargs["quantization_config"] = quantization_config - # Distributed training requires device_map=None ["auto" fails] - if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: - model_init_kwargs["device_map"] = None - model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) - model_revision = model_init_kwargs.get("revision") - model = create_model_from_path(model, **model_init_kwargs) - else: - model_revision = None - if args.model_init_kwargs is not None: - logger.warning( - "You passed `model_init_kwargs` to the `KTOConfig`, but your model is already instantiated. " - "The `model_init_kwargs` will be ignored." - ) - if quantization_config is not None: - raise ValueError( - "You passed `quantization_config` to the trainer, but your model is already instantiated. " - "Quantization can only be applied when the model is loaded from a model identifier (`str`). " - "Either pass the model as a model identifier, or omit `quantization_config`." - ) - # Non-quantized models do not have the `is_loaded_in_{8,4}bit` attributes, whereas quantized models do - _is_quantized_model = getattr(model, "is_loaded_in_4bit", False) or getattr(model, "is_loaded_in_8bit", False) - if ref_model is model: - raise ValueError( - "`model` and `ref_model` cannot be the same object. In most cases you should omit `ref_model` and " - "we'll initialize it to a copy of `model` for you." - ) - - # Processing class - if processing_class is None: - processing_class = AutoProcessor.from_pretrained( - get_config_model_id(model.config), revision=model_revision, trust_remote_code=args.trust_remote_code - ) - if isinstance(processing_class, ProcessorMixin): - self._tokenizer = processing_class.tokenizer - self._is_vlm = True - elif isinstance(processing_class, PreTrainedTokenizerBase): - self._tokenizer = processing_class - self._is_vlm = False - else: - raise TypeError("The `processing_class` must be either a `PreTrainedTokenizerBase` or a `ProcessorMixin`") - if self._tokenizer.pad_token is None: - self._tokenizer.pad_token = self._tokenizer.eos_token - - # The model must agree with the tokenizer on the pad token from construction, so mirror it onto the model - # configs. - model.config.get_text_config().pad_token_id = self._tokenizer.pad_token_id - model.generation_config.pad_token_id = self._tokenizer.pad_token_id - - # PEFT - if False: - if not is_peft_available(): - raise ImportError( - "You passed `peft_config` but the `peft` library is not installed. " - "Install it with `pip install trl[peft]`." - ) - if not isinstance(peft_config, PeftConfig): - raise TypeError( - f"`peft_config` must be a `peft.PeftConfig` instance (e.g. `peft.LoraConfig`), " - f"got {type(peft_config).__name__}." - ) - if is_peft_model(model): - raise ValueError( - "You passed a `PeftModel` instance together with a `peft_config` to the trainer. Please first merge " - "and unload the existing adapter, save the resulting base model, and then pass that base model along " - "with the new `peft_config` to the trainer." - ) - # Create PEFT model - # ZeRO-3 + PEFT for non-quantized models: - # - PEFT's default autocast_adapter_dtype=True upcasts LoRA adapter params to fp32 even when the base model is bf16. - # - ZeRO-3's _allgather_params_coalesced allocates output buffers using the dtype of the first persistent parameter, - # so mixed-dtype persistent_parameters [bf16 base + fp32 LoRA] cause a TypeError on the first optimizer step. - # - Passing autocast_adapter_dtype=False keeps adapter params in the base model dtype [bf16], fixing the mismatch. - # - This is safe: the fp32 upcast is a QLoRA-specific concern [low-bit quantized base models], not needed for - # non-quantized bf16 training. - # - See: - # - TRL issue: https://github.com/huggingface/trl/issues/6089 - # - Upstream issue: https://github.com/deepspeedai/DeepSpeed/issues/8072 - get_peft_model_kwargs = {} - if args.deepspeed_plugin is not None and args.deepspeed_plugin.zero_stage == 3 and not _is_quantized_model: - get_peft_model_kwargs["autocast_adapter_dtype"] = False - model = get_peft_model(model, peft_config, **get_peft_model_kwargs) - - elif is_peft_model(model) and ref_model is None: - # If the model is a PEFT model with a pretrained adapter, we need to create a "ref" adapter that is a copy - # of the "default" adapter, so that we can use it as the reference model during KTO training. Before PEFT - # 0.20.0, only one adapter per model was supported when the LoRA config uses `target_parameters` [see - # peft#3340, fixed in peft#3350], so in that case we skip the "ref" adapter and compute the reference log - # probs with adapters disabled, i.e. with the base model. The fix only allows adapters targeting the same - # parameters, which holds here since the "ref" adapter reuses the "default" config. `target_parameters` - # itself was only added in PEFT 0.17.0, so the version check is bounded on both sides. - default_config = model.peft_config["default"] - if ( - isinstance(default_config, LoraConfig) - and Version("0.17.0") <= Version(peft.__version__) < Version("0.20.0") - and default_config.target_parameters - ): - logger.warning( - "PEFT<0.20.0 can't add a frozen reference adapter alongside one that uses `target_parameters` " - "(peft#3340], so the reference log probs are computed from the base model [adapters disabled]. " - "Upgrade to `peft>=0.20.0` to train against a copy of your adapter instead. If you wrapped the " - "model only to apply LoRA, pass a `peft_config` to the trainer instead; if you wrapped it " - "deliberately (pretrained adapter or custom init), note that the base model matches your adapter " - "only when it's freshly zero-initialized. If it is, this warning is safe to ignore." - ) - else: - model.add_adapter("ref", default_config) - for name, param in model.named_parameters(): - if ".default." in name: - ref_name = name.replace(".default.", ".ref.") - ref_param = model.get_parameter(ref_name) - ref_param.data.copy_(param.data) - - # PEFT + DeepSpeed ZeRO-3 requires reentrant checkpointing. For more details, see - # https://github.com/huggingface/trl/issues/2514#issuecomment-2692152703. - # Can be removed once https://github.com/deepspeedai/DeepSpeed/pull/8130 is merged and released. - if ( - is_peft_model(model) - and args.deepspeed_plugin is not None - and args.deepspeed_plugin.zero_stage == 3 - and args.gradient_checkpointing - ): - args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} - use_reentrant = args.gradient_checkpointing_kwargs.get("use_reentrant") - if use_reentrant is False: - logger.warning( - "You are using PEFT with DeepSpeed ZeRO-3 and gradient checkpointing with `use_reentrant=False`. " - "`use_reentrant` is forced to `True` in this configuration to ensure correct training. To remove " - "this warning, unset `use_reentrant` in `gradient_checkpointing_kwargs` or set it to `True`." - ) - args.gradient_checkpointing_kwargs["use_reentrant"] = True - - # When using gradient checkpointing with PEFT, we need to enable input gradients. transformers.Trainer normally - # handles this, but a bug currently prevents it; see https://github.com/huggingface/transformers/issues/42489 - if is_peft_model(model) and args.gradient_checkpointing: - model.enable_input_require_grads() - - # When using QLoRA, the PEFT adapter weights are converted to bf16 to follow the recommendations from the - # original paper [see https://huggingface.co/papers/2305.14314, paragraph 3]. Normally, this can be done by - # passing `autocast_adapter_dtype=False` to `get_peft_model`, but this option is not yet supported for - # quantized models. See: https://github.com/huggingface/peft/issues/2889 - # The DoRA magnitude vector is excluded: unlike LoRA A/B, its optimizer updates can be smaller than bf16 can - # represent, silently freezing it, see #7268. - if False: - for name, param in model.named_parameters(): - if param.requires_grad and "lora_magnitude_vector" not in name: - param.data = param.data.to(torch.bfloat16) - - # Vision dataset detection - dataset_sample = next(iter(train_dataset)) - self._is_vision_dataset = "image" in dataset_sample or "images" in dataset_sample - if self._is_vision_dataset and not self._is_vlm: - raise ValueError( - "The dataset appears to be vision-related (contains 'image' or 'images' keys), but the provided " - "model does not seem to be a vision-language model. Please check your model and dataset." - ) - if self._is_vision_dataset and args.precompute_ref_log_probs: - raise ValueError( - "`precompute_ref_log_probs=True` is not supported for vision datasets. For vision-language " - "models, all data processing is performed on the fly rather than upfront. " - "Set `precompute_ref_log_probs=False`." - ) - if self._is_vision_dataset and ("chosen" in dataset_sample or "rejected" in dataset_sample): - raise ValueError( - "Vision datasets must be in unpaired format with `completion` and `label` columns. " - "Paired format (`chosen`/`rejected`) is not supported for vision datasets because " - "iterating over the full dataset to unpair it would be too expensive for large image " - "collections. Unpair your dataset first: `dataset = unpair_preference_dataset(dataset)`." - ) - - # Data collator - calculate_kl = args.loss_type not in ["apo_zero_unpaired"] - if data_collator is None and not self._is_vision_dataset: - data_collator = DataCollatorForUnpairedPreference( - pad_token_id=self._tokenizer.pad_token_id, - max_length=args.max_length, - pad_to_multiple_of=args.pad_to_multiple_of, - ) - elif data_collator is None and self._is_vision_dataset: - data_collator = DataCollatorForVisionUnpairedPreference( - processor=processing_class, - max_length=args.max_length, - calculate_kl=calculate_kl, - pad_to_multiple_of=args.pad_to_multiple_of, - ) - - # Training arguments - self.beta = args.beta - self.precompute_ref_logps = args.precompute_ref_log_probs - self.loss_type = args.loss_type - self.desirable_weight = args.desirable_weight - self.undesirable_weight = args.undesirable_weight - self.calculate_KL = calculate_kl - if self.calculate_KL and args.train_sampling_strategy != "sequential": - raise ValueError( - f"Loss type `'{args.loss_type}'` estimates the KL divergence term and requires " - f"`train_sampling_strategy='sequential'` because the KL completion for each example is precomputed " - f"against its neighbors in a fixed-order batch; any other strategy breaks that pairing. " - f"Got `train_sampling_strategy='{args.train_sampling_strategy}'`." - ) - if self.calculate_KL and args.per_device_train_batch_size <= 1: - raise ValueError( - "Actual (not effective) batch size must be > 1. KTO will not work properly because the KL term will be equivalent to the implied reward." - ) - if is_peft_model(model) and isinstance(model.get_output_embeddings(), BaseTunerLayer): - # The log-probabilities are computed by multiplying the hidden states by `lm_head.weight` directly, so an - # adapter on the LM head would be ignored and never trained. - raise ValueError( - "Applying a PEFT adapter to `lm_head` is not supported: the log-probabilities are computed from " - "`lm_head.weight` directly, so the adapter would never be trained. Remove `'lm_head'` from your " - "`target_modules`." - ) - - # Dataset - # Skip dataset preparation for VLMs: tokenization and image processing happen on-the-fly in the collator. - if not self._is_vision_dataset: - train_dataset = self._prepare_dataset(train_dataset, processing_class, args, "train") - if eval_dataset is not None: - if isinstance(eval_dataset, dict): - eval_dataset = { - key: self._prepare_dataset(dataset, processing_class, args, key) - for key, dataset in eval_dataset.items() - } - else: - eval_dataset = self._prepare_dataset(eval_dataset, processing_class, args, "eval") - - # Transformers explicitly set use_reentrant=True in the past to silence a PyTorch warning, but the default was - # never updated once PyTorch switched to recommending use_reentrant=False. Until that change lands upstream - # [see https://github.com/huggingface/transformers/pull/43203] and is released [most likely in 5.0.0], we - # default to the recommended non-reentrant behavior here, while preserving any user-provided value. - if args.gradient_checkpointing and Version(transformers.__version__) < Version("5.0.0"): - args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} - args.gradient_checkpointing_kwargs.setdefault("use_reentrant", False) - - super().__init__( - model=model, - args=args, - data_collator=data_collator, - train_dataset=train_dataset, - eval_dataset=eval_dataset, - processing_class=processing_class, - compute_metrics=compute_metrics, - callbacks=callbacks, - optimizers=optimizers, - ) - - # With several GPUs visible and no distributed launcher, `Trainer` wraps the model in `nn.DataParallel`, whose - # replicas would all run the forward bound to the original model. Checked after `Trainer.__init__`, which sets - # `n_gpu` to 1 for a model split across devices with `device_map`. - if args.n_gpu > 1: - raise ValueError( - f"{type(self).__name__} does not support `nn.DataParallel`, which `Trainer` uses when several GPUs are " - "visible to a single process. Launch the script with `accelerate launch` or `torchrun`, or make a " - "single GPU visible with `CUDA_VISIBLE_DEVICES`." - ) - - # Initialize activation offloading context - if self.args.activation_offloading: - self.maybe_activation_offload_context = get_act_offloading_ctx_manager(model=self.model) - else: - self.maybe_activation_offload_context = contextlib.nullcontext() - - # MoE load-balancing auxiliary loss. `output_router_logits` in the config means the model returns its router - # logits; `router_aux_loss_coef` is the architecture's own coefficient, 0.0 when the config declares none. - text_config = model.config.get_text_config() - coef = args.router_aux_loss_coef - self.router_aux_loss_coef = getattr(text_config, "router_aux_loss_coef", 0.0) if coef is None else coef - if coef and not hasattr(text_config, "output_router_logits"): - raise ValueError( - f"`router_aux_loss_coef` is set to {coef} but {type(model).__name__} is not a Mixture-of-Experts model " - f"that returns its router logits, so there is no auxiliary loss to weight." - ) - self.aux_loss_enabled = hasattr(text_config, "output_router_logits") and self.router_aux_loss_coef != 0.0 - - # Reference model - if ref_model is None: - if is_peft_model(self.model) or (args.precompute_ref_log_probs and not self.is_fsdp_enabled): - # If PEFT is used, the reference model is not needed since the adapter can be disabled to revert to the - # initial model. If precompute_ref_log_probs is True, the reference model does not need to be kept in - # memory during training. Under FSDP, the policy isn't sharded until `train`, so the precompute still - # loads a reference model, freed once it is done. - self.ref_model = None - else: - ref_model_init_kwargs = dict(args.model_init_kwargs or {}) # copy to avoid mutating model_init_kwargs - if quantization_config is not None: - ref_model_init_kwargs["quantization_config"] = quantization_config - # Distributed training requires device_map=None ["auto" fails] - if self.args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: - ref_model_init_kwargs["device_map"] = None - ref_model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) - ref_model_path = get_config_model_id(self.model.config) - self.ref_model = create_model_from_path(ref_model_path, **ref_model_init_kwargs) - else: - self.ref_model = ref_model - - # Disable dropout in the models - if args.disable_dropout: - disable_dropout_in_model(model) - if self.ref_model is not None: - disable_dropout_in_model(self.ref_model) - - # Liger's fused linear cross-entropy replaces `model.forward` when training starts, which would drop the fused - # LM head, so only its layer kernels are applied - if args.use_liger_kernel: - warnings.warn( - "`use_liger_kernel=True` is deprecated and will be removed in v2.0.0. Use the Hub kernels instead, " - 'with `model_init_kwargs={"use_kernels": True}`.', - FutureWarning, - stacklevel=2, - ) - liger_kernel_config = args.liger_kernel_config or {} - if liger_kernel_config.get("fused_linear_cross_entropy"): - raise ValueError( - '`liger_kernel_config={"fused_linear_cross_entropy": True}` is not supported: it replaces the ' - "model's forward, which this trainer patches with a fused LM head." - ) - args.liger_kernel_config = {**liger_kernel_config, "fused_linear_cross_entropy": False} - - # Compute the per-token log-probabilities in chunks, without materializing the full logits - # `mean_logits` feeds the `logits/*` metrics - add_fused_lm_head( - self.model.get_base_model() if is_peft_model(self.model) else self.model, - outputs=("log_probs", "entropy", "mean_logits"), - ) - if self.ref_model is not None: - add_fused_lm_head(self.ref_model) - - # Initialize the metrics - self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} - self._total_train_tokens = 0 - - # Tensor parallel ranks are given the same batch, so a token count gathered across all processes repeats - # every token once per rank in the group. Context and sequence parallelism shard the batch before the loss is - # computed, so they need no such correction. `ParallelismConfig.tp_size` requires accelerate 1.10.0. - if Version(accelerate.__version__) >= Version("1.10.0") and self.accelerator.parallelism_config is not None: - self._tp_size = self.accelerator.parallelism_config.tp_size - else: - self._tp_size = 1 - - # Add tags to the model - self.model.add_model_tags(self._tag_names) - - # Hash before the reference model is sharded, since DTensor parameters can't be hashed - self._precompute_model_hash = hash_module(self.ref_model or self.model) if self.precompute_ref_logps else None - - if self.ref_model is not None: - if self.is_deepspeed_enabled: - self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator) - elif self.is_fsdp_enabled: - self.ref_model = prepare_fsdp(self.ref_model, self.accelerator) - else: - self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True) - - if args.sync_ref_model: - if is_peft_model(self.model): - raise NotImplementedError( - "You passed `sync_ref_model=True` while using a PEFT model, which is currently not supported. " - "With PEFT, KTOTrainer does not keep a separate reference model in memory; instead, it recovers " - "reference behavior by temporarily disabling the adapter. As a result, there is no standalone " - "`ref_model` instance to synchronize. Use `sync_ref_model=False`, or opt for full fine-tuning if " - "you need a synced reference model. If you need `sync_ref_model` to work with PEFT, please open a " - "feature request at https://github.com/huggingface/trl/issues." - ) - if args.precompute_ref_log_probs: - raise ValueError( - "You cannot use `sync_ref_model=True` together with `precompute_ref_log_probs=True`. " - "`precompute_ref_log_probs=True` assumes a fixed reference model, but with `sync_ref_model=True` " - "the reference model is periodically updated during training, making any precomputed reference " - "log-probs stale. Set `precompute_ref_log_probs=False` or disable `sync_ref_model`." - ) - self.add_callback(SyncRefModelCallback(ref_model=self.ref_model, accelerator=self.accelerator)) - - # Reference forwards during precompute reuse a single DeepSpeed inference engine [see `_precompute_ref_logps`]. - self._precompute_engine = None - if self.precompute_ref_logps: - self.train_dataset = self._precompute_ref_logps( - self.train_dataset, - "train", - self.args.precompute_ref_batch_size or self.args.per_device_train_batch_size, - ) - if self.eval_dataset is not None: - if isinstance(self.eval_dataset, dict): - self.eval_dataset = { - name: self._precompute_ref_logps( - dataset, name, self.args.precompute_ref_batch_size or self.args.per_device_eval_batch_size - ) - for name, dataset in self.eval_dataset.items() - } - else: - self.eval_dataset = self._precompute_ref_logps( - self.eval_dataset, - "eval", - self.args.precompute_ref_batch_size or self.args.per_device_eval_batch_size, - ) - if self.is_fsdp_enabled and ref_model is None: - self.ref_model = None - - def _prepare_dataset( - self, - dataset: Dataset | IterableDataset, - processing_class: PreTrainedTokenizerBase | ProcessorMixin, - args: KTOConfig, - dataset_name: str, - ) -> Dataset | IterableDataset: - # Build the kwargs for the `map` function - map_kwargs = {} - if isinstance(dataset, Dataset): # IterableDataset does not support num_proc - map_kwargs["num_proc"] = args.dataset_num_proc - - with global_then_local_main_first(): - # Extract the prompt if needed - first_example = next(iter(dataset)) - if "prompt" not in first_example: - if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` - map_kwargs["desc"] = f"Extracting prompt from {dataset_name} dataset" - dataset = dataset.map(extract_prompt, **map_kwargs) - - # Unpair the dataset if needed - first_example = next(iter(dataset)) - if "chosen" in first_example and "rejected" in first_example: - if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` - map_kwargs["desc"] = f"Unpairing {dataset_name} dataset" - dataset = unpair_preference_dataset(dataset, **map_kwargs) - - # Add EOS token if needed: non-conversational only - first_example = next(iter(dataset)) - if not is_conversational(first_example): - if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` - map_kwargs["desc"] = f"Adding EOS to {dataset_name} dataset" - - def add_eos(example, eos_token): - if not example["completion"].endswith(eos_token): - example["completion"] = example["completion"] + eos_token - return example - - dataset = dataset.map(add_eos, fn_kwargs={"eos_token": self._tokenizer.eos_token}, **map_kwargs) - - # Tokenize the dataset - if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` - map_kwargs["desc"] = f"Tokenizing {dataset_name} dataset" - - def tokenize_fn(example, processing_class): - tools = example.get("tools") - tools = json.loads(tools) if isinstance(tools, str) else tools - apply_chat_template_kwargs = {"tools": tools, **example.get("chat_template_kwargs", {})} - if is_conversational(example): - prompt_ids = _tokenize( - processing_class, - example["prompt"], - add_generation_prompt=True, - **apply_chat_template_kwargs, - )["input_ids"] - prompt_completion_ids = _tokenize( - processing_class, - example["prompt"] + example["completion"], - **apply_chat_template_kwargs, - )["input_ids"] - else: - prompt_ids = _tokenize(processing_class, example["prompt"])["input_ids"] - prompt_completion_ids = _tokenize(processing_class, example["prompt"] + example["completion"])[ - "input_ids" - ] - - # The completion starts where the tokenized prompt and prompt+completion diverge, which is not always - # after the prompt (see `common_prefix_length`) - prompt_len = common_prefix_length(prompt_ids, prompt_completion_ids) - if prompt_len < len(prompt_ids): - logger.warning_once( - "The tokenized prompt is not a prefix of the tokenized prompt+completion, usually because the " - "chat template renders the prompt alone differently or because tokens merge across the " - "boundary. The completion starts where they diverge, and this end of the prompt is left out " - "of the training sequence: " - f"{processing_class.decode(prompt_ids[prompt_len:])!r}" - ". The model is trained on a context that differs from the one it sees at inference. To avoid " - "it, use a chat template that renders the prompt the same way in both cases, or end the " - "prompt on a token boundary. This warning is shown once, but it likely applies to every " - "example in the dataset." - ) - - return { - "prompt_ids": prompt_ids[:prompt_len], - "completion_ids": prompt_completion_ids[prompt_len:], - } - - dataset = dataset.map(tokenize_fn, fn_kwargs={"processing_class": processing_class}, **map_kwargs) - - # Drop examples whose prompt alone fills `max_length`: with `keep_start` truncation the collator would - # remove every completion token, leaving no learning signal. `keep_end` keeps the completion end, so - # nothing is dropped there. - if args.max_length is not None: - if isinstance(dataset, Dataset): # `IterableDataset.filter` does not support `desc` - map_kwargs["desc"] = f"Dropping fully truncated examples from {dataset_name} dataset" - dataset = dataset.filter(lambda example: len(example["prompt_ids"]) < args.max_length, **map_kwargs) - - # Add KL completions if needed. The KL term is estimated from mismatched (prompt, completion) pairs, built - # by rotating the completions by +1 within each batch of size `per_device_train_batch_size`: - # (x_1, y_1), ..., (x_n, y_n) --> (x_1, y_n), (x_2, y_1), ..., (x_n, y_{n-1}). For best results, the - # mismatched outputs y' used to estimate the KL term for a batch should be the same set as the matched - # outputs y used to estimate the rewards in that batch, just paired with different x. - if self.calculate_KL: - if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` - map_kwargs["desc"] = f"Extracting KL {dataset_name} dataset" - - def add_kl_completion_ids(examples): - examples["KL_completion_ids"] = [examples["completion_ids"][-1]] + examples["completion_ids"][:-1] - return examples - - dataset = dataset.map( - add_kl_completion_ids, batched=True, batch_size=args.per_device_train_batch_size, **map_kwargs - ) - - # Calculate dataset desirability balance - if dataset_name == "train" and isinstance(dataset, Dataset): # IterableDataset does not support len - num_desirable = max(sum(dataset["label"]), 1) - num_undesirable = max(len(dataset["label"]) - num_desirable, 1) # "label" is binary - - if num_desirable != num_undesirable: - # The lower and upper bounds come from Eq. (8) of https://huggingface.co/papers/2402.01306 - des_weight_lower_bound = round((num_undesirable * self.undesirable_weight / num_desirable) * 1, 2) - des_weight_upper_bound = round( - (num_undesirable * self.undesirable_weight / num_desirable) * 1.33, 2 - ) - und_weight_lower_bound = round((num_desirable * self.desirable_weight / num_undesirable) / 1.33, 2) - und_weight_upper_bound = round((num_desirable * self.desirable_weight / num_undesirable) / 1, 2) - - des_weight_in_range = des_weight_lower_bound <= self.desirable_weight <= des_weight_upper_bound - und_weight_in_range = und_weight_lower_bound <= self.undesirable_weight <= und_weight_upper_bound - - if not (des_weight_in_range or und_weight_in_range): - logger.warning( - "You have different amounts of desirable/positive and undesirable/negative examples but the " - "weights on the desirable and undesirable losses don't seem to be in an ideal range. Based " - f"on your data, we recommend EITHER " - f"desirable_weight in [{des_weight_lower_bound}, {des_weight_upper_bound}] or " - f"undesirable_weight in [{und_weight_lower_bound}, {und_weight_upper_bound}] (but NOT BOTH). " - "See the documentation on how to optimally set these weights.", - ) - return dataset - - def _set_signature_columns_if_needed(self): - # If `self.args.remove_unused_columns` is True, non-signature columns are removed. - # By default, this method sets `self._signature_columns` to the model's expected inputs (usually, "input_ids" - # and "attention_mask"). - if self._signature_columns is None: - if self._is_vision_dataset: - self._signature_columns = [ - "prompt", - "completion", - "image", - "images", - "label", - "tools", - "chat_template_kwargs", - ] - else: - self._signature_columns = [ - "prompt_ids", - "completion_ids", - "KL_completion_ids", - "label", - "ref_logps", - "ref_KL_logps", - ] - - def _get_train_sampler(self, train_dataset: Dataset | None = None) -> Sampler | None: - if self.calculate_KL and Version(transformers.__version__) < Version("5.2.0"): - if train_dataset is None: - train_dataset = self.train_dataset - if train_dataset is None or not has_length(train_dataset): - return None - return SequentialSampler(train_dataset) - return super()._get_train_sampler(train_dataset) - - def _precompute_ref_logps(self, dataset: Dataset, name: str, batch_size: int) -> Dataset: - if isinstance(dataset, IterableDataset): - raise ValueError( - "`precompute_ref_log_probs=True` is not supported with IterableDataset. Please use a map-style " - "Dataset or set `precompute_ref_log_probs=False`." - ) - model_hash = self._precompute_model_hash - # Both inputs are rank-dependent under distributed training (ZeRO-3 shards the model), so broadcast rank 0's - # value so all ranks share one cache file. - fingerprint = [Hasher.hash((dataset._fingerprint, model_hash, self.calculate_KL))] - broadcast_object_list(fingerprint, from_process=0) - fingerprint = fingerprint[0] - cache_file = dataset._get_cache_file_path(fingerprint) - if os.path.exists(cache_file): - return concatenate_datasets([dataset, Dataset.from_file(cache_file)], axis=1) - - dataloader = DataLoader( - dataset, - batch_size=batch_size, - collate_fn=self.data_collator, - num_workers=self.args.dataloader_num_workers, - pin_memory=self.args.dataloader_pin_memory, - shuffle=False, - ) - data_loader = self.accelerator.prepare(dataloader) - - # This runs before the parent class prepares the model in `train`, so with DeepSpeed the parameters are still - # on CPU (ZeRO-1/2) and sharded (ZeRO-3). Wrap the model in an inference engine to place and gather them. Build - # it once and reuse it across precompute passes (train, eval, and later `evaluate` calls) - if self.ref_model is None and self.is_deepspeed_enabled: - if self._precompute_engine is None: - self._precompute_engine = prepare_deepspeed(self.model, self.accelerator) - model = self._precompute_engine - else: - model = self.ref_model or self.model - - ref_logps = [] - ref_KL_logps = [] - for padded_batch in tqdm( - iterable=data_loader, - desc=f"Computing reference log probs for {name} dataset", - disable=bool(os.environ.get("TQDM_DISABLE", "")), - ): - ref_logp, ref_KL_logp = self.compute_ref_log_probs(model, padded_batch) - if self.calculate_KL: - ref_logp, ref_KL_logp = self.accelerator.gather_for_metrics((ref_logp, ref_KL_logp)) - ref_KL_logps.append(ref_KL_logp.cpu()) - else: - ref_logp = self.accelerator.gather_for_metrics(ref_logp) - ref_logps.append(ref_logp.cpu()) - - ref_logps = torch.cat(ref_logps) - if self.calculate_KL: - ref_KL_logps = torch.cat(ref_KL_logps) - - if self.accelerator.is_main_process: - - def add_ref_logps(batch, indices): - result = {"ref_logps": ref_logps[indices]} - if self.calculate_KL: - result.update({"ref_KL_logps": ref_KL_logps[indices]}) - return result - - dataset.map( - add_ref_logps, - with_indices=True, - batched=True, - remove_columns=dataset.column_names, - new_fingerprint=fingerprint, - cache_file_name=cache_file, - desc=f"Caching reference log probs for {name} dataset", - ) - self.accelerator.wait_for_everyone() - - return concatenate_datasets([dataset, Dataset.from_file(cache_file)], axis=1) - - def compute_ref_log_probs(self, model, inputs): - """Computes reference log probabilities for a single padded batch.""" - adapter_context = contextlib.nullcontext() - if self.ref_model is None and is_peft_model(self.model): - unwrapped_model = self.accelerator.unwrap_model(self.model) - adapter_context = use_adapter( - unwrapped_model, adapter_name="ref" if "ref" in unwrapped_model.peft_config else None - ) - - with ( - torch.no_grad(), - disable_gradient_checkpointing(self.model, self.args.gradient_checkpointing_kwargs), - adapter_context, - self.accelerator.autocast(), - ): - labels = inputs["input_ids"].masked_fill(inputs["completion_mask"] == 0, -100) - per_token_logps = model( - input_ids=inputs["input_ids"], - attention_mask=inputs["attention_mask"], - labels=labels, - fused_lm_head=True, - ).log_probs - if self.calculate_KL: - KL_labels = inputs["KL_input_ids"].masked_fill(inputs["KL_completion_mask"] == 0, -100) - KL_per_token_logps = model( - input_ids=inputs["KL_input_ids"], - attention_mask=inputs["KL_attention_mask"], - labels=KL_labels, - fused_lm_head=True, - ).log_probs - - # Prompt-learning PEFT prepends virtual tokens to the outputs; keep the positions of the real tokens - completion_logps = per_token_logps[:, -(inputs["input_ids"].size(1) - 1) :].sum(-1) - - if self.calculate_KL: - KL_logps = KL_per_token_logps[:, -(inputs["KL_input_ids"].size(1) - 1) :].sum(-1) - else: - KL_logps = None - - return completion_logps, KL_logps - - def _compute_kl_logps(self, model, batch): - """Compute KL log probabilities for a given batch.""" - KL_logps = None - if self.calculate_KL: - _non_model_keys = { - "input_ids", - "attention_mask", - "completion_mask", - "KL_completion_mask", - "KL_token_type_ids", - "KL_mm_token_type_ids", - "label", - "ref_logps", - "ref_KL_logps", - } - KL_model_kwargs = {k: v for k, v in batch.items() if k not in _non_model_keys} - KL_model_kwargs["input_ids"] = KL_model_kwargs.pop("KL_input_ids") - KL_model_kwargs["attention_mask"] = KL_model_kwargs.pop("KL_attention_mask") - KL_model_kwargs["use_cache"] = False - # KL sequences have different widths from the main completion after flush_left; override token-type - # tensors with the KL-specific ones the collator built for exactly this purpose. - if "KL_token_type_ids" in batch: - KL_model_kwargs["token_type_ids"] = batch["KL_token_type_ids"] - if "KL_mm_token_type_ids" in batch: - KL_model_kwargs["mm_token_type_ids"] = batch["KL_mm_token_type_ids"] - - KL_labels = batch["KL_input_ids"].masked_fill(batch["KL_completion_mask"] == 0, -100) - with torch.no_grad(): - KL_per_token_logps = model(**KL_model_kwargs, labels=KL_labels, fused_lm_head=True).log_probs - # Prompt-learning PEFT prepends virtual tokens to the outputs; keep the positions of the real tokens - KL_logps = KL_per_token_logps[:, -(KL_labels.size(1) - 1) :].sum(-1) - return KL_logps - - def _compute_loss(self, model, inputs, return_outputs): - """Compute the KTO loss and other metrics for the given batch of inputs for train or test.""" - mode = "train" if self.model.training else "eval" - batch = {k: (v.to(self.accelerator.device) if isinstance(v, torch.Tensor) else v) for k, v in inputs.items()} - - labels = batch["label"] - num_chosen = labels.sum().to(self.accelerator.device) - num_rejected = (len(labels) - num_chosen).to(self.accelerator.device) - - KL_logps = self._compute_kl_logps(model, batch) - - _non_model_keys = { - "completion_mask", - "KL_input_ids", - "KL_attention_mask", - "KL_completion_mask", - "KL_token_type_ids", - "KL_mm_token_type_ids", - "label", - "ref_logps", - "ref_KL_logps", - } - model_kwargs = {k: v for k, v in batch.items() if k not in _non_model_keys} - model_kwargs["use_cache"] = False - if self.aux_loss_enabled: - model_kwargs["output_router_logits"] = True - completion_labels = batch["input_ids"].masked_fill(batch["completion_mask"] == 0, -100) - outputs = model(**model_kwargs, labels=completion_labels, fused_lm_head=True) - # Prompt-learning PEFT prepends virtual tokens to the outputs; keep the positions of the real tokens - seq_len = batch["input_ids"].size(1) - 1 - per_token_logps = outputs.log_probs[:, -seq_len:] - completion_logps = per_token_logps.sum(-1) - - if completion_logps.shape[0] != len(batch["label"]): - raise ValueError( - "There is a mismatch between the number of examples in this batch and the number of " - "examples for which an output sequence was predicted." - ) - - device = per_token_logps.device - bool_labels = torch.as_tensor(batch["label"], dtype=torch.bool, device=device) - chosen_idx = torch.nonzero(bool_labels, as_tuple=False).view(-1) - rejected_idx = torch.nonzero(~bool_labels, as_tuple=False).view(-1) - - chosen_logps = completion_logps.index_select(0, chosen_idx) - rejected_logps = completion_logps.index_select(0, rejected_idx) - - if self.precompute_ref_logps: - ref_chosen_logps = batch["ref_logps"].index_select(0, chosen_idx) - ref_rejected_logps = batch["ref_logps"].index_select(0, rejected_idx) - if self.calculate_KL: - ref_KL_logps = batch["ref_KL_logps"] - else: - ref_KL_logps = None - else: - ref_model_kwargs = {k: v for k, v in model_kwargs.items() if k != "output_router_logits"} - with torch.no_grad(), disable_gradient_checkpointing(self.model, self.args.gradient_checkpointing_kwargs): - if is_peft_model(self.model) and self.ref_model is None: - ref_model_unwrapped = self.accelerator.unwrap_model(self.model) - with use_adapter( - ref_model_unwrapped, adapter_name="ref" if "ref" in ref_model_unwrapped.peft_config else None - ): - ref_KL_logps = self._compute_kl_logps(self.model, batch) - ref_outputs = self.model(**ref_model_kwargs, labels=completion_labels, fused_lm_head=True) - else: - ref_KL_logps = self._compute_kl_logps(self.ref_model, batch) - ref_outputs = self.ref_model(**ref_model_kwargs, labels=completion_labels, fused_lm_head=True) - ref_per_token_logps = ref_outputs.log_probs[:, -seq_len:] - ref_completion_logps = ref_per_token_logps.sum(-1) - ref_chosen_logps = ref_completion_logps.index_select(0, chosen_idx) - ref_rejected_logps = ref_completion_logps.index_select(0, rejected_idx) - - if self.calculate_KL: - kl = (KL_logps - ref_KL_logps).mean().detach() - kl = self.accelerator.gather_for_metrics(kl).mean().clamp(min=0) - else: - kl = torch.zeros(1).to(chosen_logps.device) - # Chosen losses - if chosen_logps.shape[0] != 0 or ref_chosen_logps.shape[0] != 0: - chosen_logratios = chosen_logps - ref_chosen_logps - - if self.loss_type == "kto": - # Eqn (7) of the KTO paper (https://huggingface.co/papers/2402.01306) - chosen_losses = 1 - F.sigmoid(self.beta * (chosen_logratios - kl)) - elif self.loss_type == "apo_zero_unpaired": - # Unpaired variant of Eqn (7) of the APO paper (https://huggingface.co/papers/2408.06266) - # Use this loss when you believe the chosen outputs are better than your model's default output - chosen_losses = 1 - F.sigmoid(self.beta * chosen_logratios) - - chosen_rewards = self.beta * chosen_logratios.detach() - - else: - # lists can't be empty -- if they are, then accelerate.gather will hang - chosen_losses = torch.Tensor([]).to(self.accelerator.device) - chosen_rewards = torch.Tensor([]).to(self.accelerator.device) - # Rejected losses - if rejected_logps.shape[0] != 0 or ref_rejected_logps.shape[0] != 0: - rejected_logratios = rejected_logps - ref_rejected_logps - - if self.loss_type == "kto": - rejected_losses = 1 - F.sigmoid(self.beta * (kl - rejected_logratios)) - elif self.loss_type == "apo_zero_unpaired": - rejected_losses = F.sigmoid(self.beta * rejected_logratios) - - rejected_rewards = self.beta * rejected_logratios.detach() - else: - # lists can't be empty -- if they are, then accelerate.gather will hang - rejected_losses = torch.Tensor([]).to(self.accelerator.device) - rejected_rewards = torch.Tensor([]).to(self.accelerator.device) - losses = torch.cat( - (self.desirable_weight * chosen_losses, self.undesirable_weight * rejected_losses), - 0, - ) - - self._metrics[mode]["kl"].append(kl.item()) - - # Entropy - per_token_entropy = outputs.entropy[:, -seq_len:].detach() - mask = batch["completion_mask"][:, 1:] - entropy_sum = (per_token_entropy * mask).sum() - total_tokens = mask.sum() - - # Gather counts across ranks and weight-average - entropy_sum = self.accelerator.gather_for_metrics(entropy_sum).sum() - total_tokens = self.accelerator.gather_for_metrics(total_tokens).sum() - entropy = (entropy_sum / total_tokens).item() if total_tokens > 0 else 0.0 - self._metrics[mode]["entropy"].append(entropy) - - # Number of tokens - if mode == "train": - num_tokens_in_batch = self.accelerator.gather_for_metrics(batch["attention_mask"].sum()).sum().item() - self._total_train_tokens += num_tokens_in_batch // self._tp_size - self._metrics[mode]["num_tokens"] = [self._total_train_tokens] - - # Average logits for chosen and rejected completions - shift_completion_mask = batch["completion_mask"][:, 1:] - mean_logits = outputs.mean_logits[:, -seq_len:] - chosen_mean_logits = mean_logits.index_select(0, chosen_idx) - rejected_mean_logits = mean_logits.index_select(0, rejected_idx) - chosen_mask = shift_completion_mask.index_select(0, chosen_idx) - rejected_mask = shift_completion_mask.index_select(0, rejected_idx) - total_chosen_logits = chosen_mean_logits[chosen_mask.bool()].sum() - total_chosen_tokens = chosen_mask.sum() - total_rejected_logits = rejected_mean_logits[rejected_mask.bool()].sum() - total_rejected_tokens = rejected_mask.sum() - total_chosen_logits = self.accelerator.gather_for_metrics(total_chosen_logits).sum().item() - total_chosen_tokens = self.accelerator.gather_for_metrics(total_chosen_tokens).sum().item() - total_rejected_logits = self.accelerator.gather_for_metrics(total_rejected_logits).sum().item() - total_rejected_tokens = self.accelerator.gather_for_metrics(total_rejected_tokens).sum().item() - if total_chosen_tokens > 0: - self._metrics[mode]["logits/chosen"].append(total_chosen_logits / total_chosen_tokens) - if total_rejected_tokens > 0: - self._metrics[mode]["logits/rejected"].append(total_rejected_logits / total_rejected_tokens) - - all_num_chosen = self.accelerator.gather_for_metrics(num_chosen).sum().item() - all_num_rejected = self.accelerator.gather_for_metrics(num_rejected).sum().item() - - if all_num_chosen > 0: - self._metrics[mode]["rewards/chosen"].append( - self.accelerator.gather_for_metrics(chosen_rewards.nansum()).nansum().item() / all_num_chosen - ) - self._metrics[mode]["logps/chosen"].append( - self.accelerator.gather_for_metrics(chosen_logps.nansum()).nansum().item() / all_num_chosen - ) - - if all_num_rejected > 0: - self._metrics[mode]["rewards/rejected"].append( - self.accelerator.gather_for_metrics(rejected_rewards.nansum()).nansum().item() / all_num_rejected - ) - self._metrics[mode]["logps/rejected"].append( - self.accelerator.gather_for_metrics(rejected_logps.nansum()).nansum().item() / all_num_rejected - ) - - if all_num_chosen > 0 and all_num_rejected > 0: - self._metrics[mode]["rewards/margins"].append( - self._metrics[mode]["rewards/chosen"][-1] - self._metrics[mode]["rewards/rejected"][-1] - ) - - loss = losses.nanmean() - if self.aux_loss_enabled: - aux_loss = outputs.aux_loss - loss = loss + self.router_aux_loss_coef * aux_loss - self._metrics[mode]["aux_loss"].append(self.accelerator.gather_for_metrics(aux_loss).mean().item()) - - if return_outputs: - warnings.warn( - "`return_outputs=True`, and the logits it passes to `compute_metrics`, are deprecated and will be removed " - "in v2.0.0. The fused LM head does not build the logits, so they are recomputed with an extra forward pass.", - FutureWarning, - stacklevel=2, - ) - with torch.no_grad(): - outputs = model(**model_kwargs) - return (loss, outputs) if return_outputs else loss - - def evaluate( - self, - eval_dataset: Dataset - | IterableDataset - | DatasetDict - | IterableDatasetDict - | dict[str, Dataset | IterableDataset] - | None = None, - ignore_keys: list[str] | None = None, - metric_key_prefix: str = "eval", - ) -> dict[str, float]: - # When a dataset is passed directly to `evaluate` (e.g. a held-out test set), preprocess it the same way - # `__init__` does, so that `evaluate` accepts the same dataset types as the trainer. `_prepare_dataset` is - # idempotent: it skips datasets that are already tokenized. A `str` selects a dataset that was already prepared - # at init time, so it's left untouched. - if not self._is_vision_dataset and eval_dataset is not None and not isinstance(eval_dataset, str): - # Full fine-tuning with no `ref_model` uses `self.model` as the reference, which is only valid before - # training. After a step (`global_step > 0`) it's the trained policy, so we can't precompute a correct - # reference here. (PEFT is exempt: the reference is recovered by disabling the adapter.) Checked before - # tokenizing so we fail fast. - if ( - self.precompute_ref_logps - and self.ref_model is None - and not is_peft_model(self.model) - and self.state.global_step > 0 - ): - raise ValueError( - "Cannot compute reference log-probs for a dataset passed to `evaluate()` after training has " - "started, because `precompute_ref_log_probs=True` and no `ref_model` was provided (full " - "fine-tuning). In this setup the reference model is not kept in memory, so it is only available " - "before training. Provide this dataset as `eval_dataset` at initialization, pass an explicit " - "`ref_model`, or set `precompute_ref_log_probs=False`." - ) - if isinstance(eval_dataset, dict): - eval_dataset = { - key: self._prepare_dataset(dataset, self.processing_class, self.args, key) - for key, dataset in eval_dataset.items() - } - else: - eval_dataset = self._prepare_dataset(eval_dataset, self.processing_class, self.args, "eval") - # With `precompute_ref_log_probs`, `_compute_loss` reads the reference log-probs from the batch, so they - # must be precomputed here as well, mirroring `__init__`. - if self.precompute_ref_logps: - batch_size = self.args.precompute_ref_batch_size or self.args.per_device_eval_batch_size - if isinstance(eval_dataset, dict): - eval_dataset = { - name: self._precompute_ref_logps(dataset, name, batch_size) - for name, dataset in eval_dataset.items() - } - else: - eval_dataset = self._precompute_ref_logps(eval_dataset, "eval", batch_size) - # Call `super().evaluate()` once per split ourselves instead of handing the whole dict to it: `Trainer.evaluate` - # would otherwise recurse into `self.evaluate` per split, re-entering this override on an already-prepared dataset. - if isinstance(eval_dataset, dict): - metrics = {} - for name, dataset in eval_dataset.items(): - metrics.update( - super().evaluate( - eval_dataset=dataset, - ignore_keys=ignore_keys, - metric_key_prefix=f"{metric_key_prefix}_{name}", - ) - ) - return metrics - return super().evaluate( - eval_dataset=eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix - ) - - def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): - try: - return self._compute_loss(model, inputs, return_outputs) - except ValueError as e: - if "Image features and image tokens do not match" in str(e) and self.args.max_length is not None: - raise ValueError( - f"The current `max_length` ({self.args.max_length}) is too short and causes image placeholder " - f"tokens in `input_ids` to be truncated, while the corresponding image features remain intact. " - f"Please increase `max_length` or set it to `None` to disable truncation." - ) from e - raise - - # Override training step to add activation offloading context. - def training_step(self, *args, **kwargs): - with self.maybe_activation_offload_context: - return super().training_step(*args, **kwargs) - - def log(self, logs: dict[str, float], start_time: float | None = None) -> None: - mode = "train" if self.model.training else "eval" - metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()} # average the metrics - # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` - # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. - if mode == "eval": - metrics = {f"eval_{key}": val for key, val in metrics.items()} - logs.update(metrics) - super().log(logs, start_time) - self._metrics[mode].clear() - - # During eval, Trainer calls prediction_step. If no labels are present in the inputs, it only runs forward and - # returns logits. We override prediction_step to force compute_loss, because this trainer doesn't involve labels. - def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys: list[str] | None = None): - inputs = self._prepare_inputs(inputs) - with torch.no_grad(), self.compute_loss_context_manager(): - if prediction_loss_only: - loss = self.compute_loss(model, inputs, return_outputs=False) - logits, labels = None, None - else: - loss, outputs = self.compute_loss(model, inputs, return_outputs=True) - logits, labels = outputs.logits, inputs["input_ids"] - return loss, logits, labels - - # Ensure the model card is saved along with the checkpoint - def _save_checkpoint(self, model, trial): - if self.args.hub_model_id is None: - model_name = Path(self.args.output_dir).name - else: - model_name = self.args.hub_model_id.split("/")[-1] - self.create_model_card(model_name=model_name) - super()._save_checkpoint(model, trial) -class UnslothKTOTrainer(_UnslothKTOTrainer): - """ - - Trainer for Kahneman-Tversky Optimization (KTO) method. This algorithm was initially proposed in the paper [KTO: - Model Alignment as Prospect Theoretic Optimization](https://huggingface.co/papers/2402.01306). This class is a - wrapper around the [`~transformers.Trainer`] class and inherits all of its attributes and methods. - - Example: - - ```python - >>> from trl import KTOTrainer - >>> from datasets import load_dataset - - >>> dataset = load_dataset("trl-lib/kto-mix-14k", split="train") - - >>> trainer = KTOTrainer( - ... model="Qwen/Qwen2.5-0.5B-Instruct", - ... train_dataset=dataset, - ... ) - >>> trainer.train() - ``` - - Args: - model (`str` or [`~transformers.PreTrainedModel`] or [`~peft.PeftModel`]): - Model to be trained. Can be either: - - - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a - path to a *directory* containing model weights saved using - [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded - using `.from_pretrained` (where `` is derived from the model - config) with the keyword arguments in `args.model_init_kwargs`. If `dtype` is not specified in - `args.model_init_kwargs`, it defaults to `float32`. This differs from - [`~transformers.PreTrainedModel.from_pretrained`], where (since Transformers v5) the dtype is inferred - from the model config. - - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. - - A [`~peft.PeftModel`] object. Only causal language models are supported. - ref_model ([`~transformers.PreTrainedModel`], *optional*): - Reference model used to compute the reference log probabilities. - - - If provided, this model is used directly as the reference policy. - - If `None`, the trainer will automatically use the initial policy corresponding to `model`, i.e. the model - state before KTO training starts. - args ([`KTOConfig`], *optional*): - Configuration for this trainer. If `None`, a default configuration is used. - data_collator ([`~transformers.DataCollator`], *optional*): - Function to use to form a batch from a list of elements of the processed `train_dataset` or `eval_dataset`. - Will default to [`~trainer.kto_trainer.DataCollatorForUnpairedPreference`] if the model is a language model - and [`~trainer.kto_trainer.DataCollatorForVisionUnpairedPreference`] if the model is a vision-language - model. Custom collators must truncate sequences before padding; the trainer does not apply post-collation - truncation. - train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): - Dataset to use for training. This trainer supports [unpaired preference](#unpaired-preference) type. The - format of the samples can be either: - - - [Standard](dataset_formats#standard): Each sample contains plain text. - - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role - and content). - - When `train_dataset` is an [`~datasets.IterableDataset`] (e.g. a streaming dataset), `max_steps` must be - set in the training arguments, since its length cannot be inferred and the total number of training steps - is required to bound the training loop and configure the learning rate scheduler. - eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`], [`~datasets.DatasetDict`], [`~datasets.IterableDatasetDict`] or `dict[str, Dataset | IterableDataset]`): - Dataset to use for evaluation. It must meet the same requirements as `train_dataset`. - processing_class ([`~transformers.PreTrainedTokenizerBase`] or [`~transformers.ProcessorMixin`], *optional*): - Processing class used to process the data. The padding side must be set to "left". If `None`, the - processing class is loaded from the model's name with [`~transformers.AutoProcessor.from_pretrained`]. A - padding token, `tokenizer.pad_token`, must be set. If the processing class has not set a padding token, - `tokenizer.eos_token` will be used as the default. - compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): - The function that will be used to compute metrics at evaluation. Must take a - [`~transformers.EvalPrediction`] and return a dictionary string to metric values. When passing - [`SFTConfig`] with `batch_eval_metrics` set to `True`, your `compute_metrics` function must take a boolean - `compute_result` argument. This will be triggered after the last eval batch to signal that the function - needs to calculate and return the global summary statistics rather than accumulating the batch-level - statistics. The logits it receives are recomputed with an extra forward pass. This is deprecated and will - be removed in v2.0.0. - callbacks (list of [`~transformers.TrainerCallback`], *optional*): - List of callbacks to customize the training loop. Will add those to the list of default callbacks detailed - in [here](https://huggingface.co/docs/transformers/main_classes/callback). - - If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`] - method. - optimizers (`tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None]`, *optional*, defaults to `(None, None)`): - A tuple containing the optimizer and the scheduler to use. Will default to an instance of `AdamW` on your - model and a scheduler given by [`~transformers.get_linear_schedule_with_warmup`] controlled by `args`. - quantization_config ([`~transformers.BitsAndBytesConfig`], *optional*): - Quantization configuration used when loading the model from a model identifier. Combine with `peft_config` - for QLoRA training. Ignored if the model is already instantiated. - peft_config ([`~peft.PeftConfig`], *optional*): - PEFT configuration used to wrap the model. If `None`, the model is not wrapped. - - """ - def __init__( - self, - model, - ref_model = None, - args = None, - data_collator = None, - train_dataset = None, - eval_dataset = None, - processing_class = None, - compute_metrics = None, - callbacks = None, - quantization_config = None, - peft_config = None, - **kwargs - ): - if args is None: args = UnslothKTOConfig() - _unsloth_precision_requested = tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) - use_bf16 = getattr(args, 'bf16', False) - if type(use_bf16) is not bool: use_bf16 = False - use_fp16 = getattr(args, 'fp16', False) - if type(use_fp16) is not bool: use_fp16 = False - force_float32 = False - try: - from unsloth_zoo.device_type import device_is_bf16_supported as _bf16_supported - except Exception: - _bf16_supported = torch.cuda.is_bf16_supported - try: - from unsloth_zoo.device_type import DEVICE_TYPE as _unsloth_device_type - except Exception: - _unsloth_device_type = None - full_finetuning = getattr(model, '_unsloth_full_finetuning', None) - if full_finetuning is None: full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1' - model_forced_float32 = getattr(model, '_unsloth_forced_float32', None) - if model_forced_float32 is None: model_forced_float32 = os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1' - if model_forced_float32 and not (full_finetuning and _bf16_supported()): - print('Unsloth: Switching to float32 training since model cannot work with float16') - force_float32 = True - mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') - dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None) - if dtype is None: dtype = model.get_input_embeddings().weight.dtype - from unsloth_zoo.utils import _get_dtype - dtype = _get_dtype(dtype) - float16 = dtype == torch.float16 - bfloat16 = dtype == torch.bfloat16 - float32 = dtype == torch.float32 - user_float32 = bool(getattr(model, '_unsloth_user_float32', False)) - if full_finetuning: - if bfloat16 and use_fp16: use_fp16 = False - if float16 and use_bf16: use_bf16 = False - if not force_float32 and bfloat16 and use_fp16 and _unsloth_device_type == 'xpu' and not _bf16_supported(): - print('Unsloth: The model was loaded in bfloat16 on purpose, so training uses bfloat16 mixed precision instead of float16.') - use_fp16 = False - use_bf16 = True - args.fp16 = False - args.bf16 = True - if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`') - if not force_float32 and (bfloat16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`') - if force_float32: - # Forced float32 training - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32' and float32 and user_float32 and not _bf16_supported(): - print('Unsloth: Model is in float32 and this GPU has no bfloat16 support, so training stays in float32. Pass fp16 = True to force float16 mixed precision instead.') - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32': - # Mixed precision training. bf16 only if the GPU supports it; V100/T4 use fp16. - use_bf16_amp = (not float16) and (_bf16_supported() or (bfloat16 and _unsloth_device_type == 'xpu')) - args.fp16 = not use_bf16_amp - args.bf16 = use_bf16_amp - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16_amp else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16_amp else 'fp16' - # args.mixed_precision is a new argument which needs to be set now - elif mixed_precision_dtype == 'bfloat16': - # Both False since bfloat16 full finetuning doesn't do any autocasting. - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif use_bf16 or use_fp16: - # transformers <5 exported this itself from fp16/bf16; 5.x dropped the write, so an - # explicit flag left it unset and GRPO readers defaulted to 'fp16', wrapping a - # bfloat16 model in a float16 autocast. See unslothai/unsloth#4891. - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16 else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16 else 'fp16' - - if getattr(args, 'eval_dataset', None) is not None and getattr(args, 'eval_strategy', 'no') == 'no': - args.eval_strategy = 'steps' - if getattr(args, 'eval_steps', None) is None: args.eval_steps = 0.1 - ga_steps = getattr(args, 'gradient_accumulation_steps', None) - if ga_steps is not None and ga_steps > 1: - from transformers import __version__ as transformers_version - if Version(transformers_version) <= Version('4.45.2'): - print('**** Unsloth: Please use our fixed gradient_accumulation_steps by updating transformers, TRL and Unsloth!\n' - '`pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth transformers trl unsloth_zoo`') - if getattr(args, 'eval_strategy', 'no') != 'no': - eval_bsz = getattr(args, 'per_device_eval_batch_size', 8) - eval_generations = getattr(args, 'num_generations_eval', None) or getattr(args, 'num_generations', None) or 1 - if eval_bsz == 8 and args.per_device_train_batch_size < eval_bsz and (args.per_device_train_batch_size * args.world_size) % eval_generations == 0: args.per_device_eval_batch_size = args.per_device_train_batch_size - if getattr(args, 'eval_accumulation_steps', None) is None and ga_steps is not None: args.eval_accumulation_steps = ga_steps - fp16_full_eval = getattr(args, 'fp16_full_eval', False) - if type(fp16_full_eval) is not bool: fp16_full_eval = False - bf16_full_eval = getattr(args, 'bf16_full_eval', False) - if type(bf16_full_eval) is not bool: bf16_full_eval = False - if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True - if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False - if force_float32: - args.bf16_full_eval = False - args.fp16_full_eval = False - elif os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16': - args.bf16_full_eval = True - args.fp16_full_eval = False - elif not bf16_full_eval and not fp16_full_eval: - args.bf16_full_eval = args.bf16 - args.fp16_full_eval = args.fp16 - if tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) != _unsloth_precision_requested: - from unsloth.models.loader_utils import sync_deepspeed_mixed_precision - sync_deepspeed_mixed_precision(args) - if model is not None: - from unsloth.models.loader_utils import apply_deepspeed_zero3_model_setup - apply_deepspeed_zero3_model_setup(model) - if model is not None: - _warnings_issued = getattr(model, 'warnings_issued', None) - if _warnings_issued is None: - model.warnings_issued = {} - elif not isinstance(_warnings_issued, dict): - try: - model.warnings_issued = dict(_warnings_issued) - except Exception: - model.warnings_issued = {} - if 'max_seq_length' not in locals() and not hasattr(args, 'max_seq_length'): - pass - else: - model_max_seq_length = getattr(model, 'max_seq_length', None) - args_max_seq_length = getattr(args, 'max_seq_length', None) - if args_max_seq_length is None and model_max_seq_length is not None: - max_seq_length = model.max_seq_length - if hasattr(args, 'max_seq_length'): args.max_seq_length = max_seq_length - elif args_max_seq_length is not None and model_max_seq_length is not None: - if args_max_seq_length > model_max_seq_length: - print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but ' - 'the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.') - args.max_seq_length = model_max_seq_length - _unsloth_model_msl = getattr(model, 'max_seq_length', None) - if isinstance(_unsloth_model_msl, int) and _unsloth_model_msl > 0 and hasattr(args, 'max_length'): - if args.max_length is None or args.max_length > _unsloth_model_msl: - print('Unsloth: `max_length = ' + str(args.max_length) + '` exceeds the model max_seq_length of ' + str(_unsloth_model_msl) + ', so it is reduced to ' + str(_unsloth_model_msl) + '.') - args.max_length = _unsloth_model_msl - _unsloth_mpl = getattr(args, 'max_prompt_length', 0) - if (_unsloth_mpl is None and args.max_length <= 128) or (_unsloth_mpl is not None and _unsloth_mpl >= args.max_length): - args.max_prompt_length = args.max_length // 2 - if model is not None and hasattr(model, 'for_training'): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right' - if 'processing_class' in locals(): - if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right' - if hasattr(processing_class, 'tokenizer') and hasattr(processing_class.tokenizer, 'padding_side'): processing_class.tokenizer.padding_side = 'right' - __tokenizer = processing_class if 'processing_class' in locals() else tokenizer - from unsloth_zoo.vision_utils import UnslothVisionDataCollator - if not isinstance(data_collator, UnslothVisionDataCollator): - if isinstance(data_collator, DataCollatorForSeq2Seq) and 'labels' not in _unsloth_dataset_column_names(train_dataset): - data_collator = TransformersDataCollatorForLanguageModeling( - __tokenizer, - mlm = False, - mlm_probability = 0.0, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling) and 'labels' in _unsloth_dataset_column_names(train_dataset): - data_collator = DataCollatorForSeq2Seq( - __tokenizer, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - else: - if hasattr(args, 'remove_unused_columns'): args.remove_unused_columns = False - if hasattr(args, 'dataset_text_field'): args.dataset_text_field = '' - if hasattr(args, 'dataset_kwargs'): args.dataset_kwargs = {'skip_prepare_dataset': True} - if not isinstance(data_collator, UnslothVisionDataCollator): - if not hasattr(__tokenizer, 'pad') and hasattr(__tokenizer, 'tokenizer'): - if isinstance(data_collator, DataCollatorForSeq2Seq): - data_collator = DataCollatorForSeq2Seq( - __tokenizer.tokenizer, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling): - data_collator = TransformersDataCollatorForLanguageModeling( - __tokenizer.tokenizer, - mlm = False, - mlm_probability = 0.0, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - other_metrics = [] - - from unsloth_zoo.logging_utils import PatchRLStatistics - PatchRLStatistics('kto_trainer', other_metrics) - - # [TODO] Fix up DataParallel multiplying batch sizes - # [TODO] DDP works, but DP seems to not work? [TODO] - if getattr(args, "parallel_mode", None) == ParallelMode.NOT_DISTRIBUTED and args.n_gpu > 1: - if getattr(args, "_n_gpu", 1) != 1: - args._n_gpu = 1 - if "model" in locals(): - try: - from unsloth.models._gc_config_default import honour_gradient_checkpointing_off - honour_gradient_checkpointing_off(args, model) - except Exception: - pass - if "model" in locals() and hasattr(model, "for_training"): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if "model" in locals(): - try: - from unsloth.models._ddp_defaults import unsloth_default_ddp_find_unused_parameters - unsloth_default_ddp_find_unused_parameters(args, model) - except Exception: - pass - super().__init__( - model = model, - ref_model = ref_model, - args = args, - data_collator = data_collator, - train_dataset = train_dataset, - eval_dataset = eval_dataset, - processing_class = processing_class, - compute_metrics = compute_metrics, - callbacks = callbacks, - quantization_config = quantization_config, - peft_config = peft_config,**kwargs) - # A tensor parallel model (tp_native.tp_trainer_setup) stays on its training path: for_inference would - # route the decoder layers onto decode paths outside the tensor parallel regions. - if "model" in locals() and hasattr(model, "for_inference") and not getattr(self, "_unsloth_tensor_parallel", False): - model.for_inference() - if hasattr(self, 'neftune_hook_handle'): - self.neftune_hook_handle.remove() - if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle - if getattr(args, 'neftune_noise_alpha', None) is not None: - model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha - pass - if hasattr(self, 'accelerator'): - scaler = self.accelerator.scaler - current_model = model - while hasattr(current_model, 'model'): - current_model.accelerator_scaler = scaler - current_model = current_model.model - current_model.accelerator_scaler = scaler - pass - if hasattr(self, 'train'): - self.train = MethodType(prepare_for_training_mode(self.__class__.train), self) - pass - if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'): - _vllm_tok = self.llm.get_tokenizer() - _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None) - if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None: - _vllm_tok.chat_template = _pc.chat_template - pass - if getattr(self, 'aux_loss_enabled', False) and hasattr(getattr(self, 'model', None), 'config'): - _text_config = self.model.config - if hasattr(_text_config, 'get_text_config'): _text_config = _text_config.get_text_config() - _n_experts = [getattr(_text_config, _k) for _k in ('num_local_experts', 'num_experts', 'n_routed_experts', 'moe_num_experts') if isinstance(getattr(_text_config, _k, None), int)] - if _n_experts and all(_n <= 1 for _n in _n_experts): - self.aux_loss_enabled = False - _text_config.output_router_logits = False - pass - -pass - - -if hasattr(logger, "addFilter"): - import logging - class HideLoggingMessage(logging.Filter): - def __init__(self, text): self.text = text - def filter(self, x): return not (self.text in x.getMessage()) - pass - logger.addFilter(HideLoggingMessage("`use_cache=True`")) - diff --git a/unsloth_compiled_cache/UnslothORPOTrainer.py b/unsloth_compiled_cache/UnslothORPOTrainer.py deleted file mode 100644 index 97d19c169cb5119115ad20c0ebcd24f92c730f44..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/UnslothORPOTrainer.py +++ /dev/null @@ -1,2582 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from unsloth_zoo.temporary_patches.common import torch_compile -from unsloth_zoo.temporary_patches.common import _maybe_compile -import functools -import transformers -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from trl.experimental.orpo.orpo_trainer import (Any, AutoModelForCausalLM, AutoTokenizer, BaseImageProcessor, Callable, DPODataCollatorWithPadding, DataCollator, DataLoader, Dataset, EvalLoopOutput, F, FeatureExtractionMixin, Literal, ORPOConfig, ORPOTrainer, Path, PeftConfig, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, TrainerCallback, Version, _BaseTrainer, add_bos_token_if_needed, add_eos_token_if_needed, autocast, defaultdict, disable_dropout_in_model, get_config_model_id, get_peft_model, global_then_local_main_first, inspect, is_comet_available, is_peft_available, is_peft_model, is_torch_fx_proxy, is_torch_xla_available, is_wandb_available, log1mexp, log_table_to_comet_experiment, logger, maybe_apply_chat_template, maybe_extract_prompt, nn, np, nullcontext, pad_to_length, pd, peft_module_casting_to_bf16, prepare_model_for_kbit_training, random, selective_log_softmax, set_seed, textwrap, torch, transformers, AutoModelForCausalLM, AutoTokenizer, BaseImageProcessor, Callable, DPODataCollatorWithPadding, DataCollator, Dataset, EvalLoopOutput, F, FeatureExtractionMixin, ORPOConfig, ORPOTrainer, PeftConfig, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, TrainerCallback, Version, autocast, defaultdict, disable_dropout_in_model, get_config_model_id, get_peft_model, global_then_local_main_first, inspect, is_comet_available, is_peft_available, is_peft_model, is_wandb_available, logger, maybe_apply_chat_template, maybe_extract_prompt, nn, np, pd, peft_module_casting_to_bf16, prepare_model_for_kbit_training, random, set_seed, torch, transformers, F, PreTrainedModel, is_peft_available, logger, torch, F, np, F, nn, np, torch, F, PreTrainedModel, add_bos_token_if_needed, add_eos_token_if_needed, is_torch_xla_available, logger, nn, np, torch) - - -import os -import math -import logging -from typing import * -from dataclasses import dataclass, field -from packaging.version import Version -import torch -import numpy as np -from contextlib import nullcontext -from torch.nn import functional as F -import inspect -from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling -from transformers.training_args import ParallelMode -from unsloth_zoo.device_type import DEVICE_TYPE, DEVICE_TYPE_TORCH, device_synchronize - -# Wrap trainer with padding to right and enable training mode -import functools -from types import MethodType -try: - from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers -except: - def reset_unsloth_gradient_checkpointing_buffers(): pass -# Canonical reset lives in unsloth.models._utils so the SFT auto-packing wrapper and the plain -# Trainer loop can import the same helper; fall back to a no-op only if it can't be imported. -try: - from unsloth.models._utils import _unsloth_reset_stray_compile_cache -except Exception: - def _unsloth_reset_stray_compile_cache(self): pass -try: - from unsloth.models._utils import _unsloth_dataset_column_names -except Exception: - def _unsloth_dataset_column_names(dataset): return dataset.column_names -# Drops/renames config arguments the installed TRL no longer accepts, so a -# script pinned to an older TRL keeps working after an upgrade. Falls back to -# the historical raw passthrough so this can never break trainer construction. -try: - from unsloth.models.rl_config_compat import filter_config_init_kwargs as _unsloth_filter_config_init_kwargs - # A cache file generated here can be imported by an older Unsloth whose filter - # predates `mirrored_from`, so drop the argument rather than raise TypeError. - if "mirrored_from" not in inspect.signature(_unsloth_filter_config_init_kwargs).parameters: - _unsloth_filter_config_init_kwargs_old = _unsloth_filter_config_init_kwargs - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): - return _unsloth_filter_config_init_kwargs_old(config_class, kwargs) -except Exception: - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): return kwargs -def prepare_for_training_mode(f): - @functools.wraps(f) - def wrapper(self, *args, **kwargs): - # Drop any torch.compile graph cache poisoned by a stray pre-train forward. - try: - _unsloth_reset_stray_compile_cache(self) - except Exception: - pass - # Finish the previous W&B run if this is a subsequent train() call. - # We do this at the START of train() (not the end) so that - # evaluate() / log() still work after train() completes. - # HF's WandbCallback.setup() will call wandb.init() for the new run. - # See: https://github.com/unslothai/unsloth/issues/3954 - if getattr(self, '_unsloth_training_completed', False): - try: - import wandb - if wandb.run is not None: - wandb.finish() - # Reset HF's WandbCallback so it calls wandb.init() for the new run - for cb in self.callback_handler.callbacks: - if type(cb).__name__ == 'WandbCallback': - cb._initialized = False - break - except: - pass - # Enable training mode - _was_training = None - # Restore the GC mode the model was configured with at setup; fall back to - # the training args only when it wasn't recorded (issue #4735). Use hasattr, - # not a None sentinel, so a deliberately-recorded None is restored verbatim. - _model = getattr(self, 'model', None) - if hasattr(_model, '_unsloth_gradient_checkpointing'): - use_gc = _model._unsloth_gradient_checkpointing - else: - use_gc = getattr(self.args, 'gradient_checkpointing', True) - if hasattr(self, 'model') and hasattr(self.model, "training"): - _was_training = self.model.training - if hasattr(self, 'model') and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - output = f(self, *args, **kwargs) - # Restore previous mode when possible - if hasattr(self, 'model') and hasattr(self.model, "for_inference"): - if _was_training is False: - self.model.for_inference() - elif _was_training is True and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - # Reset gradient checkpointing buffers to free memory while staying ready for next run - try: - reset_unsloth_gradient_checkpointing_buffers() - except: - pass - # Mark that training completed so the next train() call can - # finish this W&B run before starting a new one - self._unsloth_training_completed = True - return output - return wrapper -pass - -torch_compile_options = { - "epilogue_fusion" : True, - "max_autotune" : False, - "shape_padding" : True, - "trace.enabled" : False, - "triton.cudagraphs" : False, -} - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_hidden_states_selective_log_softmax( - hidden_states: torch.Tensor, - lm_head: torch.Tensor, - index: torch.Tensor, - chunks: int = 4, - logit_scale_multiply: float = 0.0, - logit_scale_divide: float = 0.0, - logit_softcapping: float = 0.0, - temperature: float = 1.0, - # Rows per chunk cap. Read HERE, in the default, not in the body: the body - # is traced with fullgraph = True, and `os.environ` is an unsupported op - # there on torch 2.4 -- Dynamo raises - # torch._dynamo.exc.Unsupported: const method call bytes.decode - # from os._Environ.__getitem__, which the eager fallback does not catch - # (it only catches recompile-limit and disabled-hook breaks), so the very - # first call would die even with the variable unset. A default is evaluated - # once when the def runs, which is import time, outside any traced region. - # It is also a plain int argument, so Dynamo guards on it instead of - # constant-folding an unguarded read (2.7+ never notice a later change). - # A non-numeric value is ignored rather than raised on, so a typo cannot - # break the import. 0 keeps the previous chunk boundaries exactly. - max_rows_per_chunk: int = ( - int(os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip()) - if os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip().isdigit() - else 0 - ), - lm_head_bias: Optional[torch.Tensor] = None, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, - # Read in the default for the same Dynamo reason as max_rows_per_chunk. Pairs with vLLM's - # head_dtype=float32 that load_vllm requests under the same flag. - fp32_lm_head: bool = ( - os.environ.get("UNSLOTH_VLLM_FP32_LM_HEAD", "0").strip().lower() in ("1", "true", "on", "yes") - ), - # Older Unsloth scores old / ref log-probs here and the policy with the fused kernel: one scorer keeps ratio at 1. - fused_linear: bool = ( - os.environ.get("UNSLOTH_FUSED_LINEAR", "1") != "0" - and os.environ.get("UNSLOTH_GRPO_CHUNKED_LOGPROBS", "1") != "0" - ), -) -> torch.Tensor: - # All Unsloth Zoo code licensed under AGPL3 - from unsloth_zoo.fused_losses import chunked_logprob as _chunked_logprob - if (fused_linear and _chunked_logprob._FUSED_LINEAR and not fp32_lm_head - and support_top_k == 0 and support_top_p == 1.0 and support_min_p == 0.0 - and support_excluded_ids is None): - return _chunked_logprob.unsloth_chunked_logprobs( - hidden_states, lm_head, index, - temperature = temperature, - logit_scale_multiply = logit_scale_multiply, - logit_scale_divide = logit_scale_divide, - logit_softcapping = logit_softcapping, - lm_head_bias = lm_head_bias, - max_rows_per_chunk = max_rows_per_chunk if max_rows_per_chunk > 0 else None, - ).reshape((hidden_states.shape[0], hidden_states.shape[1])) - # Reshape on this tensor's own last dim: a no-op, so a wrong-width caller - # cannot have its row count silently rewritten and instead fails at the - # matmul below, which prints both operands. Do not swap in a bare - # torch._check: it reports only "Expected cond to be True", naming neither - # operand, and Dynamo rejects a message-carrying one. Callers dispatch on - # the width first -- see `compute_logprobs_chunk`, the packed path and - # `_pg_grad_forward`. - flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) - flat_index = index.reshape(-1) - - # Each chunk materialises rows x vocab logits and then a float32 copy of - # them, all on the device holding the output head. With a large vocabulary - # and a fixed chunk count that grows with the batch, so the peak scales with - # the batch rather than staying bounded. max_rows_per_chunk caps the rows - # per chunk instead, which is pure loop splitting: more, smaller chunks, - # same concatenated result. 0 (the default) keeps the previous chunk - # boundaries exactly. - if max_rows_per_chunk > 0: - n_rows = flat_hidden_states.shape[0] - chunks = max(chunks, -(-n_rows // max_rows_per_chunk)) - chunks = min(chunks, max(n_rows, 1)) - - # A one-row chunk's head gradient is a K=1 matmul Inductor mis-lowers under dynamic shapes (illegal - # memory access or NaN head gradients, torch 2.13), so keep chunks at 2+ rows when the head trains. - # Exactly `chunks` row ranges, so the unrolled loop length never depends on T: torch.chunk - # returns ceil(T / ceil(T / chunks)) pieces and recompiled per new count; tensor_split guards - # on T % chunks. At most rows // 2 chunks when the head trains keeps every piece at 2+ rows. - lm_head_grad = torch.is_grad_enabled() and lm_head.requires_grad - if lm_head_grad and flat_hidden_states.shape[0] < 2 * chunks: - chunks = max(flat_hidden_states.shape[0] // 2, 1) - - n_rows = flat_hidden_states.shape[0] - chunks = max(1, min(chunks, n_rows)) - bounds = [(i * n_rows) // chunks for i in range(chunks + 1)] - chunked_hidden_states = [flat_hidden_states[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - - # Cast once, outside the loop, so every chunk shares one fp32 head copy (also for autograd). - if fp32_lm_head: - lm_head = lm_head.to(torch.float32) - if lm_head_bias is not None: - lm_head_bias = lm_head_bias.to(torch.float32) - - all_per_token_logps = [] - - for chunk_hidden_states, chunk_index in zip(chunked_hidden_states, chunked_index): - # When the model is dispatched over several devices, the output head can - # sit on a different one from the hidden states, because accelerate - # places the tail of the model on the last device it fills. Co-locate on - # the head's device before the matmul, otherwise this raises - # Unhandled FakeTensor Device Propagation for aten.mm.default, - # found two different devices cuda:0, cuda:1 - # On a single device every .to() here is a no-op and the result is - # bit-identical to before. - chunk_hidden_states = chunk_hidden_states.to(device = lm_head.device, dtype = lm_head.dtype) - chunk_index = chunk_index.to(lm_head.device) - if fp32_lm_head: - # The GRPO loss runs this under bf16 / fp16 autocast, which would downcast the matmul. - with torch.autocast(device_type = chunk_hidden_states.device.type, enabled = False): - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - elif lm_head_bias is None: - chunk_logits = chunk_hidden_states @ lm_head.t() - else: - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - - if logit_scale_multiply != 0.0: - chunk_logits = chunk_logits * logit_scale_multiply - if logit_scale_divide != 0.0: - chunk_logits = chunk_logits / logit_scale_divide - if logit_softcapping != 0.0: - chunk_logits = logit_softcapping * torch.tanh(chunk_logits / logit_softcapping) - - chunk_logits = chunk_logits.to(torch.float32) - - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - per_token_logps = _truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - ) - else: - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - # Return to the caller's device so the concatenation below and every - # downstream consumer see the device they started on. - all_per_token_logps.append(per_token_logps.to(hidden_states.device)) - - all_per_token_logps = torch.concat(all_per_token_logps) - - all_per_token_logps = all_per_token_logps.reshape((hidden_states.shape[0], hidden_states.shape[1])) - return all_per_token_logps - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_selective_log_softmax( - logits, - index, - temperature: float = 1.0, - chunks: int = 4, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, -): - # Fixed `chunks` row ranges; see chunked_hidden_states_selective_log_softmax. - flat_logits = logits.reshape(-1, logits.shape[-1]) - flat_index = index.reshape(-1) - # Fewer rows than chunks: one row per range, as torch.chunk did, never empty ranges. - chunks = max(1, min(chunks, flat_logits.shape[0])) - bounds = [(i * flat_logits.shape[0]) // chunks for i in range(chunks + 1)] - chunked_logits = [flat_logits[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - all_per_token_logps = [] - # Per-chunk selective_log_softmax. - for chunk_logits, chunk_index in zip(chunked_logits, chunked_index): - chunk_logits = chunk_logits.to(torch.float32) - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - all_per_token_logps.append(_truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - )) - continue - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - all_per_token_logps.append(per_token_logps) - pass - all_per_token_logps = torch.concat(all_per_token_logps) - all_per_token_logps = all_per_token_logps.reshape((logits.shape[0], logits.shape[1])) - return all_per_token_logps - -def calculate_pad_tokens_in_prompt( - input_ids: torch.Tensor, - logits_to_keep: int, - pad_token_id: int -) -> torch.Tensor: - """Count left-padded tokens per sequence, e.g. [pad, pad, pad, cat] -> 3.""" - if logits_to_keep >= input_ids.shape[1]: - raise ValueError("logits_to_keep must be smaller than the sequence length.") - - prompt_section = input_ids[:, :-logits_to_keep] - - padding_mask = (prompt_section == pad_token_id) - - pad_token_counts = padding_mask.sum(dim=1) - - return pad_token_counts - -def create_completion_attention_mask( - completion_input_ids: torch.Tensor, - left_pad_tokens_per_prompt: torch.Tensor, - max_left_pad: int, - pad_token_id: int -) -> torch.Tensor: - """Build a completion mask that zeros leading prompt and trailing pad tokens. - - For [p,p,p,c,c,c,pad,pad,pad] (p=sliced prompt, c=completion, pad=padding) - this returns [0,0,0,1,1,1,0,0,0]. - """ - batch_size, completion_len = completion_input_ids.shape - device = completion_input_ids.device - - num_tokens_to_mask = max_left_pad - left_pad_tokens_per_prompt - - indices = torch.arange(completion_len, device=device).unsqueeze(0) - shift_mask = indices >= num_tokens_to_mask.unsqueeze(1) - - non_padding_mask = (completion_input_ids != pad_token_id) - - final_mask = shift_mask & non_padding_mask - - return final_mask - -def left_pack_padding(tensor: torch.Tensor, pad_id: int) -> torch.Tensor: - """Move all padding tokens in each sequence to the right.""" - mask = (tensor != pad_id) - # stable=True since the binary mask is unordered. - sorted_indices = torch.argsort(mask, dim=1, descending=True, stable=True) - packed_tensor = torch.gather(tensor, 1, sorted_indices) - return packed_tensor - -def align_logprobs_with_mask( - logprob_tensor: torch.Tensor, - attention_mask: torch.Tensor, - pad_value: float = 0.0 -) -> torch.Tensor: - """Align a log probability tensor with a given attention mask.""" - - device = logprob_tensor.device - batch_size, logprob_seq_len = logprob_tensor.shape - mask_seq_len = attention_mask.shape[1] - - padded_logprobs = torch.full( - attention_mask.shape, - fill_value=pad_value, - dtype=logprob_tensor.dtype, - device=device - ) - - left_pad_counts = torch.argmax(attention_mask, dim=1) - - cols = torch.arange(logprob_seq_len, device=device) - dest_indices = left_pad_counts.unsqueeze(1) + cols - - # Destination row indices, shape [batch_size, logprob_seq_len]. - row_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand_as(dest_indices) - - # Keep only in-bounds destinations, then scatter via advanced indexing. - valid_mask = dest_indices < mask_seq_len - valid_rows = row_indices[valid_mask] - valid_cols = dest_indices[valid_mask] - valid_vals = logprob_tensor[valid_mask] - padded_logprobs[valid_rows, valid_cols] = valid_vals - - return padded_logprobs - -def align_completion_tool_mask( - tool_mask: torch.Tensor, - completion_mask: torch.Tensor, -) -> torch.Tensor: - """Align a raw completion-length tool/env mask with Unsloth's repacked loss mask.""" - if tool_mask is None: - return completion_mask - if tool_mask.shape[0] != completion_mask.shape[0]: - raise ValueError("tool_mask batch size must match completion_mask batch size.") - - tool_mask = tool_mask.to(device=completion_mask.device) - if tool_mask.shape == completion_mask.shape: - aligned_tool_mask = tool_mask - else: - aligned_tool_mask = align_logprobs_with_mask( - tool_mask, - completion_mask, - pad_value=0, - ) - return completion_mask * aligned_tool_mask.to(dtype=completion_mask.dtype) - -def autotune_batch_and_chunks( - total_input_rows, - seq_len, - hidden_size, - vocab_size, - dtype_bytes=16, - multiplier=None -): - if multiplier is None: - final_m = max(4, seq_len // 4096) - else: - final_m = multiplier - - if torch.cuda.is_available(): - free_bytes, _ = torch.cuda.mem_get_info() - limit_gb = (free_bytes / (1024**3))*.80 - elif hasattr(torch, "xpu") and torch.xpu.is_available(): - # XPU: estimate free memory as total - reserved. - total_mem = torch.xpu.get_device_properties(0).total_memory - reserved_mem = torch.xpu.memory_reserved() - free_bytes = total_mem - reserved_mem - limit_gb = (free_bytes / (1024**3)) * 0.80 - else: - # Fallback: assume 8GB available. - limit_gb = 8.0 - - bytes_to_gb = 1024**3 - - b_vals = torch.arange(total_input_rows, 0, -1, device='cpu', dtype=torch.float32) - - hidden_gb = (b_vals * seq_len * hidden_size * dtype_bytes) / bytes_to_gb - - base_logits = ((b_vals/total_input_rows) * b_vals * seq_len * vocab_size * dtype_bytes) / bytes_to_gb - logits_gb = base_logits / final_m - - total_mem_gb = hidden_gb + logits_gb - - valid_mask = total_mem_gb <= limit_gb - valid_indices = torch.nonzero(valid_mask, as_tuple=False) - - if valid_indices.shape[0] == 0: - #This means your GPU will OOM - # Capped at the row count: unsloth's no-grad pass divides rows by this without max(1, ...). - return max(1, min(4, total_input_rows)), final_m - - best_idx = valid_indices[0].item() - final_b = int(b_vals[best_idx].item()) - - return final_b, final_m - -def sanitize_logprob(logprob): - """Local port of trl.scripts.vllm_serve.sanitize_logprob. - Filters NaN logprobs from vLLM outputs.""" - value = logprob.logprob - if math.isnan(value): - logging.getLogger(__name__).warning( - f"Generated NaN logprob, token logprob '{logprob}' will be ignored" - ) - return None - return value -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_text_configs(config): - """``config`` and its text sub-config, the two places a transform can live: Gemma-4-style configs keep it on ``config.text_config``, T5Gemma-style ones only reach it via ``config.get_text_config()``.""" - if config is None: - return [] - holders = [config] - 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: - holders.append(text_cfg) - return holders - -def _unsloth_get_final_logit_softcapping(model): - """The soft cap the loss applies, under any of its three spellings: ``final_logit_softcapping`` (Gemma), ``logits_soft_cap`` (RecurrentGemma), ``output_logit_soft_cap`` (xLSTM). Returns 0 if unset.""" - config = _unsloth_get_model_config(model) - if config is None: - return 0 - for holder in _unsloth_text_configs(config): - for name in ("final_logit_softcapping", "logits_soft_cap", "output_logit_soft_cap"): - softcap = getattr(holder, name, None) - if softcap: - return softcap - return 0 - -def _unsloth_resolve_logit_scales(model_config): - """``(multiply, divide)`` for the logits, read with the same field table and per-family overrides ``detect_logit_transforms`` uses. - - Only reached when the installed unsloth_zoo predates that helper. Must stay in step with ``resolve_logit_transforms`` in unsloth/models/llama.py: if the forward applies a transform GRPO does not, the policy log-probabilities come from different logits than the ones generated, which silently shifts every importance ratio. - """ - # Granite divides, HyperCLOVA X multiplies, MiniCPM3 scales hidden states (not a logit transform). - overrides = { - ("logits_scaling", "hyperclovax"): "multiply", - ("logits_scaling", "minicpm3"): None, - } - found = {"multiply": 0, "divide": 0} - for holder in _unsloth_text_configs(model_config): - model_type = getattr(holder, "model_type", "") or "" - for bucket, names in ( - ("multiply", ("logit_scale", "lm_head_multiplier", "output_multiplier")), - ("divide", ("logits_scaling",)), - ): - for name in names: - target = overrides.get((name, model_type), bucket) - if target is None or found[target]: - continue - value = getattr(holder, name, 0) or 0 - if value: - found[target] = value - break - return found["multiply"], found["divide"] - -def _unsloth_grpo_returns_hidden_states(model, tensor, lm_head): - """Does ``tensor`` (a forward's ``.logits``) carry hidden states or real logits? - - ``_get_per_token_logps_and_entropies`` sets ``UNSLOTH_RETURN_HIDDEN_STATES=1``, but only a forward that honours the name hands hidden states back as ``.logits``; any other forward returns a real ``[.., vocab]`` tensor that must not reach the ``lm_head`` matmul. - - The primary test is an explicit signal that the forward honours the flag, both set outside this file: ``__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__`` on the generated class, written by ``unsloth_zoo.compiler.create_standalone_class`` exactly when ``apply_fused_lm_head`` gave that forward its own branch; and ``_unsloth_grpo_hidden_states_forward_wrapped``, set by ``_install_grpo_hidden_states_forward_wrapper`` in ``unsloth/models/rl.py`` for models the compiler did not rewrite. That wrapper degrades to real logits when the model cannot produce hidden states and records whether it did so in ``_unsloth_grpo_hidden_states_degraded`` before returning, so reading the pair after a forward describes the call that just finished; degradation is per call, not per model. - - The width comparison stays as the fallback, for an ``unsloth_zoo`` old enough that it never writes the marker. It is decisive whenever ``vocab_size != hidden_size``, and the signal may only overrule it when it is not, which is the one case the shape cannot answer. - """ - if tensor.shape[-1] != lm_head.shape[1]: - return False # vocab-wide: real logits, whatever any signal claims - if lm_head.shape[0] != lm_head.shape[1]: - return True # hidden-wide and vocab_size != hidden_size: hidden states - return _unsloth_grpo_hidden_states_signal(model) is not False - -def _unsloth_grpo_hidden_states_signal(model): - """``True``/``False`` if the forward honours ``UNSLOTH_RETURN_HIDDEN_STATES``, ``None`` when neither marker is present. See ``_unsloth_grpo_returns_hidden_states`` for where each marker is set. Walks the wrapper chain because the markers are set on whichever object the trainer saw, which may be the DDP module or the PEFT base model rather than the object handed to the logprob loop.""" - candidates = [] - pending = [model] - while pending and len(candidates) < 8: - candidate = pending.pop(0) - if candidate is None or any(candidate is seen for seen in candidates): - continue - candidates.append(candidate) - get_base_model = getattr(candidate, "get_base_model", None) - if callable(get_base_model): - try: - pending.append(get_base_model()) - except Exception: - pass - for _attr in ("module", "base_model", "model"): - child = getattr(candidate, _attr, None) - if child is not None and hasattr(child, "forward"): - pending.append(child) - for candidate in candidates: - if getattr(candidate, "__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__", False): - return True - if getattr(type(candidate), "__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__", False): - return True - if any( - getattr(candidate, "_unsloth_grpo_hidden_states_forward_wrapped", False) - for candidate in candidates - ): - # The wrapper honours the flag unless it recorded that this call could not. - if any( - hasattr(candidate, "_unsloth_grpo_hidden_states_degraded") for candidate in candidates - ): - return not any( - getattr(candidate, "_unsloth_grpo_hidden_states_degraded", False) - for candidate in candidates - ) - # An unsloth/models/rl.py predating the per-call attribute set only the warn-once flag; it is the best signal such a wrapper offers. - return not any( - getattr(candidate, "_unsloth_grpo_hidden_states_warning_issued", False) - for candidate in candidates - ) - return None - -def _unsloth_pref_returns_hidden_states(model, tensor, lm_head): - """``_unsloth_grpo_returns_hidden_states`` for the preference routes, which need a zoo that writes the hidden-state - markers: a square head with no signal is a forward that ignored the flag, so its output is real logits.""" - if tensor.shape[-1] != lm_head.shape[1]: - return False - if lm_head.shape[0] != lm_head.shape[1]: - return True - return _unsloth_grpo_hidden_states_signal(model) is True - -def _unsloth_gkd_dense_head(model): - """The output head when it is a plain dense ``[vocab, hidden]`` projection the chunked loss can read directly. - Rejects bnb heads (``nn.Linear`` subclasses with packed weights), PEFT heads and ZeRO-3 partitioned weights. - """ - get_output_embeddings = getattr(model, "get_output_embeddings", None) - if not callable(get_output_embeddings): - return None - try: - head = get_output_embeddings() - except Exception: - return None - if type(head) is not torch.nn.Linear: - return None - # An lm_head that is not the output embeddings runs before them (ModernBERT decoder, RoBERTa-style heads). - lm_head = getattr(model, "lm_head", None) - if isinstance(lm_head, torch.nn.Module) and lm_head is not head: - return None - weight = getattr(head, "weight", None) - if not isinstance(weight, torch.Tensor) or weight.dim() != 2 or weight.numel() == 0: - return None - if weight.dtype not in (torch.float16, torch.bfloat16, torch.float32, torch.float64): - return None - if type(weight) not in (torch.Tensor, torch.nn.Parameter): - return None - bias = getattr(head, "bias", None) - if bias is not None and (bias.dim() != 1 or bias.shape[0] != weight.shape[0]): - return None - if weight.is_meta or (bias is not None and bias.is_meta): - return None - return head - -def _unsloth_stock_logps_hook( - self, - method, - hooks = ("get_batch_logps",), -): - """False when a subclass overrides one of ``hooks`` below the class whose ``method`` is the rewritten one: TRL's - body would call the override (or it reads the full logits the chunked path replaces with stand-ins).""" - owner = None - for klass in type(self).__mro__: - # The rewritten body calls a chunked or fused-pass helper; a subclass method wrapping it via super() does not. - names = getattr(getattr(vars(klass).get(method), "__code__", None), "co_names", ()) - if any( - name.startswith("_unsloth_") and ("chunked" in name or "fused_pass" in name) - for name in names - ): - owner = klass - break - if owner is None: - return True - for name in hooks: - hook = next((klass for klass in type(self).__mro__ if name in vars(klass)), None) - if hook is not None and hook not in owner.__mro__: - return False - return True - -def _unsloth_pref_chunked_head(self, model): - """The dense output head when this concatenated_forward can score from hidden states, else ``None``.""" - if unsloth_chunked_logprobs is None or os.environ.get("UNSLOTH_FUSED_LINEAR", "1") == "0": - return None - if os.environ.get("UNSLOTH_PREFERENCE_CHUNKED", "1").lower() in ("0", "false", "no", "off"): - return None - if getattr(self, "is_encoder_decoder", False) or getattr(self, "aux_loss_enabled", False): - return None - # Listing the method itself rejects a subclass wrapping it: it would see stand-in logits. - if not _unsloth_stock_logps_hook( - self, - "concatenated_forward", - ("get_batch_logps", "get_batch_loss_metrics", "concatenated_forward"), - ): - return None - # TRL <= 0.25 scores labels against label_pad_token_id but nn.CrossEntropyLoss always ignores -100. - if getattr(self, "label_pad_token_id", -100) != -100: - return None - if getattr(self, "is_fsdp_enabled", False) or getattr(self, "is_deepspeed_enabled", False): - return None - try: - unwrapped = self.accelerator.unwrap_model(model) - except Exception: - return None - head = _unsloth_gkd_dense_head(unwrapped) - if head is None: - return None - # The kernel reads weight / bias directly: a hooked head (accelerate, forward / backward hooks) is not - # F.linear(h, W, b), and its backward hooks would never see the gradient. - if ( - hasattr(head, "_hf_hook") - or "forward" in vars(head) - or head._forward_hooks - or head._forward_pre_hooks - or head._backward_hooks - or getattr(head, "_backward_pre_hooks", None) - ): - return None - # MiniCPM3 divides hidden states by logits_scaling before the head. - for config in _unsloth_text_configs(_unsloth_get_model_config(unwrapped)): - if getattr(config, "model_type", None) == "minicpm3": - return None - # DDP(find_unused_parameters=True) marks a head skipped in forward as unused, then its grad hook fires twice. - if getattr(model, "find_unused_parameters", False) and any( - p is not None and p.requires_grad for p in (head.weight, head.bias) - ): - return None - return head - -def _unsloth_gkd_logit_transforms(model): - """``(scale, softcap)`` exactly as the model's own forward applies them to its logits.""" - model_config = _unsloth_get_model_config(model) - if detect_logit_transforms is not None: - _transforms = detect_logit_transforms(model_config) - logit_softcapping = _transforms["logit_softcapping"] - logit_scale_multiply = _transforms["logit_scale_multiply"] - logit_scale_divide = _transforms["logit_scale_divide"] - else: - logit_softcapping = _unsloth_get_final_logit_softcapping(model) - logit_scale_multiply, logit_scale_divide = _unsloth_resolve_logit_scales(model_config) - scale = 1.0 - if logit_scale_multiply: - scale = scale * float(logit_scale_multiply) - if logit_scale_divide: - scale = scale / float(logit_scale_divide) - return scale, float(logit_softcapping or 0.0) - -def _unsloth_pref_logit_transforms(model): - """``(multiply, divide, softcap)`` for the chunked kernel, from the GKD resolver (one scale, 0.0 = off).""" - scale, softcap = _unsloth_gkd_logit_transforms(model) - return (0.0 if scale == 1.0 else scale), 0.0, softcap - -def _unsloth_pref_chunked_forward(self, model, head, concatenated_batch, len_chosen, states, kind): - """``(outputs, all_logits)``: ``outputs`` is TRL's return tuple computed from hidden states, or ``None`` with - ``all_logits`` holding real logits for TRL's own body (the forward ignored the flag, or a shape TRL rejects).""" - unwrapped = self.accelerator.unwrap_model(model) - multiply, divide, softcap = _unsloth_pref_logit_transforms(unwrapped) - if not _unsloth_pref_returns_hidden_states(unwrapped, states, head.weight): - return None, states - input_ids = concatenated_batch["concatenated_input_ids"] - labels = concatenated_batch["concatenated_labels"] - if states.dim() != 3 or states.shape[:2] != labels.shape or labels.shape != input_ids.shape: - # TRL raises or reads these itself: hand it the logits it expects. - logits = torch.nn.functional.linear( - states.to(device = head.weight.device, dtype = head.weight.dtype), head.weight, head.bias - ) - if multiply: - logits = logits * multiply - if divide: - logits = logits / divide - if softcap: - logits = softcap * torch.tanh(logits / softcap) - return None, logits - - device = states.device - labels = labels.to(device) - input_ids = input_ids.to(device) - shifted = labels[:, 1:] - label_mask = shifted != -100 - # cat, not boolean slice assignment: CPU Inductor fails to vectorise that (VecMask vs VecMask). - rejected_off = torch.zeros_like(label_mask[len_chosen:]) - if kind == "orpo": - # ORPO's NLL scores every attended token of prompt + chosen, from input_ids. - nll_targets = input_ids[:, 1:] - attention_mask = concatenated_batch["concatenated_attention_mask"].to(device) - nll_mask = torch.cat((attention_mask[:len_chosen, 1:] == 1, rejected_off), dim = 0) - average_log_prob = True - # Checking for a second NLL pass needs a host sync, so compiled graphs always take it (correct either way). - is_compiling = getattr(getattr(torch, "compiler", None), "is_compiling", None) - if is_compiling is not None and is_compiling(): - split_nll = True - else: - split_nll = bool((label_mask & nll_mask & (shifted != nll_targets)).any()) - else: - nll_targets = shifted - if self.cpo_alpha != 0: - nll_mask = torch.cat((label_mask[:len_chosen], rejected_off), dim = 0) - else: - nll_mask = torch.zeros_like(label_mask) - average_log_prob = self.loss_type in ["ipo", "simpo"] - split_nll = False - targets = torch.where(label_mask, shifted, torch.where(nll_mask, nll_targets, 0)) - - # CPO's logits metric averages every position, ORPO's all but the last. - rows = states if kind == "cpo" else states[:, :-1] - kwargs = dict( - logit_scale_multiply = multiply, - logit_scale_divide = divide, - logit_softcapping = softcap, - lm_head_bias = head.bias, - ) - if _unsloth_pref_masked_logprobs is not None and _unsloth_linear_mean_logits is not None: - # Prompt and padding rows (often most of each row) never reach the head; they only feed the mean metric. - per_token = _unsloth_pref_masked_logprobs( - states[:, :-1], - head.weight, - targets, - mask = label_mask if split_nll else label_mask | nll_mask, - **kwargs, - ) - mean_logits = _unsloth_linear_mean_logits( - rows, head.weight, head.bias, multiply, divide, softcap - ) - if split_nll: - nll_per_token = _unsloth_pref_masked_logprobs( - states[:len_chosen, :-1], - head.weight, - nll_targets[:len_chosen], - mask = nll_mask[:len_chosen], - **kwargs, - ) - else: - nll_per_token = per_token[:len_chosen] - else: - if kind == "cpo": - targets = torch.nn.functional.pad(targets, (0, 1)) - per_token, mean_logits = unsloth_chunked_logprobs( - rows, head.weight, targets, return_mean_logits = True, **kwargs - ) - if kind == "cpo": - per_token = per_token[:, :-1] - if split_nll: - nll_per_token = unsloth_chunked_logprobs( - states[:len_chosen, :-1], - head.weight, - torch.where(nll_mask, nll_targets, 0)[:len_chosen], - **kwargs, - ) - else: - nll_per_token = per_token[:len_chosen] - - loss_mask = label_mask.to(per_token.dtype) - all_logps = (per_token * loss_mask).sum(-1) - if average_log_prob: - all_logps = all_logps / loss_mask.sum(-1) - chosen_nll_mask = nll_mask[:len_chosen].to(per_token.dtype) - if kind == "cpo" and self.cpo_alpha == 0: - nll_loss = torch.tensor(0.0).to(self.accelerator.device) - else: - nll_loss = -(nll_per_token * chosen_nll_mask).sum() / chosen_nll_mask.sum() - try: - self._unsloth_pref_chunked_calls = getattr(self, "_unsloth_pref_chunked_calls", 0) + 1 - except Exception: - pass - # Stand-ins: TRL only reads chosen/rejected logits as .detach().mean(). - return ( - all_logps[:len_chosen], - all_logps[len_chosen:], - mean_logits[:len_chosen], - mean_logits[len_chosen:], - nll_loss, - ), states - -try: - from unsloth_zoo.device_map_planner import detect_logit_transforms -except Exception: - detect_logit_transforms = None -try: - from unsloth_zoo.fused_losses.chunked_logprob import unsloth_chunked_logprobs -except Exception: - unsloth_chunked_logprobs = None -try: - from unsloth_zoo.fused_losses.chunked_logprob import unsloth_masked_logprobs as _unsloth_pref_masked_logprobs - from unsloth_zoo.fused_losses.preference_loss import linear_mean_logits as _unsloth_linear_mean_logits -except Exception: - _unsloth_pref_masked_logprobs = _unsloth_linear_mean_logits = None -@dataclass -class UnslothORPOConfig(ORPOConfig): - """ - - Configuration class for the [`experimental.orpo.ORPOTrainer`]. - - This class includes only the parameters that are specific to ORPO training. For a full list of training arguments, - please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may - differ from those in [`~transformers.TrainingArguments`]. - - Using [`~transformers.HfArgumentParser`] we can turn this class into - [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the - command line. - - Parameters: - max_length (`int` or `None`, *optional*, defaults to `1024`): - Maximum length of the sequences (prompt + completion) in the batch. This argument is required if you want - to use the default data collator. - max_completion_length (`int`, *optional*): - Maximum length of the completion. This argument is required if you want to use the default data collator - and your model is an encoder-decoder. - beta (`float`, *optional*, defaults to `0.1`): - Parameter controlling the relative ratio loss weight in the ORPO loss. In the - [paper](https://huggingface.co/papers/2403.07691), it is denoted by λ. In the - [code](https://github.com/xfactlab/orpo), it is denoted by `alpha`. - disable_dropout (`bool`, *optional*, defaults to `True`): - Whether to disable dropout in the model. - padding_value (`int`, *optional*): - Padding value to use. If `None`, the padding value of the tokenizer is used. - generate_during_eval (`bool`, *optional*, defaults to `False`): - If `True`, generates and logs completions from the model to W&B or Comet during evaluation. - is_encoder_decoder (`bool`, *optional*): - When using the `model_init` argument (callable) to instantiate the model instead of the `model` argument, - you need to specify if the model returned by the callable is an encoder-decoder model. - model_init_kwargs (`dict[str, Any]`, *optional*): - Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model from a - string. The `revision` value is also used when loading the processing class. - trust_remote_code (`bool`, *optional*, defaults to `False`): - Whether to allow loading models that ship custom Python code from the Hub. Forwarded to - [`~transformers.AutoModelForCausalLM.from_pretrained`]. - dataset_num_proc (`int`, *optional*): - Number of processes to use for processing the dataset. - - > [!NOTE] - > These parameters have default values different from [`~transformers.TrainingArguments`]: - > - `logging_steps`: Defaults to `10` instead of `500`. - > - `gradient_checkpointing`: Defaults to `True` instead of `False`. - > - `bf16`: Defaults to `True` if `fp16` is not set, instead of `False`. - > - `learning_rate`: Defaults to `1e-6` instead of `5e-5`. - - """ - _unsloth_accepts_adamw_flash = True - vllm_sampling_params: Optional[Any] = field( - default = None, - metadata = {'help': 'vLLM SamplingParams'}, - ) - unsloth_num_chunks : Optional[int] = field( - default = -1, - metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'}, - ) - unsloth_logit_chunk_multiplier : Optional[int] = field( - default = None, - metadata = {'help': 'Multiplier for chunked logit computations.'}, - ) - unsloth_grpo_mini_batch : Optional[int] = field( - default = None, - metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'}, - ) - max_seq_length : Optional[int] = field( - default = None, - metadata = {'help': 'Maximum sequence length to truncate to.'}, - ) - def __init__( - self, - output_dir = None, - per_device_train_batch_size = 4, - num_train_epochs = 3.0, - max_steps = -1, - learning_rate = 5e-05, - lr_scheduler_type = 'linear', - lr_scheduler_kwargs = None, - warmup_steps = 0.1, - optim = 'adamw_8bit', - optim_args = None, - weight_decay = 0.001, - adam_beta1 = 0.9, - adam_beta2 = 0.999, - adam_epsilon = 1e-08, - optim_target_modules = None, - gradient_accumulation_steps = 2, - average_tokens_across_devices = True, - max_grad_norm = 1.0, - label_smoothing_factor = 0.0, - bf16 = False, - fp16 = False, - bf16_full_eval = False, - fp16_full_eval = False, - tf32 = None, - gradient_checkpointing = True, - gradient_checkpointing_kwargs = None, - torch_compile = False, - torch_compile_backend = None, - torch_compile_mode = None, - use_liger_kernel = False, - liger_kernel_config = None, - use_cache = False, - neftune_noise_alpha = None, - torch_empty_cache_steps = 250, - auto_find_batch_size = False, - logging_strategy = 'steps', - logging_steps = 1, - logging_first_step = False, - log_on_each_node = True, - logging_nan_inf_filter = False, - include_num_input_tokens_seen = False, - log_level = 'passive', - log_level_replica = 'warning', - disable_tqdm = None, - report_to = 'none', - run_name = None, - project = 'huggingface', - trackio_space_id = 'trackio', - eval_strategy = 'no', - eval_steps = None, - eval_delay = 0, - per_device_eval_batch_size = 4, - prediction_loss_only = False, - eval_on_start = False, - eval_do_concat_batches = True, - eval_use_gather_object = False, - eval_accumulation_steps = 2, - batch_eval_metrics = False, - save_only_model = False, - save_strategy = 'steps', - save_steps = 500, - save_on_each_node = False, - save_total_limit = None, - enable_jit_checkpoint = False, - push_to_hub = False, - hub_token = None, - hub_private_repo = None, - hub_model_id = None, - hub_strategy = 'every_save', - hub_always_push = False, - hub_revision = None, - load_best_model_at_end = False, - metric_for_best_model = None, - greater_is_better = None, - ignore_data_skip = False, - restore_callback_states_from_checkpoint = False, - full_determinism = False, - seed = 3407, - data_seed = None, - use_cpu = False, - accelerator_config = None, - parallelism_config = None, - dataloader_drop_last = False, - dataloader_num_workers = 0, - dataloader_pin_memory = True, - dataloader_persistent_workers = False, - dataloader_prefetch_factor = None, - remove_unused_columns = True, - label_names = None, - train_sampling_strategy = 'random', - length_column_name = 'length', - ddp_find_unused_parameters = None, - ddp_bucket_cap_mb = None, - ddp_broadcast_buffers = None, - ddp_backend = None, - ddp_timeout = 1800, - fsdp = None, - fsdp_config = None, - deepspeed = None, - debug = '', - skip_memory_metrics = True, - do_train = False, - do_eval = False, - do_predict = False, - resume_from_checkpoint = None, - warmup_ratio = None, - logging_dir = None, - local_rank = -1, - max_length = 1024, - max_completion_length = None, - beta = 0.1, - disable_dropout = True, - padding_value = None, - generate_during_eval = False, - is_encoder_decoder = None, - model_init_kwargs = None, - trust_remote_code = False, - dataset_num_proc = None, - vllm_sampling_params = None, - unsloth_num_chunks = -1, - unsloth_logit_chunk_multiplier = None, - unsloth_grpo_mini_batch = None, - max_seq_length = None, - **kwargs, - ): - if learning_rate < 1e-7: print(f'Unsloth: Your learning rate of `{learning_rate}` is too small and less than 1e-7! Consider increasing it, otherwise gradient updates will be close to 0!') - if learning_rate > 1: print(f'Unsloth: Your learning rate of `{learning_rate}` is way too larger > 1! Consider decreasing it to 1e-1, otherwise gradient updates will explode!') - _unsloth_fsdp2 = os.environ.get('FSDP_VERSION', '').strip() == '2' - _unsloth_fsdp_cfg = fsdp_config - if fsdp and isinstance(_unsloth_fsdp_cfg, str) and os.path.isfile(_unsloth_fsdp_cfg): - try: - import json as _unsloth_json - with open(_unsloth_fsdp_cfg, encoding = 'utf-8') as _unsloth_f: - _unsloth_fsdp_cfg = _unsloth_json.load(_unsloth_f) - except Exception: - _unsloth_fsdp_cfg = None - if fsdp and _unsloth_fsdp_cfg is None: - _unsloth_fsdp_cfg = {} - if fsdp and isinstance(_unsloth_fsdp_cfg, dict): - _unsloth_fsdp2 = _unsloth_fsdp2 or str(_unsloth_fsdp_cfg.get('fsdp_version', _unsloth_fsdp_cfg.get('version', 1))) == '2' - _unsloth_flash_adamw = str(getattr(optim, 'value', optim)).strip().lower() == 'adamw_flash' - if _unsloth_flash_adamw: optim = 'adamw_torch' - if _unsloth_fsdp2 and str(getattr(optim, 'value', optim)) == 'adamw_8bit': - print("Unsloth: FSDP2 shards parameters as DTensors, which bitsandbytes' 8-bit optimizers cannot step. Switching optim = 'adamw_8bit' to 'adamw_torch_fused'.") - optim = 'adamw_torch_fused' - if num_train_epochs is None: - num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override - if output_dir is None and save_strategy == 'steps' and save_steps == 500: - output_dir = 'unsloth_training_checkpoints' - save_strategy = 'no' - try: - from unsloth_zoo.dataset_num_proc import get_dataset_num_proc as _unsloth_get_dataset_num_proc - except Exception: - try: - from unsloth.dataset_num_proc import get_dataset_num_proc as _unsloth_get_dataset_num_proc - except Exception: - _unsloth_get_dataset_num_proc = None - if _unsloth_get_dataset_num_proc is not None: - dataset_num_proc = _unsloth_get_dataset_num_proc(dataset_num_proc, serial_as_none = True) - if eval_steps is not None and eval_strategy != 'steps': - print(f'Unsloth: `eval_steps = {eval_steps}` is ignored because `eval_strategy` is {getattr(eval_strategy, "value", eval_strategy)!r}. Set `eval_strategy = "steps"` to evaluate every `eval_steps` steps.') - - - # One dict so the filter sees the mirrored parameters AND `**kwargs`: - # filtering kwargs alone would double-bind any argument TRL renamed, - # since the new name is itself a mirrored parameter. - _unsloth_config_arguments = dict( - output_dir = output_dir, - per_device_train_batch_size = per_device_train_batch_size, - num_train_epochs = num_train_epochs, - max_steps = max_steps, - learning_rate = learning_rate, - lr_scheduler_type = lr_scheduler_type, - lr_scheduler_kwargs = lr_scheduler_kwargs, - warmup_steps = warmup_steps, - optim = optim, - optim_args = optim_args, - weight_decay = weight_decay, - adam_beta1 = adam_beta1, - adam_beta2 = adam_beta2, - adam_epsilon = adam_epsilon, - optim_target_modules = optim_target_modules, - gradient_accumulation_steps = gradient_accumulation_steps, - average_tokens_across_devices = average_tokens_across_devices, - max_grad_norm = max_grad_norm, - label_smoothing_factor = label_smoothing_factor, - bf16 = bf16, - fp16 = fp16, - bf16_full_eval = bf16_full_eval, - fp16_full_eval = fp16_full_eval, - tf32 = tf32, - gradient_checkpointing = gradient_checkpointing, - gradient_checkpointing_kwargs = gradient_checkpointing_kwargs, - torch_compile = torch_compile, - torch_compile_backend = torch_compile_backend, - torch_compile_mode = torch_compile_mode, - use_liger_kernel = use_liger_kernel, - liger_kernel_config = liger_kernel_config, - use_cache = use_cache, - neftune_noise_alpha = neftune_noise_alpha, - torch_empty_cache_steps = torch_empty_cache_steps, - auto_find_batch_size = auto_find_batch_size, - logging_strategy = logging_strategy, - logging_steps = logging_steps, - logging_first_step = logging_first_step, - log_on_each_node = log_on_each_node, - logging_nan_inf_filter = logging_nan_inf_filter, - include_num_input_tokens_seen = include_num_input_tokens_seen, - log_level = log_level, - log_level_replica = log_level_replica, - disable_tqdm = disable_tqdm, - report_to = report_to, - run_name = run_name, - project = project, - trackio_space_id = trackio_space_id, - eval_strategy = eval_strategy, - eval_steps = eval_steps, - eval_delay = eval_delay, - per_device_eval_batch_size = per_device_eval_batch_size, - prediction_loss_only = prediction_loss_only, - eval_on_start = eval_on_start, - eval_do_concat_batches = eval_do_concat_batches, - eval_use_gather_object = eval_use_gather_object, - eval_accumulation_steps = eval_accumulation_steps, - batch_eval_metrics = batch_eval_metrics, - save_only_model = save_only_model, - save_strategy = save_strategy, - save_steps = save_steps, - save_on_each_node = save_on_each_node, - save_total_limit = save_total_limit, - enable_jit_checkpoint = enable_jit_checkpoint, - push_to_hub = push_to_hub, - hub_token = hub_token, - hub_private_repo = hub_private_repo, - hub_model_id = hub_model_id, - hub_strategy = hub_strategy, - hub_always_push = hub_always_push, - hub_revision = hub_revision, - load_best_model_at_end = load_best_model_at_end, - metric_for_best_model = metric_for_best_model, - greater_is_better = greater_is_better, - ignore_data_skip = ignore_data_skip, - restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint, - full_determinism = full_determinism, - seed = seed, - data_seed = data_seed, - use_cpu = use_cpu, - accelerator_config = accelerator_config, - parallelism_config = parallelism_config, - dataloader_drop_last = dataloader_drop_last, - dataloader_num_workers = dataloader_num_workers, - dataloader_pin_memory = dataloader_pin_memory, - dataloader_persistent_workers = dataloader_persistent_workers, - dataloader_prefetch_factor = dataloader_prefetch_factor, - remove_unused_columns = remove_unused_columns, - label_names = label_names, - train_sampling_strategy = train_sampling_strategy, - length_column_name = length_column_name, - ddp_find_unused_parameters = ddp_find_unused_parameters, - ddp_bucket_cap_mb = ddp_bucket_cap_mb, - ddp_broadcast_buffers = ddp_broadcast_buffers, - ddp_backend = ddp_backend, - ddp_timeout = ddp_timeout, - fsdp = fsdp, - fsdp_config = fsdp_config, - deepspeed = deepspeed, - debug = debug, - skip_memory_metrics = skip_memory_metrics, - do_train = do_train, - do_eval = do_eval, - do_predict = do_predict, - resume_from_checkpoint = resume_from_checkpoint, - warmup_ratio = warmup_ratio, - logging_dir = logging_dir, - local_rank = local_rank, - max_length = max_length, - max_completion_length = max_completion_length, - beta = beta, - disable_dropout = disable_dropout, - padding_value = padding_value, - generate_during_eval = generate_during_eval, - is_encoder_decoder = is_encoder_decoder, - model_init_kwargs = model_init_kwargs, - trust_remote_code = trust_remote_code, - dataset_num_proc = dataset_num_proc,**kwargs) - super().__init__(**_unsloth_filter_config_init_kwargs(ORPOConfig, _unsloth_config_arguments, mirrored_from = __class__)) - self.vllm_sampling_params = vllm_sampling_params - self.unsloth_num_chunks = unsloth_num_chunks - if unsloth_grpo_mini_batch is not None: - if self.generation_batch_size >= unsloth_grpo_mini_batch: - self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch - else: - raise ValueError( - f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, " - f"which is self.per_device_train_batch_size * gradient_accumulation_steps." - ) - self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier - self.max_seq_length = max_seq_length - # Unsloth: keep the reentrant checkpoint path - if getattr(self, 'gradient_checkpointing', False): - _gc_kwargs = getattr(self, 'gradient_checkpointing_kwargs', None) or {} - if _gc_kwargs.get('context_fn') is None and not _gc_kwargs.get('debug', False): - _gc_kwargs['use_reentrant'] = True - self.gradient_checkpointing_kwargs = _gc_kwargs - from unsloth.optimizers.flash_adamw import is_flash_adamw_marked, mark_flash_adamw, unmark_flash_adamw - if _unsloth_flash_adamw or is_flash_adamw_marked(self): - mark_flash_adamw(self) - else: - unmark_flash_adamw(self) - -pass - -class _UnslothORPOTrainer(_BaseTrainer): - r"""""" - - loss_is_scaled_for_ga = False - - _tag_names = ["trl", "orpo"] - _name = "ORPO" - _paper = { - "title": "ORPO: Monolithic Preference Optimization without Reference Model", - "id": "2403.07691", - # docstyle-ignore - "citation": textwrap.dedent("""\ - @article{hong2024orpo, - title = {{ORPO: Monolithic Preference Optimization without Reference Model}}, - author = {Jiwoo Hong and Noah Lee and James Thorne}, - year = 2024, - eprint = {arXiv:2403.07691} - }"""), - } - - def __init__( - self, - model: PreTrainedModel | nn.Module | str | None = None, - args: ORPOConfig | None = None, - data_collator: DataCollator | None = None, - train_dataset: Dataset | None = None, - eval_dataset: Dataset | dict[str, Dataset] | None = None, - processing_class: PreTrainedTokenizerBase - | BaseImageProcessor - | FeatureExtractionMixin - | ProcessorMixin - | None = None, - model_init: Callable[[], PreTrainedModel] | None = None, - callbacks: list[TrainerCallback] | None = None, - optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), - preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, - peft_config: "PeftConfig | None" = None, - compute_metrics: Callable[[EvalLoopOutput], dict] | None = None, - ): - if train_dataset is None: - raise ValueError("`train_dataset` is required") - - if args.model_init_kwargs is None: - model_init_kwargs = {} - elif not isinstance(model, str): - raise ValueError("You passed model_kwargs to the ORPOTrainer. But your model is already instantiated.") - else: - model_init_kwargs = args.model_init_kwargs - dtype = model_init_kwargs.get("dtype", "auto") - if dtype is not None: - # Convert to `torch.dtype` if an str is passed - if isinstance(dtype, str) and dtype != "auto": - dtype = getattr(torch, dtype) - if dtype != "auto" and not isinstance(dtype, torch.dtype): - raise ValueError( - f"Invalid `dtype` passed to the ORPOConfig. Expected a string with either `torch.dtype` or 'auto', but got {dtype}." - ) - model_init_kwargs["dtype"] = dtype - model_init_kwargs["device_map"] = model_init_kwargs.get("device_map", "auto") - - model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) - model_revision = model_init_kwargs.get("revision") if isinstance(model, str) else None - - # PEFT initializes the adapter weights randomly, so set_seed must be done before creating the model to ensure - # reproducibility. - set_seed(args.seed) - if isinstance(model, str): - model = AutoModelForCausalLM.from_pretrained(model, **model_init_kwargs) - - # PEFT - # Initialize this variable to False. This helps tracking the case when `peft_module_casting_to_bf16` - # has been called in order to properly call autocast if needed. - self._peft_has_been_casted_to_bf16 = False - if False: - if not is_peft_available(): - raise ImportError( - "You passed `peft_config` but the `peft` library is not installed. " - "Install it with `pip install trl[peft]`." - ) - if not isinstance(peft_config, PeftConfig): - raise TypeError( - f"`peft_config` must be a `peft.PeftConfig` instance (e.g. `peft.LoraConfig`), " - f"got {type(peft_config).__name__}." - ) - if is_peft_model(model): - raise ValueError( - "You passed a `PeftModel` instance together with a `peft_config` to the trainer. Please first " - "merge and unload the existing adapter, save the resulting base model, and then pass that base " - "model along with the new `peft_config` to the trainer." - ) - - if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False): - _support_gc_kwargs = hasattr( - args, "gradient_checkpointing_kwargs" - ) and "gradient_checkpointing_kwargs" in list( - inspect.signature(prepare_model_for_kbit_training).parameters - ) - - prepare_model_kwargs = {"use_gradient_checkpointing": args.gradient_checkpointing} - - if _support_gc_kwargs: - prepare_model_kwargs["gradient_checkpointing_kwargs"] = args.gradient_checkpointing_kwargs - - model = prepare_model_for_kbit_training(model, **prepare_model_kwargs) - elif args.gradient_checkpointing: - # For backward compatibility with older versions of transformers - if hasattr(model, "enable_input_require_grads"): - model.enable_input_require_grads() - else: - - def make_inputs_require_grad(module, input, output): - output.requires_grad_(True) - - model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) - - # get peft model with the given config - # ZeRO-3 + PEFT for non-quantized models: - # - PEFT's default autocast_adapter_dtype=True upcasts LoRA adapter params to fp32 even when the base model is bf16. - # - ZeRO-3's _allgather_params_coalesced allocates output buffers using the dtype of the first persistent parameter, - # so mixed-dtype persistent_parameters [bf16 base + fp32 LoRA] cause a TypeError on the first optimizer step. - # - Passing autocast_adapter_dtype=False keeps adapter params in the base model dtype [bf16], fixing the mismatch. - # - This is safe: the fp32 upcast is a QLoRA-specific concern [low-bit quantized base models], not needed for - # non-quantized bf16 training. - # - See: - # - TRL issue: https://github.com/huggingface/trl/issues/6089 - # - Upstream issue: https://github.com/deepspeedai/DeepSpeed/issues/8072 - _is_quantized_model = getattr(model, "is_loaded_in_4bit", False) or getattr( - model, "is_loaded_in_8bit", False - ) - get_peft_model_kwargs = {} - if args.deepspeed_plugin is not None and args.deepspeed_plugin.zero_stage == 3 and not _is_quantized_model: - get_peft_model_kwargs["autocast_adapter_dtype"] = False - model = get_peft_model(model, peft_config, **get_peft_model_kwargs) - if args.bf16 and getattr(model, "is_loaded_in_4bit", False): - peft_module_casting_to_bf16(model) - # If args.bf16 we need to explicitly call `generate` with torch amp autocast context manager - self._peft_has_been_casted_to_bf16 = True - - # For models that use gradient_checkpointing, we need to attach a hook that enables input - # to explicitly have `requires_grad=True`, otherwise training will either silently - # fail or completely fail. - elif args.gradient_checkpointing: - # For backward compatibility with older versions of transformers - if hasattr(model, "enable_input_require_grads"): - model.enable_input_require_grads() - else: - - def make_inputs_require_grad(module, input, output): - output.requires_grad_(True) - - model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) - - if args.generate_during_eval and not (is_wandb_available() or is_comet_available()): - raise ValueError( - "`generate_during_eval=True` requires Weights and Biases or Comet to be installed." - " Please install `wandb` or `comet-ml` to resolve." - ) - - if model is not None: - self.is_encoder_decoder = model.config.is_encoder_decoder - elif args.is_encoder_decoder is None: - raise ValueError("When no model is provided, you need to pass the parameter is_encoder_decoder.") - else: - self.is_encoder_decoder = args.is_encoder_decoder - - if self.is_encoder_decoder: - self.decoder_start_token_id = model.config.decoder_start_token_id - self.pad_token_id = model.config.pad_token_id - - if processing_class is None: - processing_class = AutoTokenizer.from_pretrained( - get_config_model_id(model.config), revision=model_revision, trust_remote_code=args.trust_remote_code - ) - if args.max_length is None: - logger.warning( - "`max_length` is not set in the ORPOConfig's init" - " it will default to `512` by default, but you should do it yourself in the future.", - ) - max_length = 512 - else: - max_length = args.max_length - - if args.max_completion_length is None and self.is_encoder_decoder: - logger.warning( - "When using an encoder decoder architecture, you should set `max_completion_length` in the ORPOConfig's init" - " it will default to `128` by default, but you should do it yourself in the future.", - ) - self.max_completion_length = 128 - else: - self.max_completion_length = args.max_completion_length - - if data_collator is None: - data_collator = DPODataCollatorWithPadding( - pad_token_id=(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)), - is_encoder_decoder=self.is_encoder_decoder, - ) - - if args.remove_unused_columns: - args.remove_unused_columns = False - # warn users - logger.warning( - "When using DPODataCollatorWithPadding, you should set `remove_unused_columns=False` in your TrainingArguments" - " we have set it for you, but you should do it yourself in the future.", - ) - - self.use_dpo_data_collator = True - else: - self.use_dpo_data_collator = False - - # Disable dropout in the model and reference model - if args.disable_dropout: - disable_dropout_in_model(model) - - self.max_length = max_length - self.generate_during_eval = args.generate_during_eval - self.padding_value = args.padding_value if args.padding_value is not None else (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)) - self.processing_class = processing_class - - self.beta = args.beta - self.aux_loss_enabled = getattr(model.config, "output_router_logits", False) - self.aux_loss_coef = getattr(model.config, "router_aux_loss_coef", 0.0) - if self.aux_loss_enabled and self.aux_loss_coef == 0.0: - logger.warning( - "You set `output_router_logits` to `True` in the model config, but `router_aux_loss_coef` is set to " - "`0.0`, meaning the auxiliary loss will not be used. Either set `router_aux_loss_coef` to a value " - "greater than `0.0`, or set `output_router_logits` to `False` if you don't want to use the auxiliary " - "loss.", - ) - - self._stored_metrics = defaultdict(lambda: defaultdict(list)) - - with global_then_local_main_first(): - # Extract the prompt if needed, and apply the chat template if needed - train_dataset = train_dataset.map(maybe_extract_prompt, num_proc=args.dataset_num_proc) - train_dataset = train_dataset.map( - maybe_apply_chat_template, - fn_kwargs={"processing_class": processing_class}, - num_proc=args.dataset_num_proc, - ) - train_dataset = train_dataset.map(self.tokenize_row, num_proc=args.dataset_num_proc) - if eval_dataset is not None: - eval_dataset = eval_dataset.map(maybe_extract_prompt, num_proc=args.dataset_num_proc) - eval_dataset = eval_dataset.map( - maybe_apply_chat_template, - fn_kwargs={"processing_class": processing_class}, - num_proc=args.dataset_num_proc, - ) - eval_dataset = eval_dataset.map(self.tokenize_row, num_proc=args.dataset_num_proc) - - # Transformers explicitly set use_reentrant=True in the past to silence a PyTorch warning, but the default was - # never updated once PyTorch switched to recommending use_reentrant=False. Until that change lands upstream - # [see https://github.com/huggingface/transformers/pull/43203] and is released [most likely in 5.0.0], we - # default to the recommended non-reentrant behavior here, while preserving any user-provided value. - if args.gradient_checkpointing and Version(transformers.__version__) < Version("5.0.0"): - args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} - args.gradient_checkpointing_kwargs.setdefault("use_reentrant", False) - - super().__init__( - model=model, - args=args, - data_collator=data_collator, - train_dataset=train_dataset, - eval_dataset=eval_dataset, - processing_class=processing_class, - model_init=model_init, - compute_metrics=compute_metrics, - callbacks=callbacks, - optimizers=optimizers, - preprocess_logits_for_metrics=preprocess_logits_for_metrics, - ) - - # Add tags for models that have been loaded with the correct transformers version - if hasattr(self.model, "add_model_tags"): - self.model.add_model_tags(self._tag_names) - - if not hasattr(self, "accelerator"): - raise AttributeError( - "Your `Trainer` does not have an `accelerator` object. Consider upgrading `transformers`." - ) - - def build_tokenized_answer(self, prompt, answer): - """ - Llama tokenizer does satisfy `enc(a + b) = enc(a) + enc(b)`. It does ensure `enc(a + b) = enc(a) + enc(a + - b)[len(enc(a)):]`. Reference: - https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257 - """ - - tokenizer = getattr(self.processing_class, "tokenizer", self.processing_class) - full_tokenized = tokenizer(prompt + answer, add_special_tokens=False) - prompt_input_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"] - - answer_input_ids = full_tokenized["input_ids"][len(prompt_input_ids) :] - answer_attention_mask = full_tokenized["attention_mask"][len(prompt_input_ids) :] - - # Concat tokens to form `enc(a) + enc(a + b)[len(enc(a)):]` - full_concat_input_ids = np.concatenate([prompt_input_ids, answer_input_ids]) - - # Prepare input tokens for token by token comparison - full_input_ids = np.array(full_tokenized["input_ids"]) - - if len(full_input_ids) != len(full_concat_input_ids): - raise ValueError("Prompt input ids and answer input ids should have the same length.") - - # On some tokenizers, like Llama-2 tokenizer, there are occasions where tokens - # can be merged together when tokenizing prompt+answer. This could result - # on the last token from the prompt being different when tokenized on its own - # vs when done as prompt+answer. - response_token_ids_start_idx = len(prompt_input_ids) - - # If tokenized prompt is different than both prompt+answer, then it means the - # last token has changed due to merging. - if prompt_input_ids != full_tokenized["input_ids"][:response_token_ids_start_idx]: - response_token_ids_start_idx -= 1 - - prompt_input_ids = full_tokenized["input_ids"][:response_token_ids_start_idx] - prompt_attention_mask = full_tokenized["attention_mask"][:response_token_ids_start_idx] - - if len(prompt_input_ids) != len(prompt_attention_mask): - raise ValueError("Prompt input ids and attention mask should have the same length.") - - answer_input_ids = full_tokenized["input_ids"][response_token_ids_start_idx:] - answer_attention_mask = full_tokenized["attention_mask"][response_token_ids_start_idx:] - - return dict( - prompt_input_ids=prompt_input_ids, - prompt_attention_mask=prompt_attention_mask, - input_ids=answer_input_ids, - attention_mask=answer_attention_mask, - ) - - def tokenize_row(self, feature, model: PreTrainedModel | nn.Module | None = None) -> dict: - """Tokenize a single row from a ORPO specific dataset. - - At this stage, we don't convert to PyTorch tensors yet; we just handle the truncation in case the prompt + - chosen or prompt + rejected responses is/are too long. First we truncate the prompt; if we're still too long, - we truncate the chosen/rejected. - - We also create the labels for the chosen/rejected responses, which are of length equal to the sum of the length - of the prompt and the chosen/rejected response, with `-100` for the prompt tokens. - """ - batch = {} - tokenizer = getattr(self.processing_class, "tokenizer", self.processing_class) - prompt = feature["prompt"] - chosen = feature["chosen"] - rejected = feature["rejected"] - - if not self.is_encoder_decoder: - # Check issues below for more details - # 1. https://github.com/huggingface/trl/issues/907 - # 2. https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257 - # 3. https://github.com/LianjiaTech/BELLE/issues/337 - - if not isinstance(prompt, str): - raise ValueError(f"prompt should be an str but got {type(prompt)}") - prompt_tokens = tokenizer(prompt, add_special_tokens=False) - prompt_tokens = {f"prompt_{k}": v for k, v in prompt_tokens.items()} - - if not isinstance(chosen, str): - raise ValueError(f"chosen should be an str but got {type(chosen)}") - chosen_tokens = self.build_tokenized_answer(prompt, chosen) - - if not isinstance(rejected, str): - raise ValueError(f"rejected should be an str but got {type(rejected)}") - rejected_tokens = self.build_tokenized_answer(prompt, rejected) - - # Last prompt token might get merged by tokenizer and - # it should not be included for generation if that happens - prompt_len_input_ids = len(prompt_tokens["prompt_input_ids"]) - - chosen_prompt_len_input_ids = len(chosen_tokens["prompt_input_ids"]) - rejected_prompt_len_input_ids = len(rejected_tokens["prompt_input_ids"]) - prompt_len_input_ids = min(chosen_prompt_len_input_ids, rejected_prompt_len_input_ids) - - for k, v in prompt_tokens.items(): - prompt_tokens[k] = v[:prompt_len_input_ids] - - # Make sure prompts only have one different token at most an - # and length only differs by 1 at most - num_diff_tokens = sum( - a != b - for a, b in zip(chosen_tokens["prompt_input_ids"], rejected_tokens["prompt_input_ids"], strict=False) - ) - num_diff_len = abs(chosen_prompt_len_input_ids - rejected_prompt_len_input_ids) - if num_diff_tokens > 1 or num_diff_len > 1: - raise ValueError( - "Chosen and rejected prompt_input_ids might only differ on the " - "last token due to tokenizer merge ops." - ) - - # add BOS token to head of prompt. Avoid adding if it's already there - prompt_tokens, chosen_tokens, rejected_tokens = add_bos_token_if_needed( - tokenizer.bos_token_id, - prompt_len_input_ids, - prompt_tokens, - chosen_prompt_len_input_ids, - chosen_tokens, - rejected_prompt_len_input_ids, - rejected_tokens, - ) - - # add EOS token to end of answer. Avoid adding if it's already there - chosen_tokens, rejected_tokens = add_eos_token_if_needed( - tokenizer.eos_token_id, chosen_tokens, rejected_tokens - ) - - # if combined sequence is too long, truncate the response - _unsloth_ul = max(len(chosen_tokens['input_ids']), len(rejected_tokens['input_ids'])) - _unsloth_pl = max(len(chosen_tokens['prompt_input_ids']), len(rejected_tokens['prompt_input_ids'])) - if self.max_length is not None and _unsloth_pl + _unsloth_ul > self.max_length: - _unsloth_keep = min(_unsloth_pl, max(self.max_length - _unsloth_ul, self.max_length // 2)) - _unsloth_keep_start = getattr(self, 'truncation_mode', 'keep_end') == 'keep_start' - for answer_tokens in [chosen_tokens, rejected_tokens, prompt_tokens]: - for k in ['prompt_input_ids', 'prompt_attention_mask']: - answer_tokens[k] = answer_tokens[k][:_unsloth_keep] if _unsloth_keep_start else answer_tokens[k][max(0, len(answer_tokens[k]) - _unsloth_keep):] - for answer_tokens in [chosen_tokens, rejected_tokens]: - for k in ['input_ids', 'attention_mask']: - answer_tokens[k] = answer_tokens[k][: self.max_length - _unsloth_keep] - for answer_tokens in [chosen_tokens, rejected_tokens]: - if len(answer_tokens["prompt_input_ids"]) + len(answer_tokens["input_ids"]) > self.max_length: - for k in ["input_ids", "attention_mask"]: - answer_tokens[k] = answer_tokens[k][ - : max(0, self.max_length - len(answer_tokens["prompt_input_ids"])) - ] - if len(answer_tokens["input_ids"]) == 0: - logger.warning_once( - "The prompt alone fills `max_length`, so truncation leaves an empty completion. Consider " - "increasing `max_length` or filtering out long prompts." - ) - - # Create labels - chosen_sequence_tokens = { - k: chosen_tokens[f"prompt_{k}"] + chosen_tokens[k] for k in ["input_ids", "attention_mask"] - } - rejected_sequence_tokens = { - k: rejected_tokens[f"prompt_{k}"] + rejected_tokens[k] for k in ["input_ids", "attention_mask"] - } - chosen_sequence_tokens["labels"] = chosen_sequence_tokens["input_ids"][:] - chosen_sequence_tokens["labels"][: len(chosen_tokens["prompt_input_ids"])] = [-100] * len( - chosen_tokens["prompt_input_ids"] - ) - rejected_sequence_tokens["labels"] = rejected_sequence_tokens["input_ids"][:] - rejected_sequence_tokens["labels"][: len(rejected_tokens["prompt_input_ids"])] = [-100] * len( - rejected_tokens["prompt_input_ids"] - ) - - for k, toks in { - "chosen_": chosen_sequence_tokens, - "rejected_": rejected_sequence_tokens, - "": prompt_tokens, - }.items(): - for type_key, tokens in toks.items(): - if type_key == "token_type_ids": - continue - batch[f"{k}{type_key}"] = tokens - - else: - chosen_tokens = tokenizer( - chosen, truncation=True, max_length=self.max_completion_length, add_special_tokens=True - ) - rejected_tokens = tokenizer( - rejected, truncation=True, max_length=self.max_completion_length, add_special_tokens=True - ) - prompt_tokens = tokenizer(prompt, add_special_tokens=True) - - batch["chosen_labels"] = chosen_tokens["input_ids"] - batch["rejected_labels"] = rejected_tokens["input_ids"] - batch["prompt_input_ids"] = prompt_tokens["input_ids"] - batch["prompt_attention_mask"] = prompt_tokens["attention_mask"] - - if model is not None and hasattr(model, "prepare_decoder_input_ids_from_labels"): - batch["rejected_decoder_input_ids"] = model.prepare_decoder_input_ids_from_labels( - labels=torch.tensor(batch["rejected_labels"]) - ) - batch["chosen_decoder_input_ids"] = model.prepare_decoder_input_ids_from_labels( - labels=torch.tensor(batch["chosen_labels"]) - ) - - if is_torch_xla_available(): - # Pad the sequences to global max_length to avoid TorchXLA recompilation - for k in batch: - if "labels" in k or self.is_encoder_decoder: - pad_value = -100 - elif k.endswith("_input_ids"): - pad_value = self.padding_value - elif k.endswith("_attention_mask"): - pad_value = 0 - batch[k] = batch[k] + [pad_value] * (self.max_length - len(batch[k])) - return batch - - @staticmethod - def concatenated_inputs( - batch: dict[str, list | torch.LongTensor], - is_encoder_decoder: bool = False, - padding_value: int = 0, - device: torch.device | None = None, - ) -> dict[str, torch.LongTensor]: - """Concatenate the chosen and rejected inputs into a single tensor. - - Args: - batch: - A batch of data. Must contain the keys 'chosen_input_ids' and 'rejected_input_ids', which are tensors - of shape (batch_size, sequence_length). - is_encoder_decoder: - Whether the model is an encoder-decoder model. - padding_value: - The padding value to use for the concatenated inputs_ids. - device: - The device for the concatenated inputs. - - Returns: - A dictionary containing the concatenated inputs under the key 'concatenated_input_ids'. - """ - concatenated_batch = {} - - if is_encoder_decoder: - max_length = max(batch["chosen_labels"].shape[1], batch["rejected_labels"].shape[1]) - else: - max_length = max(batch["chosen_input_ids"].shape[1], batch["rejected_input_ids"].shape[1]) - - for k in batch: - if k.startswith("chosen") and isinstance(batch[k], torch.Tensor): - if "labels" in k or is_encoder_decoder: - pad_value = -100 - elif k.endswith("_input_ids"): - pad_value = padding_value - elif k.endswith("_attention_mask"): - pad_value = 0 - concatenated_key = k.replace("chosen", "concatenated") - concatenated_batch[concatenated_key] = pad_to_length(batch[k], max_length, pad_value=pad_value) - for k in batch: - if k.startswith("rejected") and isinstance(batch[k], torch.Tensor): - if "labels" in k or is_encoder_decoder: - pad_value = -100 - elif k.endswith("_input_ids"): - pad_value = padding_value - elif k.endswith("_attention_mask"): - pad_value = 0 - concatenated_key = k.replace("rejected", "concatenated") - concatenated_batch[concatenated_key] = torch.cat( - ( - concatenated_batch[concatenated_key], - pad_to_length(batch[k], max_length, pad_value=pad_value), - ), - dim=0, - ).to(device=device) - - if is_encoder_decoder: - concatenated_batch["concatenated_input_ids"] = batch["prompt_input_ids"].repeat(2, 1).to(device=device) - concatenated_batch["concatenated_attention_mask"] = ( - batch["prompt_attention_mask"].repeat(2, 1).to(device=device) - ) - - return concatenated_batch - - def odds_ratio_loss( - self, - policy_chosen_logps: torch.FloatTensor, - policy_rejected_logps: torch.FloatTensor, - ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: - """Compute ORPO's odds ratio (OR) loss for a batch of policy and reference model log probabilities. - - Args: - policy_chosen_logps: - Log probabilities of the policy model for the chosen responses. Shape: (batch_size,) - policy_rejected_logps: - Log probabilities of the policy model for the rejected responses. Shape: (batch_size,) - - Returns: - A tuple of three tensors: (losses, chosen_rewards, rejected_rewards). The losses tensor contains the ORPO - loss for each example in the batch. The chosen_rewards and rejected_rewards tensors contain the rewards for - the chosen and rejected responses, respectively. The log odds ratio of the chosen responses over the - rejected responses ratio for logging purposes. The `log(sigmoid(log_odds_chosen))` for logging purposes. - """ - - # Derived from Eqs. (4) and (7) from https://huggingface.co/papers/2403.07691 by using log identities and exp(log(P(y|x)) = P(y|x) - policy_chosen_logps = policy_chosen_logps.float() - policy_rejected_logps = policy_rejected_logps.float() - log_odds = (policy_chosen_logps - policy_rejected_logps) - ( - log1mexp(policy_chosen_logps) - log1mexp(policy_rejected_logps) - ) - ratio = F.logsigmoid(log_odds) - losses = self.beta * ratio - - chosen_rewards = self.beta * (policy_chosen_logps.to(self.accelerator.device)).detach() - rejected_rewards = self.beta * (policy_rejected_logps.to(self.accelerator.device)).detach() - - return losses, chosen_rewards, rejected_rewards, torch.mean(ratio), torch.mean(log_odds) - - @staticmethod - def get_batch_logps( - logits: torch.FloatTensor, - labels: torch.LongTensor, - average_log_prob: bool = False, - is_encoder_decoder: bool = False, - ) -> torch.FloatTensor: - """Compute the log probabilities of the given labels under the given logits. - - Args: - logits: Logits of the model (unnormalized). Shape: (batch_size, sequence_length, vocab_size) - labels: - Labels for which to compute the log probabilities. Label tokens with a value of `-100` are ignored. - Shape: (batch_size, sequence_length) - average_log_prob: - If True, return the average log probability per (non-masked) token. Otherwise, return the sum of the - log probabilities of the (non-masked) tokens. - is_encoder_decoder: Whether the model is an encoder-decoder model. - - Returns: - A tensor of shape (batch_size,) containing the average/sum log probabilities of the given labels under the - given logits. - """ - if logits.shape[:-1] != labels.shape: - raise ValueError("Logits (batch and sequence length dim) and labels must have the same shape.") - - if not is_encoder_decoder: - labels = labels[:, 1:].clone() - logits = logits[:, :-1, :] - loss_mask = labels != -100 - - # dummy token; we'll ignore the losses on these tokens later - labels = torch.where(labels == -100, 0, labels) - - per_token_logps = selective_log_softmax(logits, labels) - - if average_log_prob: - return (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1) - else: - return (per_token_logps * loss_mask).sum(-1) - - def concatenated_forward( - self, model: nn.Module, batch: dict[str, list | torch.LongTensor] - ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: - """Run the given model on the given batch of inputs, concatenating the chosen and rejected inputs together. - - We do this to avoid doing two forward passes, because it's faster for FSDP. - """ - concatenated_batch = self.concatenated_inputs( - batch, - is_encoder_decoder=self.is_encoder_decoder, - padding_value=self.padding_value, - device=self.accelerator.device, - ) - len_chosen = batch["chosen_labels"].shape[0] - - model_kwargs = ( - { - "decoder_input_ids": self._shift_right(concatenated_batch["concatenated_labels"]), - } - if self.is_encoder_decoder - else {} - ) - - if self.aux_loss_enabled: - model_kwargs["output_router_logits"] = True - - _unsloth_pref_head = _unsloth_pref_chunked_head(self, model) - _unsloth_pref_prior = os.environ.get('UNSLOTH_RETURN_HIDDEN_STATES') - if _unsloth_pref_head is not None: - os.environ['UNSLOTH_RETURN_HIDDEN_STATES'] = '1' - try: - outputs = model( - concatenated_batch["concatenated_input_ids"], - attention_mask=concatenated_batch["concatenated_attention_mask"], - use_cache=False, - **model_kwargs, - ) - finally: - if _unsloth_pref_head is not None: - if _unsloth_pref_prior is None: - os.environ.pop('UNSLOTH_RETURN_HIDDEN_STATES', None) - else: - os.environ['UNSLOTH_RETURN_HIDDEN_STATES'] = _unsloth_pref_prior - all_logits = outputs.logits - if _unsloth_pref_head is not None: - _unsloth_pref_out, all_logits = _unsloth_pref_chunked_forward( - self, model, _unsloth_pref_head, concatenated_batch, len_chosen, all_logits, 'orpo' - ) - if _unsloth_pref_out is not None: - return _unsloth_pref_out - - def cross_entropy_loss(logits, labels): - if not self.is_encoder_decoder: - # Shift so that tokens < n predict n - logits = logits[..., :-1, :].contiguous() - labels = labels[..., 1:].contiguous() - # Flatten the tokens - loss_fct = nn.CrossEntropyLoss() - logits = logits.view(-1, logits.shape[-1]) - labels = labels.view(-1) - # Enable model parallelism - labels = labels.to(logits.device) - loss = loss_fct(logits, labels) - return loss - - if self.is_encoder_decoder: - labels = concatenated_batch["concatenated_labels"].clone() - else: - labels = concatenated_batch["concatenated_input_ids"].clone() - attention_mask = concatenated_batch["concatenated_attention_mask"] - labels = torch.where(attention_mask == 1, labels, -100) - # orpo chosen nll loss is computed over the full prompt and response - chosen_nll_loss = cross_entropy_loss(all_logits[:len_chosen], labels[:len_chosen]) - - all_logps = self.get_batch_logps( - all_logits, - concatenated_batch["concatenated_labels"], - average_log_prob=True, - is_encoder_decoder=self.is_encoder_decoder, - ) - - chosen_logps = all_logps[:len_chosen] - rejected_logps = all_logps[len_chosen:] - - if not self.is_encoder_decoder: - chosen_logits = all_logits[:len_chosen, :-1, :] - rejected_logits = all_logits[len_chosen:, :-1, :] - else: - chosen_logits = all_logits[:len_chosen] - rejected_logits = all_logits[len_chosen:] - - if self.aux_loss_enabled: - return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, chosen_nll_loss, outputs.aux_loss) - - return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, chosen_nll_loss) - - def get_batch_loss_metrics( - self, - model, - batch: dict[str, list | torch.LongTensor], - train_eval: Literal["train", "eval"] = "train", - ): - """Compute the ORPO loss and other metrics for the given batch of inputs for train or test.""" - metrics = {} - - forward_output = self.concatenated_forward(model, batch) - ( - policy_chosen_logps, - policy_rejected_logps, - policy_chosen_logits, - policy_rejected_logits, - policy_nll_loss, - ) = forward_output[:5] - if self.aux_loss_enabled: - aux_loss = forward_output[5] - - losses, chosen_rewards, rejected_rewards, log_odds_ratio, log_odds_chosen = self.odds_ratio_loss( - policy_chosen_logps, policy_rejected_logps - ) - # full ORPO loss - loss = policy_nll_loss - losses.mean() - - reward_accuracies = (chosen_rewards > rejected_rewards).float() - - prefix = "eval_" if train_eval == "eval" else "" - metrics[f"{prefix}rewards/chosen"] = self.accelerator.gather_for_metrics(chosen_rewards).mean() - metrics[f"{prefix}rewards/rejected"] = self.accelerator.gather_for_metrics(rejected_rewards).mean() - metrics[f"{prefix}rewards/accuracies"] = self.accelerator.gather_for_metrics(reward_accuracies).mean() - metrics[f"{prefix}rewards/margins"] = self.accelerator.gather_for_metrics( - chosen_rewards - rejected_rewards - ).mean() - metrics[f"{prefix}logps/rejected"] = self.accelerator.gather_for_metrics(policy_rejected_logps).detach().mean() - metrics[f"{prefix}logps/chosen"] = self.accelerator.gather_for_metrics(policy_chosen_logps).detach().mean() - metrics[f"{prefix}logits/rejected"] = self.accelerator.gather_for_metrics( - policy_rejected_logits.detach().mean() - ).mean() - metrics[f"{prefix}logits/chosen"] = self.accelerator.gather_for_metrics( - policy_chosen_logits.detach().mean() - ).mean() - metrics[f"{prefix}nll_loss"] = self.accelerator.gather_for_metrics(policy_nll_loss).detach().mean() - metrics[f"{prefix}log_odds_ratio"] = self.accelerator.gather_for_metrics(log_odds_ratio).detach().mean() - metrics[f"{prefix}log_odds_chosen"] = self.accelerator.gather_for_metrics(log_odds_chosen).detach().mean() - if is_torch_xla_available(): - xm.mark_step() # needed because .item() calls - for k, v in metrics.items(): - metrics[k] = v.item() - if self.aux_loss_enabled: - loss += self.aux_loss_coef * aux_loss - - return loss, metrics - - def compute_loss( - self, - model: PreTrainedModel | nn.Module, - inputs: dict[str, torch.Tensor | Any], - return_outputs=False, - num_items_in_batch=None, - ) -> torch.Tensor | tuple[torch.Tensor, dict[str, torch.Tensor]]: - compute_loss_context_manager = ( - autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() - ) - - with compute_loss_context_manager: - loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval="train") - - # Make sure to move the loss to the device the original accumulating loss is at back in the `Trainer` class: - loss = loss.to(self.args.device) - - # force log the metrics - self.store_metrics(metrics, train_eval="train") - - if return_outputs: - return (loss, metrics) - return loss - - def generate_from_model(self, model, batch: dict[str, torch.LongTensor]) -> str: - """Generate samples from the model and reference model for the given batch of inputs.""" - - # If one uses `generate_during_eval` with peft + bf16, we need to explicitly call generate with - # the torch amp context manager as some hidden states are silently casted to full precision. - generate_context_manager = ( - autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() - ) - - with generate_context_manager: - policy_output = model.generate( - input_ids=batch["prompt_input_ids"], - attention_mask=batch["prompt_attention_mask"], - max_length=self.max_length, - do_sample=True, - pad_token_id=self.processing_class.pad_token_id, - ) - - policy_output = pad_to_length(policy_output, self.max_length, self.processing_class.pad_token_id) - policy_output_decoded = self.processing_class.batch_decode(policy_output, skip_special_tokens=True) - - return policy_output_decoded - - def prediction_step( - self, - model: PreTrainedModel | nn.Module, - inputs: dict[str, torch.Tensor | Any], - prediction_loss_only: bool, - ignore_keys: list[str] | None = None, - ): - if not self.use_dpo_data_collator: - logger.warning( - "prediction_step is only implemented for DPODataCollatorWithPadding, and you passed a datacollator that is different than " - "DPODataCollatorWithPadding - you might see unexpected behavior. Alternatively, you can implement your own prediction_step method if you are using a custom data collator" - ) - if ignore_keys is None: - if hasattr(model, "config"): - ignore_keys = getattr(model.config, "keys_to_ignore_at_inference", []) - else: - ignore_keys = [] - - prediction_context_manager = ( - autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() - ) - - with torch.no_grad(), prediction_context_manager: - loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval="eval") - - # force log the metrics - self.store_metrics(metrics, train_eval="eval") - - if prediction_loss_only: - return (loss.detach(), None, None) - - # logits for the chosen and rejected samples from model - logits_dict = { - "eval_logits/chosen": metrics["eval_logits/chosen"], - "eval_logits/rejected": metrics["eval_logits/rejected"], - } - logits = [v for k, v in logits_dict.items() if k not in ignore_keys] - logits = torch.tensor(logits, device=self.accelerator.device) - labels = torch.zeros(logits.shape[0], device=self.accelerator.device) - - return (loss.detach(), logits, labels) - - def store_metrics(self, metrics: dict[str, float], train_eval: Literal["train", "eval"] = "train") -> None: - for key, value in metrics.items(): - self._stored_metrics[train_eval][key].append(value) - - def evaluation_loop( - self, - dataloader: DataLoader, - description: str, - prediction_loss_only: bool | None = None, - ignore_keys: list[str] | None = None, - metric_key_prefix: str = "eval", - ) -> EvalLoopOutput: - """ - Overriding built-in evaluation loop to store metrics for each batch. Prediction/evaluation loop, shared by - `Trainer.evaluate()` and `Trainer.predict()`. - - Works both with or without labels. - """ - - # Sample and save to game log if requested (for one batch to save time) - if self.generate_during_eval: - # Generate random indices within the range of the total number of samples - num_samples = len(dataloader.dataset) - random_indices = random.sample(range(num_samples), k=self.args.eval_batch_size) - - # Use dataloader.dataset.select to get the random batch without iterating over the DataLoader - random_batch_dataset = dataloader.dataset.select(random_indices) - random_batch = self.data_collator(random_batch_dataset) - random_batch = self._prepare_inputs(random_batch) - - policy_output_decoded = self.generate_from_model(self.model, random_batch) - - table = pd.DataFrame( - columns=["Prompt", "Policy"], - data=[ - [prompt, pol[len(prompt) :]] - for prompt, pol in zip(random_batch["prompt"], policy_output_decoded, strict=True) - ], - ) - if "wandb" in self.args.report_to: - wandb.log({"game_log": wandb.Table(data=table)}) - - if "comet_ml" in self.args.report_to: - log_table_to_comet_experiment( - name="game_log.csv", - table=table, - ) - - # Base evaluation - initial_output = super().evaluation_loop( - dataloader, description, prediction_loss_only, ignore_keys, metric_key_prefix - ) - - return initial_output - - def log(self, logs: dict[str, float], start_time: float | None = None) -> None: - """ - Log `logs` on the various objects watching training, including stored metrics. - - Args: - logs (`dict[str, float]`): - The values to log. - start_time (`float`, *optional*): - Start time of the training. - """ - # logs either has 'loss' or 'eval_loss' - train_eval = "train" if "loss" in logs else "eval" - # Add averaged stored metrics to logs - for key, metrics in self._stored_metrics[train_eval].items(): - logs[key] = torch.tensor(metrics).mean().item() - del self._stored_metrics[train_eval] - return super().log(logs, start_time) - - def _shift_right(self, input_ids): - if self.decoder_start_token_id is None: - raise ValueError( - "model.config.decoder_start_token_id has to be defined. It is usually set to the pad_token_id." - ) - - # shift inputs to the right - if is_torch_fx_proxy(input_ids): - # Item assignment is not supported natively for proxies. - shifted_input_ids = torch.full(input_ids.shape[:-1] + (1,), self.decoder_start_token_id) - shifted_input_ids = torch.cat([shifted_input_ids, input_ids[..., :-1]], dim=-1) - else: - shifted_input_ids = input_ids.new_zeros(input_ids.shape) - shifted_input_ids[..., 1:] = input_ids[..., :-1].clone() - shifted_input_ids[..., 0] = self.decoder_start_token_id - - if self.pad_token_id is None: - raise ValueError("model.config.pad_token_id has to be defined.") - # replace possible -100 values in labels by `pad_token_id` - shifted_input_ids.masked_fill_(shifted_input_ids == -100, self.pad_token_id) - - return shifted_input_ids - - # Ensure the model card is saved along with the checkpoint - def _save_checkpoint(self, model, trial): - if self.args.hub_model_id is None: - model_name = Path(self.args.output_dir).name - else: - model_name = self.args.hub_model_id.split("/")[-1] - self.create_model_card(model_name=model_name) - super()._save_checkpoint(model, trial) -class UnslothORPOTrainer(_UnslothORPOTrainer): - """ - - Initialize ORPOTrainer. - - Args: - model ([`~transformers.PreTrainedModel`]): - The model to train, preferably an [`~transformers.AutoModelForSequenceClassification`]. - args ([`experimental.orpo.ORPOConfig`]): - The ORPO config arguments to use for training. - data_collator ([`~transformers.DataCollator`]): - The data collator to use for training. If None is specified, the default data collator - ([`experimental.utils.DPODataCollatorWithPadding`]) will be used which will pad the sequences to the - maximum length of the sequences in the batch, given a dataset of paired sequences. - train_dataset ([`~datasets.Dataset`]): - The dataset to use for training. - eval_dataset ([`~datasets.Dataset`]): - The dataset to use for evaluation. - processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.BaseImageProcessor`], [`~transformers.FeatureExtractionMixin`] or [`~transformers.ProcessorMixin`], *optional*): - Processing class used to process the data. If provided, will be used to automatically process the inputs - for the model, and it will be saved along the model to make it easier to rerun an interrupted training or - reuse the fine-tuned model. - model_init (`Callable[[], transformers.PreTrainedModel]`): - The model initializer to use for training. If None is specified, the default model initializer will be - used. - callbacks (`list[transformers.TrainerCallback]`): - The callbacks to use for training. - optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`): - The optimizer and scheduler to use for training. - preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): - The function to use to preprocess the logits before computing the metrics. - peft_config ([`~peft.PeftConfig`], *optional*): - The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped in - a PEFT model. - compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): - The function to use to compute the metrics. Must take a `EvalPrediction` and return a dictionary string to - metric values. - - """ - def __init__( - self, - model = None, - args = None, - data_collator = None, - train_dataset = None, - eval_dataset = None, - processing_class = None, - model_init = None, - callbacks = None, - preprocess_logits_for_metrics = None, - peft_config = None, - compute_metrics = None, - **kwargs - ): - if args is None: args = UnslothORPOConfig() - _unsloth_precision_requested = tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) - use_bf16 = getattr(args, 'bf16', False) - if type(use_bf16) is not bool: use_bf16 = False - use_fp16 = getattr(args, 'fp16', False) - if type(use_fp16) is not bool: use_fp16 = False - force_float32 = False - try: - from unsloth_zoo.device_type import device_is_bf16_supported as _bf16_supported - except Exception: - _bf16_supported = torch.cuda.is_bf16_supported - try: - from unsloth_zoo.device_type import DEVICE_TYPE as _unsloth_device_type - except Exception: - _unsloth_device_type = None - full_finetuning = getattr(model, '_unsloth_full_finetuning', None) - if full_finetuning is None: full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1' - model_forced_float32 = getattr(model, '_unsloth_forced_float32', None) - if model_forced_float32 is None: model_forced_float32 = os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1' - if model_forced_float32 and not (full_finetuning and _bf16_supported()): - print('Unsloth: Switching to float32 training since model cannot work with float16') - force_float32 = True - mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') - dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None) - if dtype is None: dtype = model.get_input_embeddings().weight.dtype - from unsloth_zoo.utils import _get_dtype - dtype = _get_dtype(dtype) - float16 = dtype == torch.float16 - bfloat16 = dtype == torch.bfloat16 - float32 = dtype == torch.float32 - user_float32 = bool(getattr(model, '_unsloth_user_float32', False)) - if full_finetuning: - if bfloat16 and use_fp16: use_fp16 = False - if float16 and use_bf16: use_bf16 = False - if not force_float32 and bfloat16 and use_fp16 and _unsloth_device_type == 'xpu' and not _bf16_supported(): - print('Unsloth: The model was loaded in bfloat16 on purpose, so training uses bfloat16 mixed precision instead of float16.') - use_fp16 = False - use_bf16 = True - args.fp16 = False - args.bf16 = True - if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`') - if not force_float32 and (bfloat16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`') - if force_float32: - # Forced float32 training - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32' and float32 and user_float32 and not _bf16_supported(): - print('Unsloth: Model is in float32 and this GPU has no bfloat16 support, so training stays in float32. Pass fp16 = True to force float16 mixed precision instead.') - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32': - # Mixed precision training. bf16 only if the GPU supports it; V100/T4 use fp16. - use_bf16_amp = (not float16) and (_bf16_supported() or (bfloat16 and _unsloth_device_type == 'xpu')) - args.fp16 = not use_bf16_amp - args.bf16 = use_bf16_amp - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16_amp else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16_amp else 'fp16' - # args.mixed_precision is a new argument which needs to be set now - elif mixed_precision_dtype == 'bfloat16': - # Both False since bfloat16 full finetuning doesn't do any autocasting. - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif use_bf16 or use_fp16: - # transformers <5 exported this itself from fp16/bf16; 5.x dropped the write, so an - # explicit flag left it unset and GRPO readers defaulted to 'fp16', wrapping a - # bfloat16 model in a float16 autocast. See unslothai/unsloth#4891. - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16 else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16 else 'fp16' - - if getattr(args, 'eval_dataset', None) is not None and getattr(args, 'eval_strategy', 'no') == 'no': - args.eval_strategy = 'steps' - if getattr(args, 'eval_steps', None) is None: args.eval_steps = 0.1 - ga_steps = getattr(args, 'gradient_accumulation_steps', None) - if ga_steps is not None and ga_steps > 1: - from transformers import __version__ as transformers_version - if Version(transformers_version) <= Version('4.45.2'): - print('**** Unsloth: Please use our fixed gradient_accumulation_steps by updating transformers, TRL and Unsloth!\n' - '`pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth transformers trl unsloth_zoo`') - if getattr(args, 'eval_strategy', 'no') != 'no': - eval_bsz = getattr(args, 'per_device_eval_batch_size', 8) - eval_generations = getattr(args, 'num_generations_eval', None) or getattr(args, 'num_generations', None) or 1 - if eval_bsz == 8 and args.per_device_train_batch_size < eval_bsz and (args.per_device_train_batch_size * args.world_size) % eval_generations == 0: args.per_device_eval_batch_size = args.per_device_train_batch_size - if getattr(args, 'eval_accumulation_steps', None) is None and ga_steps is not None: args.eval_accumulation_steps = ga_steps - fp16_full_eval = getattr(args, 'fp16_full_eval', False) - if type(fp16_full_eval) is not bool: fp16_full_eval = False - bf16_full_eval = getattr(args, 'bf16_full_eval', False) - if type(bf16_full_eval) is not bool: bf16_full_eval = False - if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True - if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False - if force_float32: - args.bf16_full_eval = False - args.fp16_full_eval = False - elif os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16': - args.bf16_full_eval = True - args.fp16_full_eval = False - elif not bf16_full_eval and not fp16_full_eval: - args.bf16_full_eval = args.bf16 - args.fp16_full_eval = args.fp16 - if tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) != _unsloth_precision_requested: - from unsloth.models.loader_utils import sync_deepspeed_mixed_precision - sync_deepspeed_mixed_precision(args) - if model is not None: - from unsloth.models.loader_utils import apply_deepspeed_zero3_model_setup - apply_deepspeed_zero3_model_setup(model) - if model is not None: - _warnings_issued = getattr(model, 'warnings_issued', None) - if _warnings_issued is None: - model.warnings_issued = {} - elif not isinstance(_warnings_issued, dict): - try: - model.warnings_issued = dict(_warnings_issued) - except Exception: - model.warnings_issued = {} - if 'max_seq_length' not in locals() and not hasattr(args, 'max_seq_length'): - pass - else: - model_max_seq_length = getattr(model, 'max_seq_length', None) - args_max_seq_length = getattr(args, 'max_seq_length', None) - if args_max_seq_length is None and model_max_seq_length is not None: - max_seq_length = model.max_seq_length - if hasattr(args, 'max_seq_length'): args.max_seq_length = max_seq_length - elif args_max_seq_length is not None and model_max_seq_length is not None: - if args_max_seq_length > model_max_seq_length: - print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but ' - 'the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.') - args.max_seq_length = model_max_seq_length - _unsloth_model_msl = getattr(model, 'max_seq_length', None) - if isinstance(_unsloth_model_msl, int) and _unsloth_model_msl > 0 and hasattr(args, 'max_length'): - if args.max_length is None or args.max_length > _unsloth_model_msl: - print('Unsloth: `max_length = ' + str(args.max_length) + '` exceeds the model max_seq_length of ' + str(_unsloth_model_msl) + ', so it is reduced to ' + str(_unsloth_model_msl) + '.') - args.max_length = _unsloth_model_msl - _unsloth_mpl = getattr(args, 'max_prompt_length', 0) - if (_unsloth_mpl is None and args.max_length <= 128) or (_unsloth_mpl is not None and _unsloth_mpl >= args.max_length): - args.max_prompt_length = args.max_length // 2 - if model is not None and hasattr(model, 'for_training'): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right' - if 'processing_class' in locals(): - if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right' - if hasattr(processing_class, 'tokenizer') and hasattr(processing_class.tokenizer, 'padding_side'): processing_class.tokenizer.padding_side = 'right' - __tokenizer = processing_class if 'processing_class' in locals() else tokenizer - from unsloth_zoo.vision_utils import UnslothVisionDataCollator - if not isinstance(data_collator, UnslothVisionDataCollator): - if isinstance(data_collator, DataCollatorForSeq2Seq) and 'labels' not in _unsloth_dataset_column_names(train_dataset): - data_collator = TransformersDataCollatorForLanguageModeling( - __tokenizer, - mlm = False, - mlm_probability = 0.0, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling) and 'labels' in _unsloth_dataset_column_names(train_dataset): - data_collator = DataCollatorForSeq2Seq( - __tokenizer, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - else: - if hasattr(args, 'remove_unused_columns'): args.remove_unused_columns = False - if hasattr(args, 'dataset_text_field'): args.dataset_text_field = '' - if hasattr(args, 'dataset_kwargs'): args.dataset_kwargs = {'skip_prepare_dataset': True} - if not isinstance(data_collator, UnslothVisionDataCollator): - if not hasattr(__tokenizer, 'pad') and hasattr(__tokenizer, 'tokenizer'): - if isinstance(data_collator, DataCollatorForSeq2Seq): - data_collator = DataCollatorForSeq2Seq( - __tokenizer.tokenizer, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling): - data_collator = TransformersDataCollatorForLanguageModeling( - __tokenizer.tokenizer, - mlm = False, - mlm_probability = 0.0, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - other_metrics = [] - - from unsloth_zoo.logging_utils import PatchRLStatistics - PatchRLStatistics('orpo_trainer', other_metrics) - - # [TODO] Fix up DataParallel multiplying batch sizes - # [TODO] DDP works, but DP seems to not work? [TODO] - if getattr(args, "parallel_mode", None) == ParallelMode.NOT_DISTRIBUTED and args.n_gpu > 1: - if getattr(args, "_n_gpu", 1) != 1: - args._n_gpu = 1 - if "model" in locals(): - try: - from unsloth.models._gc_config_default import honour_gradient_checkpointing_off - honour_gradient_checkpointing_off(args, model) - except Exception: - pass - if "model" in locals() and hasattr(model, "for_training"): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if "model" in locals(): - try: - from unsloth.models._ddp_defaults import unsloth_default_ddp_find_unused_parameters - unsloth_default_ddp_find_unused_parameters(args, model) - except Exception: - pass - super().__init__( - model = model, - args = args, - data_collator = data_collator, - train_dataset = train_dataset, - eval_dataset = eval_dataset, - processing_class = processing_class, - model_init = model_init, - callbacks = callbacks, - preprocess_logits_for_metrics = preprocess_logits_for_metrics, - peft_config = peft_config, - compute_metrics = compute_metrics,**kwargs) - # A tensor parallel model (tp_native.tp_trainer_setup) stays on its training path: for_inference would - # route the decoder layers onto decode paths outside the tensor parallel regions. - if "model" in locals() and hasattr(model, "for_inference") and not getattr(self, "_unsloth_tensor_parallel", False): - model.for_inference() - if hasattr(self, 'neftune_hook_handle'): - self.neftune_hook_handle.remove() - if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle - if getattr(args, 'neftune_noise_alpha', None) is not None: - model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha - pass - if hasattr(self, 'accelerator'): - scaler = self.accelerator.scaler - current_model = model - while hasattr(current_model, 'model'): - current_model.accelerator_scaler = scaler - current_model = current_model.model - current_model.accelerator_scaler = scaler - pass - if hasattr(self, 'train'): - self.train = MethodType(prepare_for_training_mode(self.__class__.train), self) - pass - if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'): - _vllm_tok = self.llm.get_tokenizer() - _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None) - if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None: - _vllm_tok.chat_template = _pc.chat_template - pass - if getattr(self, 'aux_loss_enabled', False) and hasattr(getattr(self, 'model', None), 'config'): - _text_config = self.model.config - if hasattr(_text_config, 'get_text_config'): _text_config = _text_config.get_text_config() - _n_experts = [getattr(_text_config, _k) for _k in ('num_local_experts', 'num_experts', 'n_routed_experts', 'moe_num_experts') if isinstance(getattr(_text_config, _k, None), int)] - if _n_experts and all(_n <= 1 for _n in _n_experts): - self.aux_loss_enabled = False - _text_config.output_router_logits = False - pass - -pass - - -if hasattr(logger, "addFilter"): - import logging - class HideLoggingMessage(logging.Filter): - def __init__(self, text): self.text = text - def filter(self, x): return not (self.text in x.getMessage()) - pass - logger.addFilter(HideLoggingMessage("`use_cache=True`")) - diff --git a/unsloth_compiled_cache/UnslothOnlineDPOTrainer.py b/unsloth_compiled_cache/UnslothOnlineDPOTrainer.py deleted file mode 100644 index f31df022a8ce6cc30511cf2903c5abc49facd2a0..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/UnslothOnlineDPOTrainer.py +++ /dev/null @@ -1,2941 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from unsloth_zoo.temporary_patches.common import torch_compile -from unsloth_zoo.temporary_patches.common import _maybe_compile -import functools -import transformers -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from trl.experimental.online_dpo.online_dpo_trainer import (Any, AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer, Callable, DPODataCollatorWithPadding, DataCollator, Dataset, EvalPrediction, F, FSDP, GenerationConfig, IterableDataset, MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES, OnlineDPOConfig, OnlineDPOTrainer, OptimizerNames, Path, PeftConfig, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, RewardFunc, TrainerCallback, Version, _BaseTrainer, apply_chat_template, bnb, broadcast_object_list, create_reference_model, disable_dropout_in_model, empty_cache, gather_object, get_callable_name, get_config_model_id, is_conversational, is_peft_available, is_peft_model, maybe_apply_chat_template, nested_gather, nn, nullcontext, os, prepare_deepspeed, prepare_fsdp, prepare_multimodal_messages, profiling_context, re, set_seed, textwrap, torch, transformers, truncate_right, unwrap_model_for_generation, AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer, Callable, DPODataCollatorWithPadding, DataCollator, Dataset, EvalPrediction, F, FSDP, GenerationConfig, IterableDataset, MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES, OnlineDPOConfig, PeftConfig, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, RewardFunc, TrainerCallback, Version, create_reference_model, disable_dropout_in_model, get_callable_name, get_config_model_id, is_peft_available, is_peft_model, nn, os, prepare_deepspeed, prepare_fsdp, re, set_seed, torch, transformers, re, torch, F, apply_chat_template, empty_cache, is_conversational, re, F, FSDP, is_peft_model, nn, nullcontext, os, re, F, PreTrainedModel, is_peft_available, os, re, torch, F, FSDP, nn, os, re, F, FSDP, is_peft_model, nn, re) - - -import os -import math -import logging -from typing import * -from dataclasses import dataclass, field -from packaging.version import Version -import torch -import numpy as np -from contextlib import nullcontext -from torch.nn import functional as F -import inspect -from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling -from transformers.training_args import ParallelMode -from unsloth_zoo.device_type import DEVICE_TYPE, DEVICE_TYPE_TORCH, device_synchronize - -# Wrap trainer with padding to right and enable training mode -import functools -from types import MethodType -try: - from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers -except: - def reset_unsloth_gradient_checkpointing_buffers(): pass -# Canonical reset lives in unsloth.models._utils so the SFT auto-packing wrapper and the plain -# Trainer loop can import the same helper; fall back to a no-op only if it can't be imported. -try: - from unsloth.models._utils import _unsloth_reset_stray_compile_cache -except Exception: - def _unsloth_reset_stray_compile_cache(self): pass -try: - from unsloth.models._utils import _unsloth_dataset_column_names -except Exception: - def _unsloth_dataset_column_names(dataset): return dataset.column_names -# Drops/renames config arguments the installed TRL no longer accepts, so a -# script pinned to an older TRL keeps working after an upgrade. Falls back to -# the historical raw passthrough so this can never break trainer construction. -try: - from unsloth.models.rl_config_compat import filter_config_init_kwargs as _unsloth_filter_config_init_kwargs - # A cache file generated here can be imported by an older Unsloth whose filter - # predates `mirrored_from`, so drop the argument rather than raise TypeError. - if "mirrored_from" not in inspect.signature(_unsloth_filter_config_init_kwargs).parameters: - _unsloth_filter_config_init_kwargs_old = _unsloth_filter_config_init_kwargs - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): - return _unsloth_filter_config_init_kwargs_old(config_class, kwargs) -except Exception: - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): return kwargs -def prepare_for_training_mode(f): - @functools.wraps(f) - def wrapper(self, *args, **kwargs): - # Drop any torch.compile graph cache poisoned by a stray pre-train forward. - try: - _unsloth_reset_stray_compile_cache(self) - except Exception: - pass - # Finish the previous W&B run if this is a subsequent train() call. - # We do this at the START of train() (not the end) so that - # evaluate() / log() still work after train() completes. - # HF's WandbCallback.setup() will call wandb.init() for the new run. - # See: https://github.com/unslothai/unsloth/issues/3954 - if getattr(self, '_unsloth_training_completed', False): - try: - import wandb - if wandb.run is not None: - wandb.finish() - # Reset HF's WandbCallback so it calls wandb.init() for the new run - for cb in self.callback_handler.callbacks: - if type(cb).__name__ == 'WandbCallback': - cb._initialized = False - break - except: - pass - # Enable training mode - _was_training = None - # Restore the GC mode the model was configured with at setup; fall back to - # the training args only when it wasn't recorded (issue #4735). Use hasattr, - # not a None sentinel, so a deliberately-recorded None is restored verbatim. - _model = getattr(self, 'model', None) - if hasattr(_model, '_unsloth_gradient_checkpointing'): - use_gc = _model._unsloth_gradient_checkpointing - else: - use_gc = getattr(self.args, 'gradient_checkpointing', True) - if hasattr(self, 'model') and hasattr(self.model, "training"): - _was_training = self.model.training - if hasattr(self, 'model') and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - output = f(self, *args, **kwargs) - # Restore previous mode when possible - if hasattr(self, 'model') and hasattr(self.model, "for_inference"): - if _was_training is False: - self.model.for_inference() - elif _was_training is True and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - # Reset gradient checkpointing buffers to free memory while staying ready for next run - try: - reset_unsloth_gradient_checkpointing_buffers() - except: - pass - # Mark that training completed so the next train() call can - # finish this W&B run before starting a new one - self._unsloth_training_completed = True - return output - return wrapper -pass - -torch_compile_options = { - "epilogue_fusion" : True, - "max_autotune" : False, - "shape_padding" : True, - "trace.enabled" : False, - "triton.cudagraphs" : False, -} - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_hidden_states_selective_log_softmax( - hidden_states: torch.Tensor, - lm_head: torch.Tensor, - index: torch.Tensor, - chunks: int = 4, - logit_scale_multiply: float = 0.0, - logit_scale_divide: float = 0.0, - logit_softcapping: float = 0.0, - temperature: float = 1.0, - # Rows per chunk cap. Read HERE, in the default, not in the body: the body - # is traced with fullgraph = True, and `os.environ` is an unsupported op - # there on torch 2.4 -- Dynamo raises - # torch._dynamo.exc.Unsupported: const method call bytes.decode - # from os._Environ.__getitem__, which the eager fallback does not catch - # (it only catches recompile-limit and disabled-hook breaks), so the very - # first call would die even with the variable unset. A default is evaluated - # once when the def runs, which is import time, outside any traced region. - # It is also a plain int argument, so Dynamo guards on it instead of - # constant-folding an unguarded read (2.7+ never notice a later change). - # A non-numeric value is ignored rather than raised on, so a typo cannot - # break the import. 0 keeps the previous chunk boundaries exactly. - max_rows_per_chunk: int = ( - int(os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip()) - if os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip().isdigit() - else 0 - ), - lm_head_bias: Optional[torch.Tensor] = None, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, - # Read in the default for the same Dynamo reason as max_rows_per_chunk. Pairs with vLLM's - # head_dtype=float32 that load_vllm requests under the same flag. - fp32_lm_head: bool = ( - os.environ.get("UNSLOTH_VLLM_FP32_LM_HEAD", "0").strip().lower() in ("1", "true", "on", "yes") - ), - # Older Unsloth scores old / ref log-probs here and the policy with the fused kernel: one scorer keeps ratio at 1. - fused_linear: bool = ( - os.environ.get("UNSLOTH_FUSED_LINEAR", "1") != "0" - and os.environ.get("UNSLOTH_GRPO_CHUNKED_LOGPROBS", "1") != "0" - ), -) -> torch.Tensor: - # All Unsloth Zoo code licensed under AGPL3 - from unsloth_zoo.fused_losses import chunked_logprob as _chunked_logprob - if (fused_linear and _chunked_logprob._FUSED_LINEAR and not fp32_lm_head - and support_top_k == 0 and support_top_p == 1.0 and support_min_p == 0.0 - and support_excluded_ids is None): - return _chunked_logprob.unsloth_chunked_logprobs( - hidden_states, lm_head, index, - temperature = temperature, - logit_scale_multiply = logit_scale_multiply, - logit_scale_divide = logit_scale_divide, - logit_softcapping = logit_softcapping, - lm_head_bias = lm_head_bias, - max_rows_per_chunk = max_rows_per_chunk if max_rows_per_chunk > 0 else None, - ).reshape((hidden_states.shape[0], hidden_states.shape[1])) - # Reshape on this tensor's own last dim: a no-op, so a wrong-width caller - # cannot have its row count silently rewritten and instead fails at the - # matmul below, which prints both operands. Do not swap in a bare - # torch._check: it reports only "Expected cond to be True", naming neither - # operand, and Dynamo rejects a message-carrying one. Callers dispatch on - # the width first -- see `compute_logprobs_chunk`, the packed path and - # `_pg_grad_forward`. - flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) - flat_index = index.reshape(-1) - - # Each chunk materialises rows x vocab logits and then a float32 copy of - # them, all on the device holding the output head. With a large vocabulary - # and a fixed chunk count that grows with the batch, so the peak scales with - # the batch rather than staying bounded. max_rows_per_chunk caps the rows - # per chunk instead, which is pure loop splitting: more, smaller chunks, - # same concatenated result. 0 (the default) keeps the previous chunk - # boundaries exactly. - if max_rows_per_chunk > 0: - n_rows = flat_hidden_states.shape[0] - chunks = max(chunks, -(-n_rows // max_rows_per_chunk)) - chunks = min(chunks, max(n_rows, 1)) - - # A one-row chunk's head gradient is a K=1 matmul Inductor mis-lowers under dynamic shapes (illegal - # memory access or NaN head gradients, torch 2.13), so keep chunks at 2+ rows when the head trains. - # Exactly `chunks` row ranges, so the unrolled loop length never depends on T: torch.chunk - # returns ceil(T / ceil(T / chunks)) pieces and recompiled per new count; tensor_split guards - # on T % chunks. At most rows // 2 chunks when the head trains keeps every piece at 2+ rows. - lm_head_grad = torch.is_grad_enabled() and lm_head.requires_grad - if lm_head_grad and flat_hidden_states.shape[0] < 2 * chunks: - chunks = max(flat_hidden_states.shape[0] // 2, 1) - - n_rows = flat_hidden_states.shape[0] - chunks = max(1, min(chunks, n_rows)) - bounds = [(i * n_rows) // chunks for i in range(chunks + 1)] - chunked_hidden_states = [flat_hidden_states[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - - # Cast once, outside the loop, so every chunk shares one fp32 head copy (also for autograd). - if fp32_lm_head: - lm_head = lm_head.to(torch.float32) - if lm_head_bias is not None: - lm_head_bias = lm_head_bias.to(torch.float32) - - all_per_token_logps = [] - - for chunk_hidden_states, chunk_index in zip(chunked_hidden_states, chunked_index): - # When the model is dispatched over several devices, the output head can - # sit on a different one from the hidden states, because accelerate - # places the tail of the model on the last device it fills. Co-locate on - # the head's device before the matmul, otherwise this raises - # Unhandled FakeTensor Device Propagation for aten.mm.default, - # found two different devices cuda:0, cuda:1 - # On a single device every .to() here is a no-op and the result is - # bit-identical to before. - chunk_hidden_states = chunk_hidden_states.to(device = lm_head.device, dtype = lm_head.dtype) - chunk_index = chunk_index.to(lm_head.device) - if fp32_lm_head: - # The GRPO loss runs this under bf16 / fp16 autocast, which would downcast the matmul. - with torch.autocast(device_type = chunk_hidden_states.device.type, enabled = False): - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - elif lm_head_bias is None: - chunk_logits = chunk_hidden_states @ lm_head.t() - else: - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - - if logit_scale_multiply != 0.0: - chunk_logits = chunk_logits * logit_scale_multiply - if logit_scale_divide != 0.0: - chunk_logits = chunk_logits / logit_scale_divide - if logit_softcapping != 0.0: - chunk_logits = logit_softcapping * torch.tanh(chunk_logits / logit_softcapping) - - chunk_logits = chunk_logits.to(torch.float32) - - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - per_token_logps = _truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - ) - else: - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - # Return to the caller's device so the concatenation below and every - # downstream consumer see the device they started on. - all_per_token_logps.append(per_token_logps.to(hidden_states.device)) - - all_per_token_logps = torch.concat(all_per_token_logps) - - all_per_token_logps = all_per_token_logps.reshape((hidden_states.shape[0], hidden_states.shape[1])) - return all_per_token_logps - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_selective_log_softmax( - logits, - index, - temperature: float = 1.0, - chunks: int = 4, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, -): - # Fixed `chunks` row ranges; see chunked_hidden_states_selective_log_softmax. - flat_logits = logits.reshape(-1, logits.shape[-1]) - flat_index = index.reshape(-1) - # Fewer rows than chunks: one row per range, as torch.chunk did, never empty ranges. - chunks = max(1, min(chunks, flat_logits.shape[0])) - bounds = [(i * flat_logits.shape[0]) // chunks for i in range(chunks + 1)] - chunked_logits = [flat_logits[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - all_per_token_logps = [] - # Per-chunk selective_log_softmax. - for chunk_logits, chunk_index in zip(chunked_logits, chunked_index): - chunk_logits = chunk_logits.to(torch.float32) - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - all_per_token_logps.append(_truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - )) - continue - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - all_per_token_logps.append(per_token_logps) - pass - all_per_token_logps = torch.concat(all_per_token_logps) - all_per_token_logps = all_per_token_logps.reshape((logits.shape[0], logits.shape[1])) - return all_per_token_logps - -def calculate_pad_tokens_in_prompt( - input_ids: torch.Tensor, - logits_to_keep: int, - pad_token_id: int -) -> torch.Tensor: - """Count left-padded tokens per sequence, e.g. [pad, pad, pad, cat] -> 3.""" - if logits_to_keep >= input_ids.shape[1]: - raise ValueError("logits_to_keep must be smaller than the sequence length.") - - prompt_section = input_ids[:, :-logits_to_keep] - - padding_mask = (prompt_section == pad_token_id) - - pad_token_counts = padding_mask.sum(dim=1) - - return pad_token_counts - -def create_completion_attention_mask( - completion_input_ids: torch.Tensor, - left_pad_tokens_per_prompt: torch.Tensor, - max_left_pad: int, - pad_token_id: int -) -> torch.Tensor: - """Build a completion mask that zeros leading prompt and trailing pad tokens. - - For [p,p,p,c,c,c,pad,pad,pad] (p=sliced prompt, c=completion, pad=padding) - this returns [0,0,0,1,1,1,0,0,0]. - """ - batch_size, completion_len = completion_input_ids.shape - device = completion_input_ids.device - - num_tokens_to_mask = max_left_pad - left_pad_tokens_per_prompt - - indices = torch.arange(completion_len, device=device).unsqueeze(0) - shift_mask = indices >= num_tokens_to_mask.unsqueeze(1) - - non_padding_mask = (completion_input_ids != pad_token_id) - - final_mask = shift_mask & non_padding_mask - - return final_mask - -def left_pack_padding(tensor: torch.Tensor, pad_id: int) -> torch.Tensor: - """Move all padding tokens in each sequence to the right.""" - mask = (tensor != pad_id) - # stable=True since the binary mask is unordered. - sorted_indices = torch.argsort(mask, dim=1, descending=True, stable=True) - packed_tensor = torch.gather(tensor, 1, sorted_indices) - return packed_tensor - -def align_logprobs_with_mask( - logprob_tensor: torch.Tensor, - attention_mask: torch.Tensor, - pad_value: float = 0.0 -) -> torch.Tensor: - """Align a log probability tensor with a given attention mask.""" - - device = logprob_tensor.device - batch_size, logprob_seq_len = logprob_tensor.shape - mask_seq_len = attention_mask.shape[1] - - padded_logprobs = torch.full( - attention_mask.shape, - fill_value=pad_value, - dtype=logprob_tensor.dtype, - device=device - ) - - left_pad_counts = torch.argmax(attention_mask, dim=1) - - cols = torch.arange(logprob_seq_len, device=device) - dest_indices = left_pad_counts.unsqueeze(1) + cols - - # Destination row indices, shape [batch_size, logprob_seq_len]. - row_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand_as(dest_indices) - - # Keep only in-bounds destinations, then scatter via advanced indexing. - valid_mask = dest_indices < mask_seq_len - valid_rows = row_indices[valid_mask] - valid_cols = dest_indices[valid_mask] - valid_vals = logprob_tensor[valid_mask] - padded_logprobs[valid_rows, valid_cols] = valid_vals - - return padded_logprobs - -def align_completion_tool_mask( - tool_mask: torch.Tensor, - completion_mask: torch.Tensor, -) -> torch.Tensor: - """Align a raw completion-length tool/env mask with Unsloth's repacked loss mask.""" - if tool_mask is None: - return completion_mask - if tool_mask.shape[0] != completion_mask.shape[0]: - raise ValueError("tool_mask batch size must match completion_mask batch size.") - - tool_mask = tool_mask.to(device=completion_mask.device) - if tool_mask.shape == completion_mask.shape: - aligned_tool_mask = tool_mask - else: - aligned_tool_mask = align_logprobs_with_mask( - tool_mask, - completion_mask, - pad_value=0, - ) - return completion_mask * aligned_tool_mask.to(dtype=completion_mask.dtype) - -def autotune_batch_and_chunks( - total_input_rows, - seq_len, - hidden_size, - vocab_size, - dtype_bytes=16, - multiplier=None -): - if multiplier is None: - final_m = max(4, seq_len // 4096) - else: - final_m = multiplier - - if torch.cuda.is_available(): - free_bytes, _ = torch.cuda.mem_get_info() - limit_gb = (free_bytes / (1024**3))*.80 - elif hasattr(torch, "xpu") and torch.xpu.is_available(): - # XPU: estimate free memory as total - reserved. - total_mem = torch.xpu.get_device_properties(0).total_memory - reserved_mem = torch.xpu.memory_reserved() - free_bytes = total_mem - reserved_mem - limit_gb = (free_bytes / (1024**3)) * 0.80 - else: - # Fallback: assume 8GB available. - limit_gb = 8.0 - - bytes_to_gb = 1024**3 - - b_vals = torch.arange(total_input_rows, 0, -1, device='cpu', dtype=torch.float32) - - hidden_gb = (b_vals * seq_len * hidden_size * dtype_bytes) / bytes_to_gb - - base_logits = ((b_vals/total_input_rows) * b_vals * seq_len * vocab_size * dtype_bytes) / bytes_to_gb - logits_gb = base_logits / final_m - - total_mem_gb = hidden_gb + logits_gb - - valid_mask = total_mem_gb <= limit_gb - valid_indices = torch.nonzero(valid_mask, as_tuple=False) - - if valid_indices.shape[0] == 0: - #This means your GPU will OOM - # Capped at the row count: unsloth's no-grad pass divides rows by this without max(1, ...). - return max(1, min(4, total_input_rows)), final_m - - best_idx = valid_indices[0].item() - final_b = int(b_vals[best_idx].item()) - - return final_b, final_m - -def sanitize_logprob(logprob): - """Local port of trl.scripts.vllm_serve.sanitize_logprob. - Filters NaN logprobs from vLLM outputs.""" - value = logprob.logprob - if math.isnan(value): - logging.getLogger(__name__).warning( - f"Generated NaN logprob, token logprob '{logprob}' will be ignored" - ) - return None - return value -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_text_configs(config): - """``config`` and its text sub-config, the two places a transform can live: Gemma-4-style configs keep it on ``config.text_config``, T5Gemma-style ones only reach it via ``config.get_text_config()``.""" - if config is None: - return [] - holders = [config] - 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: - holders.append(text_cfg) - return holders - -def _unsloth_get_final_logit_softcapping(model): - """The soft cap the loss applies, under any of its three spellings: ``final_logit_softcapping`` (Gemma), ``logits_soft_cap`` (RecurrentGemma), ``output_logit_soft_cap`` (xLSTM). Returns 0 if unset.""" - config = _unsloth_get_model_config(model) - if config is None: - return 0 - for holder in _unsloth_text_configs(config): - for name in ("final_logit_softcapping", "logits_soft_cap", "output_logit_soft_cap"): - softcap = getattr(holder, name, None) - if softcap: - return softcap - return 0 - -def _unsloth_resolve_logit_scales(model_config): - """``(multiply, divide)`` for the logits, read with the same field table and per-family overrides ``detect_logit_transforms`` uses. - - Only reached when the installed unsloth_zoo predates that helper. Must stay in step with ``resolve_logit_transforms`` in unsloth/models/llama.py: if the forward applies a transform GRPO does not, the policy log-probabilities come from different logits than the ones generated, which silently shifts every importance ratio. - """ - # Granite divides, HyperCLOVA X multiplies, MiniCPM3 scales hidden states (not a logit transform). - overrides = { - ("logits_scaling", "hyperclovax"): "multiply", - ("logits_scaling", "minicpm3"): None, - } - found = {"multiply": 0, "divide": 0} - for holder in _unsloth_text_configs(model_config): - model_type = getattr(holder, "model_type", "") or "" - for bucket, names in ( - ("multiply", ("logit_scale", "lm_head_multiplier", "output_multiplier")), - ("divide", ("logits_scaling",)), - ): - for name in names: - target = overrides.get((name, model_type), bucket) - if target is None or found[target]: - continue - value = getattr(holder, name, 0) or 0 - if value: - found[target] = value - break - return found["multiply"], found["divide"] - -def _unsloth_grpo_returns_hidden_states(model, tensor, lm_head): - """Does ``tensor`` (a forward's ``.logits``) carry hidden states or real logits? - - ``_get_per_token_logps_and_entropies`` sets ``UNSLOTH_RETURN_HIDDEN_STATES=1``, but only a forward that honours the name hands hidden states back as ``.logits``; any other forward returns a real ``[.., vocab]`` tensor that must not reach the ``lm_head`` matmul. - - The primary test is an explicit signal that the forward honours the flag, both set outside this file: ``__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__`` on the generated class, written by ``unsloth_zoo.compiler.create_standalone_class`` exactly when ``apply_fused_lm_head`` gave that forward its own branch; and ``_unsloth_grpo_hidden_states_forward_wrapped``, set by ``_install_grpo_hidden_states_forward_wrapper`` in ``unsloth/models/rl.py`` for models the compiler did not rewrite. That wrapper degrades to real logits when the model cannot produce hidden states and records whether it did so in ``_unsloth_grpo_hidden_states_degraded`` before returning, so reading the pair after a forward describes the call that just finished; degradation is per call, not per model. - - The width comparison stays as the fallback, for an ``unsloth_zoo`` old enough that it never writes the marker. It is decisive whenever ``vocab_size != hidden_size``, and the signal may only overrule it when it is not, which is the one case the shape cannot answer. - """ - if tensor.shape[-1] != lm_head.shape[1]: - return False # vocab-wide: real logits, whatever any signal claims - if lm_head.shape[0] != lm_head.shape[1]: - return True # hidden-wide and vocab_size != hidden_size: hidden states - return _unsloth_grpo_hidden_states_signal(model) is not False - -def _unsloth_grpo_hidden_states_signal(model): - """``True``/``False`` if the forward honours ``UNSLOTH_RETURN_HIDDEN_STATES``, ``None`` when neither marker is present. See ``_unsloth_grpo_returns_hidden_states`` for where each marker is set. Walks the wrapper chain because the markers are set on whichever object the trainer saw, which may be the DDP module or the PEFT base model rather than the object handed to the logprob loop.""" - candidates = [] - pending = [model] - while pending and len(candidates) < 8: - candidate = pending.pop(0) - if candidate is None or any(candidate is seen for seen in candidates): - continue - candidates.append(candidate) - get_base_model = getattr(candidate, "get_base_model", None) - if callable(get_base_model): - try: - pending.append(get_base_model()) - except Exception: - pass - for _attr in ("module", "base_model", "model"): - child = getattr(candidate, _attr, None) - if child is not None and hasattr(child, "forward"): - pending.append(child) - for candidate in candidates: - if getattr(candidate, "__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__", False): - return True - if getattr(type(candidate), "__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__", False): - return True - if any( - getattr(candidate, "_unsloth_grpo_hidden_states_forward_wrapped", False) - for candidate in candidates - ): - # The wrapper honours the flag unless it recorded that this call could not. - if any( - hasattr(candidate, "_unsloth_grpo_hidden_states_degraded") for candidate in candidates - ): - return not any( - getattr(candidate, "_unsloth_grpo_hidden_states_degraded", False) - for candidate in candidates - ) - # An unsloth/models/rl.py predating the per-call attribute set only the warn-once flag; it is the best signal such a wrapper offers. - return not any( - getattr(candidate, "_unsloth_grpo_hidden_states_warning_issued", False) - for candidate in candidates - ) - return None - -def _unsloth_pref_returns_hidden_states(model, tensor, lm_head): - """``_unsloth_grpo_returns_hidden_states`` for the preference routes, which need a zoo that writes the hidden-state - markers: a square head with no signal is a forward that ignored the flag, so its output is real logits.""" - if tensor.shape[-1] != lm_head.shape[1]: - return False - if lm_head.shape[0] != lm_head.shape[1]: - return True - return _unsloth_grpo_hidden_states_signal(model) is True - -def _unsloth_gkd_dense_head(model): - """The output head when it is a plain dense ``[vocab, hidden]`` projection the chunked loss can read directly. - Rejects bnb heads (``nn.Linear`` subclasses with packed weights), PEFT heads and ZeRO-3 partitioned weights. - """ - get_output_embeddings = getattr(model, "get_output_embeddings", None) - if not callable(get_output_embeddings): - return None - try: - head = get_output_embeddings() - except Exception: - return None - if type(head) is not torch.nn.Linear: - return None - # An lm_head that is not the output embeddings runs before them (ModernBERT decoder, RoBERTa-style heads). - lm_head = getattr(model, "lm_head", None) - if isinstance(lm_head, torch.nn.Module) and lm_head is not head: - return None - weight = getattr(head, "weight", None) - if not isinstance(weight, torch.Tensor) or weight.dim() != 2 or weight.numel() == 0: - return None - if weight.dtype not in (torch.float16, torch.bfloat16, torch.float32, torch.float64): - return None - if type(weight) not in (torch.Tensor, torch.nn.Parameter): - return None - bias = getattr(head, "bias", None) - if bias is not None and (bias.dim() != 1 or bias.shape[0] != weight.shape[0]): - return None - if weight.is_meta or (bias is not None and bias.is_meta): - return None - return head - -def _unsloth_gkd_logit_transforms(model): - """``(scale, softcap)`` exactly as the model's own forward applies them to its logits.""" - model_config = _unsloth_get_model_config(model) - if detect_logit_transforms is not None: - _transforms = detect_logit_transforms(model_config) - logit_softcapping = _transforms["logit_softcapping"] - logit_scale_multiply = _transforms["logit_scale_multiply"] - logit_scale_divide = _transforms["logit_scale_divide"] - else: - logit_softcapping = _unsloth_get_final_logit_softcapping(model) - logit_scale_multiply, logit_scale_divide = _unsloth_resolve_logit_scales(model_config) - scale = 1.0 - if logit_scale_multiply: - scale = scale * float(logit_scale_multiply) - if logit_scale_divide: - scale = scale / float(logit_scale_divide) - return scale, float(logit_softcapping or 0.0) - -def _unsloth_pref_logit_transforms(model): - """``(multiply, divide, softcap)`` for the chunked kernel, from the GKD resolver (one scale, 0.0 = off).""" - scale, softcap = _unsloth_gkd_logit_transforms(model) - return (0.0 if scale == 1.0 else scale), 0.0, softcap - -def _unsloth_stock_logps_hook( - self, - method, - hooks = ("get_batch_logps",), -): - """False when a subclass overrides one of ``hooks`` below the class whose ``method`` is the rewritten one: TRL's - body would call the override (or it reads the full logits the chunked path replaces with stand-ins).""" - owner = None - for klass in type(self).__mro__: - # The rewritten body calls a chunked or fused-pass helper; a subclass method wrapping it via super() does not. - names = getattr(getattr(vars(klass).get(method), "__code__", None), "co_names", ()) - if any( - name.startswith("_unsloth_") and ("chunked" in name or "fused_pass" in name) - for name in names - ): - owner = klass - break - if owner is None: - return True - for name in hooks: - hook = next((klass for klass in type(self).__mro__ if name in vars(klass)), None) - if hook is not None and hook not in owner.__mro__: - return False - return True - -def _unsloth_dpo_chunked_head( - self, - model, - concatenated_batch, - method = "concatenated_forward", -): - """``(unwrapped_model, head, (multiply, divide, softcap))`` when concatenated_forward can score from hidden - states, else ``None`` (TRL's own logits path runs). Decided before the forward, which then skips the logits.""" - reason = None - if _unsloth_masked_logprobs is None: - reason = "unsloth_zoo has no chunked_logprob" - elif os.environ.get("UNSLOTH_FUSED_LINEAR", "1") == "0": - reason = "UNSLOTH_FUSED_LINEAR=0" - elif os.environ.get("UNSLOTH_DPO_CHUNKED", "1").lower() in ("0", "false", "no", "off"): - reason = "UNSLOTH_DPO_CHUNKED=0" - elif getattr(self, "is_encoder_decoder", False): - reason = "encoder-decoder" - elif getattr(self, "aux_loss_enabled", False): - reason = "router aux loss needs the stock forward" - elif not _unsloth_stock_logps_hook( - self, - method, - ("get_batch_logps",) - if method == "concatenated_forward" - else ("get_batch_logps", "get_batch_loss_metrics", method), - ): - # KTO hands its callers logits stand-ins; DPO's outputs are exact, so a DPO wrapper is fine. - reason = "get_batch_logps / get_batch_loss_metrics overridden" - elif getattr(self, "padding_free", False): - reason = "padding_free" - elif getattr(self, "use_weighting", False): - reason = "use_weighting needs full-vocabulary log-probs" - elif getattr(self, "is_fsdp_enabled", False) or getattr(self, "is_deepspeed_enabled", False): - reason = "FSDP / DeepSpeed" - elif any( - k in concatenated_batch - for k in ("pixel_values", "pixel_attention_mask", "image_sizes", "token_type_ids") - + tuple(_DPO_VISION_KEYS) - ): - reason = "vision inputs" - if reason is None: - try: - unwrapped = self.accelerator.unwrap_model(model) - except Exception: - unwrapped = None - reason = "unwrap_model" - if reason is None: - head = _unsloth_gkd_dense_head(unwrapped) - if head is None: - reason = "output head is not a dense nn.Linear" - # The kernel reads weight / bias directly: a hooked head (accelerate, forward / backward hooks) is not - # F.linear(h, W, b), and its backward hooks would never see the gradient. - elif ( - hasattr(head, "_hf_hook") - or "forward" in vars(head) - or head._forward_hooks - or head._forward_pre_hooks - or head._backward_hooks - or getattr(head, "_backward_pre_hooks", None) - ): - reason = "output head has hooks or a patched forward" - if reason is None: - for config in _unsloth_text_configs(_unsloth_get_model_config(unwrapped)): - if getattr(config, "model_type", None) == "minicpm3": - reason = "minicpm3 scales hidden states before the head" - if reason is None and getattr(model, "find_unused_parameters", False): - if any(p is not None and p.requires_grad for p in (head.weight, head.bias)): - reason = "DDP find_unused_parameters with a trainable head" - if reason is not None: - try: - fallbacks = self.__dict__.setdefault("_unsloth_dpo_chunked_fallbacks", {}) - fallbacks[reason] = fallbacks.get(reason, 0) + 1 - except Exception: - pass - return None - return unwrapped, head, _unsloth_pref_logit_transforms(unwrapped) - -def _unsloth_dpo_model_forward(model, input_ids, model_kwargs, want_hidden_states): - if not want_hidden_states: - return model(input_ids, **model_kwargs) - prior = os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES") - os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" - try: - return model(input_ids, **model_kwargs) - finally: - if prior is None: - os.environ.pop("UNSLOTH_RETURN_HIDDEN_STATES", None) - else: - os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = prior - -def _unsloth_dpo_returns_hidden_states(self, chunked, logits): - if chunked is None: - return False - unwrapped, head, _ = chunked - hidden = _unsloth_pref_returns_hidden_states(unwrapped, logits, head.weight) - if not hidden: - fallbacks = self.__dict__.setdefault("_unsloth_dpo_chunked_fallbacks", {}) - reason = "forward returned logits, not hidden states" - fallbacks[reason] = fallbacks.get(reason, 0) + 1 - return hidden - -def _unsloth_preference_chunked_off(): - return os.environ.get("UNSLOTH_PREFERENCE_CHUNKED", "1").lower() in ("0", "false", "no", "off") - -def _unsloth_odpo_chunked_head(self, model, batch, method): - """``(unwrapped, head, transforms)`` when this pass can score from hidden states, else ``None``.""" - reason = None - if _unsloth_fused_linear_logprobs is None: - reason = "unsloth_zoo has no fused_linear_logprobs" - elif os.environ.get("UNSLOTH_FUSED_LINEAR", "1") == "0": - reason = "UNSLOTH_FUSED_LINEAR=0" - elif _unsloth_preference_chunked_off(): - reason = "UNSLOTH_PREFERENCE_CHUNKED=0" - elif any(k != "attention_mask" and k != "input_ids" for k in batch): - # Vision rows (pixel_values, image_grid_thw, ...) keep TRL's layout. - reason = "extra model inputs" - if reason is not None: - try: - fallbacks = self.__dict__.setdefault("_unsloth_dpo_chunked_fallbacks", {}) - fallbacks[reason] = fallbacks.get(reason, 0) + 1 - except Exception: - pass - return None - return _unsloth_dpo_chunked_head(self, model, batch, method) - -def _unsloth_odpo_hidden_logps( - self, - chunked, - hidden_states, - completion_ids, - mask = None, -): - """TRL's ``log_softmax(head(hidden_states))[completion_ids]`` per position (float32). Rows outside ``mask`` are 0 - (the caller zeroes them anyway). A shape TRL would reject or broadcast gets real logits and TRL's own math.""" - _, head, (multiply, divide, softcap) = chunked - if hidden_states.dim() != 3 or hidden_states.shape[:2] != completion_ids.shape: - logits = torch.nn.functional.linear( - hidden_states.to(device = head.weight.device, dtype = head.weight.dtype), - head.weight, - head.bias, - ) - if multiply: - logits = logits * multiply - if divide: - logits = logits / divide - if softcap: - logits = softcap * torch.tanh(logits / softcap) - return torch.take_along_dim( - logits.log_softmax(dim = -1), completion_ids.unsqueeze(-1), dim = 2 - ).squeeze(-1) - logps = _unsloth_fused_linear_logprobs( - hidden_states, - head.weight, - completion_ids, - bias = head.bias, - mask = None if mask is None else mask.bool(), - logit_scale_multiply = multiply, - logit_scale_divide = divide, - logit_softcapping = softcap, - ) - try: - self._unsloth_dpo_chunked_calls = getattr(self, "_unsloth_dpo_chunked_calls", 0) + 1 - except Exception: - pass - return logps.to(completion_ids.device) - -try: - from unsloth_zoo.device_map_planner import detect_logit_transforms -except Exception: - detect_logit_transforms = None -try: - from unsloth_zoo.fused_losses.chunked_logprob import ( - unsloth_chunked_logprobs as _unsloth_chunked_logprobs, - unsloth_masked_logprobs as _unsloth_masked_logprobs, - ) -except Exception: - _unsloth_chunked_logprobs = _unsloth_masked_logprobs = None -try: - from unsloth_zoo.fused_losses.preference_loss import linear_mean_logits as _unsloth_linear_mean_logits -except Exception: - _unsloth_linear_mean_logits = None -try: - from unsloth_zoo.fused_losses.fused_linear import ( - fused_linear_logprobs as _unsloth_fused_linear_logprobs, - ) -except Exception: - _unsloth_fused_linear_logprobs = None -_DPO_VISION_KEYS = ('pixel_position_ids', 'image_position_ids', 'mm_token_type_ids') - -def vLLMSamplingParams(**kwargs): - from vllm import SamplingParams - - sampling_params = SamplingParams(**kwargs) - sampling_params._set_kwargs = kwargs - return sampling_params -@dataclass -class UnslothOnlineDPOConfig(OnlineDPOConfig): - """ - - Configuration class for the [`experimental.online_dpo.OnlineDPOTrainer`]. - - This class includes only the parameters that are specific to Online DPO training. For a full list of training - arguments, please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this - class may differ from those in [`~transformers.TrainingArguments`]. - - Using [`~transformers.HfArgumentParser`] we can turn this class into - [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the - command line. - - Parameters: - reward_model_path (`str`, *optional*): - Path to the reward model. - max_new_tokens (`int`, *optional*, defaults to `64`): - Maximum number of tokens to generate per completion. - max_length (`int`, *optional*, defaults to `512`): - Maximum total length of the sequence (prompt + completion) used to compute log probabilities. If the - sequence exceeds this limit, the leftmost tokens will be truncated to preserve as much of the completion as - possible. - temperature (`float`, *optional*, defaults to `0.9`): - Temperature for sampling. The higher the temperature, the more random the completions. - missing_eos_penalty (`float`, *optional*): - Penalty applied to the score when the model fails to generate an EOS token. This is useful to encourage to - generate completions shorter than the maximum length (`max_new_tokens`). The penalty must be a positive - value. - beta (`float` or `list[float]`, *optional*, defaults to `0.1`): - Parameter controlling the deviation from the reference model. Higher β means less deviation from the - reference model. For the IPO loss (`loss_type="ipo"`), β is the regularization parameter denoted by τ in - the [paper](https://huggingface.co/papers/2310.12036). If a list of floats is provided then the β is - selected for each new epoch and the last β is used for the rest of the epochs. - loss_type (`str`, *optional*, defaults to `"sigmoid"`): - Type of loss to use. Possible values are: - - - `"sigmoid"`: sigmoid loss from the original [DPO](https://huggingface.co/papers/2305.18290) paper. - - `"ipo"`: IPO loss from the [IPO](https://huggingface.co/papers/2310.12036) paper. - disable_dropout (`bool`, *optional*, defaults to `True`): - Whether to disable dropout in the model and reference model. - - > Parameters that control generation - - top_p (`float`, *optional*, defaults to `1.0`): - Float that controls the cumulative probability of the top tokens to consider. Must be in (0, 1]. Set to - `1.0` to consider all tokens. - top_k (`int`, *optional*, defaults to `0`): - Number of highest probability vocabulary tokens to keep for top-k-filtering. If `0`, top-k-filtering is - disabled and all tokens are considered. - min_p (`float`, *optional*): - Minimum token probability, which will be scaled by the probability of the most likely token. It must be a - value between `0.0` and `1.0`. Typical values are in the `0.01-0.2` range. - repetition_penalty (`float`, *optional*, defaults to `1.0`): - Float that penalizes new tokens based on whether they appear in the prompt and the generated text so far. - Values > `1.0` encourage the model to use new tokens, while values < `1.0` encourage the model to repeat - tokens. - cache_implementation (`str`, *optional*): - Implementation of the cache method for faster generation when `use_vllm` is set to `False`. - generation_kwargs (`dict[str, Any]`, *optional*): - Additional keyword arguments to pass to [`~transformers.GenerationConfig`] (if using transformers) or - `SamplingParams` (if using vLLM) when sampling completions. This can be used to further customize the - generation behavior, such as setting `suppress_tokens`, `num_beams`, etc. If it contains keys that conflict - with the other generation parameters (like `min_p`, `top_p`, etc.), they will override them. - - > Parameters that control generation acceleration powered by vLLM - - use_vllm (`bool`, *optional*, defaults to `False`): - Whether to use vLLM for generating completions. If set to `True`, the trainer will use vLLM for generation - instead of the default model.generate(). Requires `vllm` to be installed. - vllm_model_impl (`str`, *optional*, defaults to `"vllm"`): - Model implementation to use for vLLM. Must be one of `"transformers"` or `"vllm"`. `"transformers"`: Use - the `transformers` backend for model implementation. `"vllm"`: Use the `vllm` library for model - implementation. - vllm_mode (`str`, *optional*, defaults to `"colocate"`): - Mode to use for vLLM integration when `use_vllm` is set to `True`. Must be one of `"server"` or - `"colocate"`. - - - `"server"`: The trainer will send generation requests to a separate vLLM server. Make sure a vLLM server - is running (start with `vllm serve`). - - `"colocate"`: vLLM will run in the same process and share the training GPUs. This avoids the need for a - separate server but may cause resource contention with training. - vllm_structured_outputs_regex (`str`, *optional*): - Regex for vLLM structured outputs. If `None` (default), structured outputs is disabled. - - > Parameters that control the vLLM server (only used when `vllm_mode` is `"server"`) - - vllm_server_base_url (`str`, *optional*): - Base URL for the vLLM server (e.g., `"http://localhost:8000"`). If provided, `vllm_server_host` and - `vllm_server_port` are ignored. - vllm_server_host (`str`, *optional*, defaults to `"0.0.0.0"`): - Host of the vLLM server to connect to. Ignored if `vllm_server_base_url` is provided. - vllm_server_port (`int`, *optional*, defaults to `8000`): - Port of the vLLM server to connect to. Ignored if `vllm_server_base_url` is provided. - vllm_server_timeout (`float`, *optional*, defaults to `240.0`): - Total timeout duration in seconds to wait for the vLLM server to be up. If the server is not up after the - timeout, a `ConnectionError` is raised. - vllm_group_port (`int`, *optional*, defaults to `51216`): - Port number for the weight update group. This is used to communicate with the vLLM server. Unless the port - is occupied, there is no need to change it. - - > Parameters that control colocated vLLM execution (only used when `vllm_mode` is `"colocate"`) - - vllm_gpu_memory_utilization (`float`, *optional*, defaults to `0.55`): - Control the GPU memory utilization for vLLM. This setting only applies when `vllm_mode` is set to - `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when - launching the vLLM server via the `--vllm_gpu_memory_utilization` flag. - vllm_tensor_parallel_size (`int`, *optional*, defaults to `1`): - Control the tensor parallel size for vLLM. This setting only applies when `vllm_mode` is set to - `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when - launching the vLLM server via the `--vllm_tensor_parallel_size` flag. - vllm_enable_sleep_mode (`bool`, *optional*, defaults to `False`): - Enable vLLM sleep mode to offload weights/cache during the optimizer step. Keeps GPU memory usage low, but - waking the engine adds host–device transfer latency. - - > Other parameters - - ds3_gather_for_generation (`bool`, *optional*, defaults to `True`): - This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for generation, - improving generation speed. However, disabling this option allows training models that exceed the VRAM - capacity of a single GPU, albeit at the cost of slower generation. Disabling this option is not compatible - with vLLM generation. - model_init_kwargs (`dict[str, Any]`, *optional*): - Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model from a - string. The `revision` value is also used when loading reward processing classes. - trust_remote_code (`bool`, *optional*, defaults to `False`): - Whether to allow loading models that ship custom Python code from the Hub. Forwarded to - [`~transformers.AutoModelForCausalLM.from_pretrained`]. Also applied to reward-model and reward-tokenizer - loads. - reward_weights (`list[float]`, *optional*): - Weights for combining multiple reward functions. Must match the number of reward functions. If `None`, all - reward functions are equally weighted. - - > [!NOTE] - > These parameters have default values different from [`~transformers.TrainingArguments`]: - > - `logging_steps`: Defaults to `10` instead of `500`. - > - `gradient_checkpointing`: Defaults to `True` instead of `False`. - > - `bf16`: Defaults to `True` if `fp16` is not set, instead of `False`. - > - `learning_rate`: Defaults to `5e-7` instead of `5e-5`. - > - `remove_unused_columns`: Defaults to `False` instead of `True`. - - """ - _unsloth_accepts_adamw_flash = True - vllm_sampling_params: Optional[Any] = field( - default = None, - metadata = {'help': 'vLLM SamplingParams'}, - ) - unsloth_num_chunks : Optional[int] = field( - default = -1, - metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'}, - ) - unsloth_logit_chunk_multiplier : Optional[int] = field( - default = None, - metadata = {'help': 'Multiplier for chunked logit computations.'}, - ) - unsloth_grpo_mini_batch : Optional[int] = field( - default = None, - metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'}, - ) - max_seq_length : Optional[int] = field( - default = None, - metadata = {'help': 'Maximum sequence length to truncate to.'}, - ) - def __init__( - self, - output_dir = None, - per_device_train_batch_size = 4, - num_train_epochs = 3.0, - max_steps = -1, - learning_rate = 5e-05, - lr_scheduler_type = 'linear', - lr_scheduler_kwargs = None, - warmup_steps = 0.1, - optim = 'adamw_8bit', - optim_args = None, - weight_decay = 0.001, - adam_beta1 = 0.9, - adam_beta2 = 0.999, - adam_epsilon = 1e-08, - optim_target_modules = None, - gradient_accumulation_steps = 2, - average_tokens_across_devices = True, - max_grad_norm = 1.0, - label_smoothing_factor = 0.0, - bf16 = False, - fp16 = False, - bf16_full_eval = False, - fp16_full_eval = False, - tf32 = None, - gradient_checkpointing = True, - gradient_checkpointing_kwargs = None, - torch_compile = False, - torch_compile_backend = None, - torch_compile_mode = None, - use_liger_kernel = False, - liger_kernel_config = None, - use_cache = False, - neftune_noise_alpha = None, - torch_empty_cache_steps = 250, - auto_find_batch_size = False, - logging_strategy = 'steps', - logging_steps = 1, - logging_first_step = False, - log_on_each_node = True, - logging_nan_inf_filter = False, - include_num_input_tokens_seen = False, - log_level = 'passive', - log_level_replica = 'warning', - disable_tqdm = None, - report_to = 'none', - run_name = None, - project = 'huggingface', - trackio_space_id = 'trackio', - eval_strategy = 'no', - eval_steps = None, - eval_delay = 0, - per_device_eval_batch_size = 4, - prediction_loss_only = False, - eval_on_start = False, - eval_do_concat_batches = True, - eval_use_gather_object = False, - eval_accumulation_steps = 2, - batch_eval_metrics = False, - save_only_model = False, - save_strategy = 'steps', - save_steps = 500, - save_on_each_node = False, - save_total_limit = None, - enable_jit_checkpoint = False, - push_to_hub = False, - hub_token = None, - hub_private_repo = None, - hub_model_id = None, - hub_strategy = 'every_save', - hub_always_push = False, - hub_revision = None, - load_best_model_at_end = False, - metric_for_best_model = None, - greater_is_better = None, - ignore_data_skip = False, - restore_callback_states_from_checkpoint = False, - full_determinism = False, - seed = 3407, - data_seed = None, - use_cpu = False, - accelerator_config = None, - parallelism_config = None, - dataloader_drop_last = False, - dataloader_num_workers = 0, - dataloader_pin_memory = True, - dataloader_persistent_workers = False, - dataloader_prefetch_factor = None, - remove_unused_columns = False, - label_names = None, - train_sampling_strategy = 'random', - length_column_name = 'length', - ddp_find_unused_parameters = None, - ddp_bucket_cap_mb = None, - ddp_broadcast_buffers = None, - ddp_backend = None, - ddp_timeout = 1800, - fsdp = None, - fsdp_config = None, - deepspeed = None, - debug = '', - skip_memory_metrics = True, - do_train = False, - do_eval = False, - do_predict = False, - resume_from_checkpoint = None, - warmup_ratio = None, - logging_dir = None, - local_rank = -1, - reward_model_path = None, - max_new_tokens = 64, - max_length = 512, - temperature = 0.9, - top_p = 1.0, - top_k = 0, - min_p = None, - repetition_penalty = 1.0, - generation_kwargs = {}, - cache_implementation = None, - missing_eos_penalty = None, - loss_type = 'sigmoid', - disable_dropout = True, - use_vllm = False, - vllm_model_impl = 'vllm', - vllm_structured_outputs_regex = None, - vllm_gpu_memory_utilization = 0.55, - vllm_mode = 'colocate', - vllm_server_base_url = None, - vllm_server_host = '0.0.0.0', - vllm_server_port = 8000, - vllm_server_timeout = 240.0, - vllm_group_port = 51216, - vllm_tensor_parallel_size = 1, - vllm_enable_sleep_mode = False, - ds3_gather_for_generation = True, - model_init_kwargs = None, - trust_remote_code = False, - reward_weights = None, - vllm_sampling_params = None, - unsloth_num_chunks = -1, - unsloth_logit_chunk_multiplier = None, - unsloth_grpo_mini_batch = None, - max_seq_length = None, - **kwargs, - ): - if learning_rate < 1e-7: print(f'Unsloth: Your learning rate of `{learning_rate}` is too small and less than 1e-7! Consider increasing it, otherwise gradient updates will be close to 0!') - if learning_rate > 1: print(f'Unsloth: Your learning rate of `{learning_rate}` is way too larger > 1! Consider decreasing it to 1e-1, otherwise gradient updates will explode!') - _unsloth_fsdp2 = os.environ.get('FSDP_VERSION', '').strip() == '2' - _unsloth_fsdp_cfg = fsdp_config - if fsdp and isinstance(_unsloth_fsdp_cfg, str) and os.path.isfile(_unsloth_fsdp_cfg): - try: - import json as _unsloth_json - with open(_unsloth_fsdp_cfg, encoding = 'utf-8') as _unsloth_f: - _unsloth_fsdp_cfg = _unsloth_json.load(_unsloth_f) - except Exception: - _unsloth_fsdp_cfg = None - if fsdp and _unsloth_fsdp_cfg is None: - _unsloth_fsdp_cfg = {} - if fsdp and isinstance(_unsloth_fsdp_cfg, dict): - _unsloth_fsdp2 = _unsloth_fsdp2 or str(_unsloth_fsdp_cfg.get('fsdp_version', _unsloth_fsdp_cfg.get('version', 1))) == '2' - _unsloth_flash_adamw = str(getattr(optim, 'value', optim)).strip().lower() == 'adamw_flash' - if _unsloth_flash_adamw: optim = 'adamw_torch' - if _unsloth_fsdp2 and str(getattr(optim, 'value', optim)) == 'adamw_8bit': - print("Unsloth: FSDP2 shards parameters as DTensors, which bitsandbytes' 8-bit optimizers cannot step. Switching optim = 'adamw_8bit' to 'adamw_torch_fused'.") - optim = 'adamw_torch_fused' - if num_train_epochs is None: - num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override - if output_dir is None and save_strategy == 'steps' and save_steps == 500: - output_dir = 'unsloth_training_checkpoints' - save_strategy = 'no' - if eval_steps is not None and eval_strategy != 'steps': - print(f'Unsloth: `eval_steps = {eval_steps}` is ignored because `eval_strategy` is {getattr(eval_strategy, "value", eval_strategy)!r}. Set `eval_strategy = "steps"` to evaluate every `eval_steps` steps.') - - if temperature <= 0: - raise ValueError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.') - elif temperature >= 10: - raise ValueError('Unsloth: Please set a positive non-zero temperature less than 10, since sampling will be quite erratic.') - - - # One dict so the filter sees the mirrored parameters AND `**kwargs`: - # filtering kwargs alone would double-bind any argument TRL renamed, - # since the new name is itself a mirrored parameter. - _unsloth_config_arguments = dict( - output_dir = output_dir, - per_device_train_batch_size = per_device_train_batch_size, - num_train_epochs = num_train_epochs, - max_steps = max_steps, - learning_rate = learning_rate, - lr_scheduler_type = lr_scheduler_type, - lr_scheduler_kwargs = lr_scheduler_kwargs, - warmup_steps = warmup_steps, - optim = optim, - optim_args = optim_args, - weight_decay = weight_decay, - adam_beta1 = adam_beta1, - adam_beta2 = adam_beta2, - adam_epsilon = adam_epsilon, - optim_target_modules = optim_target_modules, - gradient_accumulation_steps = gradient_accumulation_steps, - average_tokens_across_devices = average_tokens_across_devices, - max_grad_norm = max_grad_norm, - label_smoothing_factor = label_smoothing_factor, - bf16 = bf16, - fp16 = fp16, - bf16_full_eval = bf16_full_eval, - fp16_full_eval = fp16_full_eval, - tf32 = tf32, - gradient_checkpointing = gradient_checkpointing, - gradient_checkpointing_kwargs = gradient_checkpointing_kwargs, - torch_compile = torch_compile, - torch_compile_backend = torch_compile_backend, - torch_compile_mode = torch_compile_mode, - use_liger_kernel = use_liger_kernel, - liger_kernel_config = liger_kernel_config, - use_cache = use_cache, - neftune_noise_alpha = neftune_noise_alpha, - torch_empty_cache_steps = torch_empty_cache_steps, - auto_find_batch_size = auto_find_batch_size, - logging_strategy = logging_strategy, - logging_steps = logging_steps, - logging_first_step = logging_first_step, - log_on_each_node = log_on_each_node, - logging_nan_inf_filter = logging_nan_inf_filter, - include_num_input_tokens_seen = include_num_input_tokens_seen, - log_level = log_level, - log_level_replica = log_level_replica, - disable_tqdm = disable_tqdm, - report_to = report_to, - run_name = run_name, - project = project, - trackio_space_id = trackio_space_id, - eval_strategy = eval_strategy, - eval_steps = eval_steps, - eval_delay = eval_delay, - per_device_eval_batch_size = per_device_eval_batch_size, - prediction_loss_only = prediction_loss_only, - eval_on_start = eval_on_start, - eval_do_concat_batches = eval_do_concat_batches, - eval_use_gather_object = eval_use_gather_object, - eval_accumulation_steps = eval_accumulation_steps, - batch_eval_metrics = batch_eval_metrics, - save_only_model = save_only_model, - save_strategy = save_strategy, - save_steps = save_steps, - save_on_each_node = save_on_each_node, - save_total_limit = save_total_limit, - enable_jit_checkpoint = enable_jit_checkpoint, - push_to_hub = push_to_hub, - hub_token = hub_token, - hub_private_repo = hub_private_repo, - hub_model_id = hub_model_id, - hub_strategy = hub_strategy, - hub_always_push = hub_always_push, - hub_revision = hub_revision, - load_best_model_at_end = load_best_model_at_end, - metric_for_best_model = metric_for_best_model, - greater_is_better = greater_is_better, - ignore_data_skip = ignore_data_skip, - restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint, - full_determinism = full_determinism, - seed = seed, - data_seed = data_seed, - use_cpu = use_cpu, - accelerator_config = accelerator_config, - parallelism_config = parallelism_config, - dataloader_drop_last = dataloader_drop_last, - dataloader_num_workers = dataloader_num_workers, - dataloader_pin_memory = dataloader_pin_memory, - dataloader_persistent_workers = dataloader_persistent_workers, - dataloader_prefetch_factor = dataloader_prefetch_factor, - remove_unused_columns = remove_unused_columns, - label_names = label_names, - train_sampling_strategy = train_sampling_strategy, - length_column_name = length_column_name, - ddp_find_unused_parameters = ddp_find_unused_parameters, - ddp_bucket_cap_mb = ddp_bucket_cap_mb, - ddp_broadcast_buffers = ddp_broadcast_buffers, - ddp_backend = ddp_backend, - ddp_timeout = ddp_timeout, - fsdp = fsdp, - fsdp_config = fsdp_config, - deepspeed = deepspeed, - debug = debug, - skip_memory_metrics = skip_memory_metrics, - do_train = do_train, - do_eval = do_eval, - do_predict = do_predict, - resume_from_checkpoint = resume_from_checkpoint, - warmup_ratio = warmup_ratio, - logging_dir = logging_dir, - local_rank = local_rank, - reward_model_path = reward_model_path, - max_new_tokens = max_new_tokens, - max_length = max_length, - temperature = temperature, - top_p = top_p, - top_k = top_k, - min_p = min_p, - repetition_penalty = repetition_penalty, - generation_kwargs = generation_kwargs, - cache_implementation = cache_implementation, - missing_eos_penalty = missing_eos_penalty, - loss_type = loss_type, - disable_dropout = disable_dropout, - use_vllm = use_vllm, - vllm_model_impl = vllm_model_impl, - vllm_structured_outputs_regex = vllm_structured_outputs_regex, - vllm_gpu_memory_utilization = vllm_gpu_memory_utilization, - vllm_mode = vllm_mode, - vllm_server_base_url = vllm_server_base_url, - vllm_server_host = vllm_server_host, - vllm_server_port = vllm_server_port, - vllm_server_timeout = vllm_server_timeout, - vllm_group_port = vllm_group_port, - vllm_tensor_parallel_size = vllm_tensor_parallel_size, - vllm_enable_sleep_mode = vllm_enable_sleep_mode, - ds3_gather_for_generation = ds3_gather_for_generation, - model_init_kwargs = model_init_kwargs, - trust_remote_code = trust_remote_code, - reward_weights = reward_weights,**kwargs) - super().__init__(**_unsloth_filter_config_init_kwargs(OnlineDPOConfig, _unsloth_config_arguments, mirrored_from = __class__)) - self.vllm_sampling_params = vllm_sampling_params - self.unsloth_num_chunks = unsloth_num_chunks - if unsloth_grpo_mini_batch is not None: - if self.generation_batch_size >= unsloth_grpo_mini_batch: - self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch - else: - raise ValueError( - f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, " - f"which is self.per_device_train_batch_size * gradient_accumulation_steps." - ) - self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier - self.max_seq_length = max_seq_length - # Unsloth: keep the reentrant checkpoint path - if getattr(self, 'gradient_checkpointing', False): - _gc_kwargs = getattr(self, 'gradient_checkpointing_kwargs', None) or {} - if _gc_kwargs.get('context_fn') is None and not _gc_kwargs.get('debug', False): - _gc_kwargs['use_reentrant'] = True - self.gradient_checkpointing_kwargs = _gc_kwargs - from unsloth.optimizers.flash_adamw import is_flash_adamw_marked, mark_flash_adamw, unmark_flash_adamw - if _unsloth_flash_adamw or is_flash_adamw_marked(self): - mark_flash_adamw(self) - else: - unmark_flash_adamw(self) - -pass - -class _UnslothOnlineDPOTrainer(_BaseTrainer): - r"""""" - - _tag_names = ["trl", "online-dpo"] - _name = "Online DPO" - _paper = { - "title": "Direct Language Model Alignment from Online AI Feedback", - "id": "2402.04792", - # docstyle-ignore - "citation": textwrap.dedent("""\ - @article{guo2024direct, - title = {{Direct Language Model Alignment from Online AI Feedback}}, - author = {Shangmin Guo and Biao Zhang and Tianlin Liu and Tianqi Liu and Misha Khalman and Felipe Llinares and Alexandre Ram{\'{e}} and Thomas Mesnard and Yao Zhao and Bilal Piot and Johan Ferret and Mathieu Blondel}, - year = 2024, - eprint = {arXiv:2402.04792} - }"""), - } - - def __init__( - self, - model: PreTrainedModel | nn.Module | str, - ref_model: PreTrainedModel | nn.Module | None = None, - reward_funcs: RewardFunc | list[RewardFunc] | None = None, - args: OnlineDPOConfig | None = None, - data_collator: DataCollator | None = None, - train_dataset: Dataset | IterableDataset | None = None, - eval_dataset: Dataset | IterableDataset | dict[str, Dataset | IterableDataset] | None = None, - processing_class: PreTrainedTokenizerBase | ProcessorMixin | None = None, - reward_processing_classes: PreTrainedTokenizerBase | list[PreTrainedTokenizerBase] | None = None, - peft_config: "PeftConfig | None" = None, - compute_metrics: Callable[[EvalPrediction], dict] | None = None, - callbacks: list[TrainerCallback] | None = None, - optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), - preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, - ) -> None: - - if hasattr(model, 'vllm_engine') and hasattr(args, 'use_vllm'): - _unsloth_fsdp_cfg = getattr(args, 'fsdp_config', None) - if getattr(args, 'fsdp', None) and isinstance(_unsloth_fsdp_cfg, dict) and str(_unsloth_fsdp_cfg.get('version', _unsloth_fsdp_cfg.get('fsdp_version', 1))) == '2': - raise NotImplementedError('Unsloth: `fast_inference = True` (vLLM) is not supported with FSDP2.\nLoad the model with `fast_inference = False` to train with FSDP2, or launch without FSDP2 to use vLLM.') - if (getattr(args, 'use_vllm', False) == False): - args.use_vllm = True - if getattr(args, 'top_k', -1) is None or getattr(args, 'top_k', -1) == 0: - args.top_k = -1 - if train_dataset is None: - raise ValueError("`train_dataset` is required") - - if ref_model is model: - raise ValueError( - "`model` and `ref_model` cannot be the same object. If you want `ref_model` to be the " - "same as `model`, either omit the `ref_model` argument or pass `None`." - ) - - self.ref_model = ref_model - - if reward_funcs is None: - raise ValueError("`reward_funcs` must be provided.") - - # Handle reward_funcs - if not isinstance(reward_funcs, list): - reward_funcs = [reward_funcs] - self.reward_func_names = [] - - # Process reward functions [convert strings to models, collect names] - model_init_kwargs = args.model_init_kwargs or {} - model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) - reward_model_revisions = [None] * len(reward_funcs) - for i, reward_func in enumerate(reward_funcs): - if isinstance(reward_func, str): - # Load model from string path - reward_model_revisions[i] = model_init_kwargs.get("revision") - reward_funcs[i] = AutoModelForSequenceClassification.from_pretrained( - reward_func, num_labels=1, **model_init_kwargs - ) - if isinstance(reward_funcs[i], nn.Module): - self.reward_func_names.append(get_config_model_id(reward_funcs[i].config).split("/")[-1]) - else: - self.reward_func_names.append(get_callable_name(reward_funcs[i])) - self.reward_funcs = reward_funcs - - # Handle reward processing classes for reward_funcs - if reward_processing_classes is None: - reward_processing_classes = [None] * len(reward_funcs) - elif not isinstance(reward_processing_classes, list): - reward_processing_classes = [reward_processing_classes] - else: - if len(reward_processing_classes) != len(reward_funcs): - raise ValueError("The number of reward processing classes must match the number of reward functions.") - - self.reward_processing_classes = [] - for i, (reward_processing_class_i, reward_func) in enumerate( - zip(reward_processing_classes, reward_funcs, strict=True) - ): - if isinstance(reward_func, PreTrainedModel): - if reward_processing_class_i is None: - reward_processing_class_i = AutoTokenizer.from_pretrained( - reward_func.config._name_or_path, - revision=reward_model_revisions[i], - trust_remote_code=args.trust_remote_code, - ) - if reward_processing_class_i.pad_token_id is None: - reward_processing_class_i.pad_token = reward_processing_class_i.eos_token - # Set pad token ID on reward model config - reward_func.config.get_text_config().pad_token_id = reward_processing_class_i.pad_token_id - self.reward_processing_classes.append(reward_processing_class_i) - - # Handle reward_weights - if args.reward_weights is not None: - if len(args.reward_weights) != len(self.reward_funcs): - raise ValueError( - f"Number of reward weights ({len(args.reward_weights)}) must match number of reward " - f"functions ({len(self.reward_funcs)})" - ) - self.reward_weights = torch.tensor(args.reward_weights, dtype=torch.float32) - else: - self.reward_weights = torch.ones(len(self.reward_funcs), dtype=torch.float32) - - if args is None: - raise ValueError("`args` must be provided.") - - # Check that the processing_class is provided - if processing_class is None: - raise ValueError("`processing_class` must be provided.") - - # PEFT initializes the adapter weights randomly, so set_seed must be done before creating the model to ensure - # reproducibility. - set_seed(args.seed) - model_init_kwargs = args.model_init_kwargs or {} - if isinstance(model, str): - model_id = model - - # Handle dtype in model_init_kwargs - dtype = model_init_kwargs.get("dtype", "auto") - if isinstance(dtype, torch.dtype) or dtype == "auto" or dtype is None: - pass - elif isinstance(dtype, str): - dtype = getattr(torch, dtype) - model_init_kwargs["dtype"] = dtype - else: - raise ValueError( - "Invalid `dtype` passed to `OnlineDPOConfig`. Expected either 'auto' or a string " - f"representing a `torch.dtype` (e.g., 'float32'), but got {dtype}." - ) - model_init_kwargs["device_map"] = model_init_kwargs.get("device_map", "auto") - model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) - - model = AutoModelForCausalLM.from_pretrained(model_id, **model_init_kwargs) - else: - if args.model_init_kwargs is not None: - raise ValueError( - "You passed `model_init_kwargs` to the `OnlineDPOConfig`, but your model is already instantiated. " - "This argument can only be used when the `model` argument is a string." - ) - self.is_encoder_decoder = model.config.is_encoder_decoder - self.is_vision_model = model.config.model_type in MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES.keys() - - # PEFT - if False: - if not is_peft_available(): - raise ImportError( - "You passed `peft_config` but the `peft` library is not installed. " - "Install it with `pip install trl[peft]`." - ) - if not isinstance(peft_config, PeftConfig): - raise TypeError( - f"`peft_config` must be a `peft.PeftConfig` instance (e.g. `peft.LoraConfig`), " - f"got {type(peft_config).__name__}." - ) - if peft_config is not None or is_peft_model(model): - pass - - # Enable gradient checkpointing if requested - if args.gradient_checkpointing: - model = self._enable_gradient_checkpointing(model, args) - - # Disable dropout in the model and reference model - if args.disable_dropout: - disable_dropout_in_model(model) - if self.ref_model is not None: - disable_dropout_in_model(self.ref_model) - - # Handle the ref_model - # Usually, the user wants the ref model to be the initial version of the model. When using PEFT, it's easy to - # get the ref model, as it's just the model with a disabled adapter. When not using PEFT, we need to create - # the ref model from the model by copying it and disable the gradients and set it in evaluation mode. - if ref_model is None: # No ref model provided, the most common case - if False: - self.ref_model = create_reference_model(model) # copy, disable gradients, set eval mode - else: - self.ref_model = None # we don't need a ref model here, we can just disable the adapter. - else: # rare case, the user provided a ref model - self.ref_model = ref_model - self.ref_model.eval() - - # Disable the gradient and set the reward model in eval mode - if reward_funcs is not None: - for reward_func in reward_funcs: - if isinstance(reward_func, PreTrainedModel): - reward_func.eval() - - self.max_length = args.max_length - - self.stats = { - "objective/kl": [], - "objective/entropy": [], - "objective/non_score_reward": [], - "rewards/chosen": [], - "rewards/rejected": [], - "rewards/accuracies": [], - "rewards/margins": [], - "logps/chosen": [], - "logps/rejected": [], - "val/contain_eos_token": [], - "beta": [], - } - if self.reward_funcs is not None: - self.stats["objective/rlhf_reward"] = [] - self.stats["objective/scores_margin"] = [] - self.stats["objective/scores"] = [] - - # Store generation parameters for later use - self.use_vllm = args.use_vllm - self.num_generations = 2 # Generate 2 completions per prompt for Online DPO - self.temperature = args.temperature - self.top_p = args.top_p - self.top_k = args.top_k - self.min_p = args.min_p - self.repetition_penalty = args.repetition_penalty - self.vllm_mode = args.vllm_mode if args.use_vllm else None - self.vllm_gpu_memory_utilization = args.vllm_gpu_memory_utilization - self.vllm_tensor_parallel_size = args.vllm_tensor_parallel_size - self.vllm_model_impl = args.vllm_model_impl - - # Handle pad token for processors or tokenizers - if isinstance(processing_class, ProcessorMixin): - self._tokenizer = processing_class.tokenizer - elif isinstance(processing_class, PreTrainedTokenizerBase): - self._tokenizer = processing_class - else: - raise TypeError("The `processing_class` must be either a `PreTrainedTokenizerBase` or a `ProcessorMixin`") - - if self._tokenizer.pad_token is None: - self._tokenizer.pad_token = self._tokenizer.eos_token - # The model must agree with the tokenizer on the pad token from construction, so mirror it onto the model - # configs. - model.config.get_text_config().pad_token_id = self._tokenizer.pad_token_id - model.generation_config.pad_token_id = self._tokenizer.pad_token_id - - # Every eos id the model declares [e.g. an end-of-turn token] ends a completion, not only the tokenizer's - eos_token_ids = model.generation_config.eos_token_id - if eos_token_ids is None: - eos_token_ids = [] - elif isinstance(eos_token_ids, int): - eos_token_ids = [eos_token_ids] - if self._tokenizer.eos_token_id not in eos_token_ids: - eos_token_ids = [self._tokenizer.eos_token_id, *eos_token_ids] - self.eos_token_ids = eos_token_ids - - # Vision tokens for VLM support - self.image_token_id = getattr(processing_class, "image_token_id", None) - self.vision_start_token_id = getattr(processing_class, "vision_start_token_id", None) - self.vision_end_token_id = getattr(processing_class, "vision_end_token_id", None) - # Get the image token string for token collapsing - self.image_token = None - if self.image_token_id is not None: - self.image_token = self._tokenizer.decode([self.image_token_id]) - - # Define the collator if not provided - if data_collator is None: - data_collator = DPODataCollatorWithPadding(pad_token_id=self._tokenizer.pad_token_id) - - # Transformers explicitly set use_reentrant=True in the past to silence a PyTorch warning, but the default was - # never updated once PyTorch switched to recommending use_reentrant=False. Until that change lands upstream - # [see https://github.com/huggingface/transformers/pull/43203] and is released [most likely in 5.0.0], we - # default to the recommended non-reentrant behavior here, while preserving any user-provided value. - if args.gradient_checkpointing and Version(transformers.__version__) < Version("5.0.0"): - args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} - args.gradient_checkpointing_kwargs.setdefault("use_reentrant", False) - - super().__init__( - model=model, - args=args, - data_collator=data_collator, - train_dataset=train_dataset, - eval_dataset=eval_dataset, - processing_class=processing_class, - compute_metrics=compute_metrics, - callbacks=callbacks, - optimizers=optimizers, - preprocess_logits_for_metrics=preprocess_logits_for_metrics, - ) - - # Add tags for models that have been loaded with the correct transformers version - if hasattr(self.model, "add_model_tags"): - self.model.add_model_tags(self._tag_names) - - self._beta = args.beta - - # Set up generation configuration and vLLM after super[].__init__ - if self.use_vllm: - if getattr(getattr(model, 'vllm_engine', None), 'shared_weights', False): self.llm = model.vllm_engine; self._last_loaded_step = 0 - self.generation_config = SamplingParams(**generation_kwargs, - **getattr(getattr(args, 'vllm_sampling_params', vLLMSamplingParams()), '_set_kwargs', {}), - ) - else: - # Set up transformers generation config - generation_kwargs = { - "max_new_tokens": args.max_new_tokens, - "do_sample": True, - "pad_token_id": self._tokenizer.pad_token_id, - "bos_token_id": self._tokenizer.bos_token_id, - "eos_token_id": self.eos_token_ids, - "temperature": self.temperature, - "top_k": self.top_k, - "top_p": self.top_p, - "repetition_penalty": self.repetition_penalty, - "use_cache": True if not self.args.gradient_checkpointing else False, - } - # Add min_p if supported - if self.min_p is not None: - generation_kwargs["min_p"] = self.min_p - if args.generation_kwargs is not None: - generation_kwargs.update(args.generation_kwargs) - # Remove None values - generation_kwargs = {k: v for k, v in generation_kwargs.items() if v is not None} - self.generation_config = GenerationConfig(**generation_kwargs) - # Keep training-specific generation kwargs to overwrite model's original generation config - self.generation_kwargs = generation_kwargs - - if self.ref_model is not None: - if self.is_deepspeed_enabled: - self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator) - elif self.is_fsdp_enabled: - self.ref_model = prepare_fsdp(self.ref_model, self.accelerator) - else: - self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True) - if self.reward_funcs is not None: - for i, reward_func in enumerate(self.reward_funcs): - if isinstance(reward_func, PreTrainedModel): - if self.is_deepspeed_enabled: - self.reward_funcs[i] = prepare_deepspeed(reward_func, self.accelerator) - else: - # set device placement to True to make `prepare_model` move `reward_func` to device when using fsdp - self.reward_funcs[i] = self.accelerator.prepare_model( - reward_func, evaluation_mode=True, device_placement=True - ) - - @property - def beta(self): - if isinstance(self._beta, list): - epoch = self.state.epoch - return self._beta[epoch] if epoch < len(self._beta) else self._beta[-1] - else: - return self._beta - - @staticmethod - def tokenize_row(feature, is_encoder_decoder: bool, tokenizer: PreTrainedTokenizerBase) -> dict[str, Any]: - """Tokenize a single row from a DPO specific dataset.""" - if not is_encoder_decoder: - batch = tokenizer(feature["prompt"], add_special_tokens=False) - # Add BOS token to head of prompt. Avoid adding if it's already there - if tokenizer.bos_token_id is not None: - prompt_len_input_ids = len(batch["input_ids"]) - if prompt_len_input_ids == 0 or tokenizer.bos_token_id != batch["input_ids"][0]: - batch["input_ids"] = [tokenizer.bos_token_id] + batch["input_ids"] - batch["attention_mask"] = [1] + batch["attention_mask"] - else: - batch = tokenizer(feature["prompt"], add_special_tokens=True) - batch = {f"prompt_{key}": value for key, value in batch.items()} - return batch - - def _enable_gradient_checkpointing(self, model: PreTrainedModel, args: OnlineDPOConfig) -> PreTrainedModel: - """Enables gradient checkpointing for the model.""" - # Ensure use_cache is disabled - model.config.use_cache = False - - # Enable gradient checkpointing on the base model for PEFT - if is_peft_model(model): - model.base_model.gradient_checkpointing_enable() - # Enable gradient checkpointing for non-PEFT models - else: - model.gradient_checkpointing_enable() - - model.enable_input_require_grads() - return model - - def _generate_vllm(self, prompts, images=None): - pad_token_id = self._tokenizer.pad_token_id - - # Generate completion_ids and prompt_ids based on mode - if self.vllm_mode == "server": - completion_ids, prompt_ids = self._generate_vllm_server(prompts, images) - elif self.vllm_mode == "colocate": - completion_ids, prompt_ids = self._generate_vllm_colocate(prompts, images) - - # Shared padding, masking, and tensor conversion logic - max_prompt_length = max(len(ids) for ids in prompt_ids) - prompt_mask = [[0] * (max_prompt_length - len(ids)) + [1] * len(ids) for ids in prompt_ids] - prompt_ids = [[pad_token_id] * (max_prompt_length - len(ids)) + ids for ids in prompt_ids] - max_tokens = self.generation_config.max_tokens - completion_mask = [[1] * len(ids) + [0] * (max_tokens - len(ids)) for ids in completion_ids] - completion_ids = [ - ids + [self.eos_token_ids[0]] if ids[-1] not in self.eos_token_ids and len(ids) < max_tokens else ids - for ids in completion_ids - ] - completion_ids = [ids + [pad_token_id] * (max_tokens - len(ids)) for ids in completion_ids] - - # Convert to tensors - prompt_ids = torch.tensor(prompt_ids, device=self.accelerator.device) - prompt_mask = torch.tensor(prompt_mask, device=self.accelerator.device) - completion_ids = torch.tensor(completion_ids, device=self.accelerator.device) - completion_mask = torch.tensor(completion_mask, device=self.accelerator.device) - - return prompt_ids, prompt_mask, completion_ids, completion_mask - - def _generate_vllm_server(self, prompts, images=None): - """Generate completions using vLLM server mode""" - has_images = images is not None - - # Update vLLM server weights if needed - if hasattr(self, "_last_loaded_step") and self.state.global_step != self._last_loaded_step: - self._move_model_to_vllm() - self._last_loaded_step = self.state.global_step - elif not hasattr(self, "_last_loaded_step"): - self._move_model_to_vllm() - self._last_loaded_step = self.state.global_step - - # Apply chat template if conversational - if is_conversational({"prompt": prompts[0]}): - prompts_text = [apply_chat_template({"prompt": p}, self.processing_class)["prompt"] for p in prompts] - else: - prompts_text = prompts - # Gather all prompts to main process - all_prompts = gather_object(prompts_text) - if has_images: - # The server can't take images alongside text prompts, so multimodal prompts are sent as messages, with - # the images inlined in place of their placeholders. - messages = [ - prepare_multimodal_messages(prompt, images=[image] if image is not None else None) - for prompt, image in zip(prompts, images, strict=True) - ] - all_messages = gather_object(messages) - - if self.accelerator.is_main_process: - sampling_kwargs = { - "n": self.num_generations, - "repetition_penalty": self.repetition_penalty, - "temperature": self.temperature, - "top_p": self.top_p, - "top_k": -1 if self.top_k is None else self.top_k, - "min_p": 0.0 if self.min_p is None else self.min_p, - "max_tokens": self.generation_config.max_tokens, - "structured_outputs_regex": self.structured_outputs_regex - if hasattr(self, "structured_outputs_regex") - else None, - "generation_kwargs": self.args.generation_kwargs, - } - if has_images: - completion_ids = self.vllm_client.chat(messages=all_messages, **sampling_kwargs)["completion_ids"] - else: - completion_ids = self.vllm_client.generate(prompts=all_prompts, **sampling_kwargs)["completion_ids"] - else: - completion_ids = [None] * (len(all_prompts) * 2) - - # Broadcast completions to all processes - completion_ids = broadcast_object_list(completion_ids, from_process=0) - - # Slice to keep only the local part of the data - process_slice = slice( - self.accelerator.process_index * len(prompts) * 2, - (self.accelerator.process_index + 1) * len(prompts) * 2, - ) - completion_ids = completion_ids[process_slice] - # Reorder to block layout ([p0c0, p1c0, ..., p0c1, p1c1, ...]) to match the colocate and transformers - # generation paths, which is what the loss expects (`rewards.split(batch_size)`). - completion_ids = completion_ids[0::2] + completion_ids[1::2] - - # Create prompt_ids by tokenizing locally, in the same block layout (2 copies per prompt) - prompt_inputs = self.processing_class( - text=prompts_text, - return_tensors="pt", - padding=True, - padding_side="left", - add_special_tokens=False, - ) - prompt_ids = [prompt_tokens.tolist() for prompt_tokens in prompt_inputs["input_ids"]] - prompt_ids = prompt_ids + prompt_ids # 2 copies for 2 completions - return completion_ids, prompt_ids - - def _generate_vllm_colocate(self, prompts, images=None): - """Generate completions using vLLM colocate mode""" - if self.args.vllm_enable_sleep_mode: - # wake up colocated vLLM instances if needed - empty_cache() # required to avoid OOM in some cases - self.llm.wake_up(tags=["weights"]) - - # Update model weights if needed - only after gradient accumulation completes - if self.state.global_step != self._last_loaded_step: - self._move_model_to_vllm() - self._last_loaded_step = self.state.global_step - - # Apply chat template if conversational - if is_conversational({"prompt": prompts[0]}): - prompts_text = [apply_chat_template({"prompt": p}, self.processing_class)["prompt"] for p in prompts] - else: - prompts_text = prompts - - # Prepare vLLM inputs with images if available - if images is not None: - vllm_inputs = [] - for prompt, image in zip(prompts_text, images, strict=True): - if image is not None: - vllm_inputs.append({"prompt": prompt, "multi_modal_data": {"image": image}}) - else: - vllm_inputs.append(prompt) - else: - vllm_inputs = prompts_text - - if self.args.vllm_enable_sleep_mode: - self.llm.wake_up(tags=["kv_cache"]) - - outputs = self.llm.generate(vllm_inputs, self.generation_config, use_tqdm=False, lora_request = self.model.load_lora('online_dpo_trainer_lora_model', load_tensors = True) if getattr(self.llm, 'shared_weights', False) else None) - - completion_ids = [list(output.outputs[i].token_ids) for i in range(2) for output in outputs] - prompt_ids = [list(output.prompt_token_ids) for _ in range(2) for output in outputs] - if self.args.vllm_enable_sleep_mode: - self.llm.sleep(level=2) - - return completion_ids, prompt_ids - - def _sync_fsdp2_params_to_vllm(self, module: nn.Module): - # For FSDP2, module already covers all parameters, so no need for recursion - for name, param in module.items(): - # When using PEFT, we need to recover the original parameter name - name = name.removeprefix("base_model.model.").replace(".base_layer", "") - # Skip PEFT layers: they don’t exist in vLLM, and they are merged already. - if is_peft_model(module) and module.prefix in name: - continue - # When module to save, remove its prefix and discard the original module - if "original_module" in name: - continue - name = self._fix_param_name_to_vllm(name, extra_prefixes=["modules_to_save.default."]) - - if param.is_cpu: - param = param.to(self.accelerator.device) - param = param.full_tensor() - - if self.vllm_mode == "server" and self.accelerator.is_main_process: - self.vllm_client.update_named_param(name, param) - elif self.vllm_mode == "colocate": - - pass - - pass - - def _move_model_to_vllm(self): - if self.vllm_mode == "server" and self.accelerator.is_main_process: - # Announce one weight update for the whole model: the server prepares and finalizes it once, rather than - # once per tensor pushed below. Only the main process holds a client; the other ranks only take part in - # the parameter gathers below. - with self.vllm_client.weight_update(): - self._move_model_to_vllm_inner() - else: - self._move_model_to_vllm_inner() - - # Reset cache on vLLM - if self.vllm_mode == "server" and self.accelerator.is_main_process: - self.vllm_client.reset_prefix_cache() - elif self.vllm_mode == "colocate": - self.llm.reset_prefix_cache() - - def _move_model_to_vllm_inner(self): - # For DeepSpeed ZeRO-3 and FSDP, we need to gather all parameters before operations - deepspeed_plugin = self.accelerator.state.deepspeed_plugin - zero_stage_3 = deepspeed_plugin is not None and deepspeed_plugin.zero_stage == 3 - if zero_stage_3: - import deepspeed - - gather_if_zero3 = deepspeed.zero.GatheredParameters - else: - gather_if_zero3 = nullcontext - - if is_peft_model(self.model): - # With PEFT and FSDP/DeepSpeed ZeRO Stage 3, we must gather the full model at once before merging, as - # merging adapters in a sharded manner is not supported. - # TODO: does this work with FSDP? - with gather_if_zero3(list(self.model.parameters())): - self.model.merge_adapter() - - # Update vLLM weights while parameters are gathered - if self.is_fsdp_enabled: # note if using FSDP, gather_if_zero3 is nullcontext - # Update vLLM weights while parameters are gathered - # For PEFT with FSDP we need to use the memory efficient post-order traversal - fsdp_plugin = getattr(self.accelerator.state, "fsdp_plugin", None) - fsdp_version = getattr(fsdp_plugin, "fsdp_version", 1) if fsdp_plugin else 1 - if fsdp_version == 1: - self._sync_fsdp1_params_to_vllm( - self.model - ) # use memory-efficient post-order traversal for FSDP - elif fsdp_version == 2: - self._sync_fsdp2_params_to_vllm(self.model) - else: - # DeepSpeed ZeRO-3 with PEFT - for name, param in self.model.named_parameters(): - # When using PEFT, we need to recover the original parameter name - name = name.removeprefix("base_model.model.").replace(".base_layer", "") - # Skip PEFT layers: they don’t exist in vLLM, and they are merged already. - if self.model.prefix in name: - continue - # When module to save, remove its prefix and discard the original module - if "original_module" in name: - continue - name = self._fix_param_name_to_vllm(name, extra_prefixes=["modules_to_save.default."]) - - if self.vllm_mode == "server" and self.accelerator.is_main_process: - self.vllm_client.update_named_param(name, param.data) - elif self.vllm_mode == "colocate": - - pass - - pass - # Unmerge adapters while parameters are still gathered - self.model.unmerge_adapter() - # Parameters will automatically be repartitioned when exiting the context - else: - # For non-PEFT models, simply gather (if needed) and update each parameter individually. - if self.is_fsdp_enabled: - fsdp_plugin = getattr(self.accelerator.state, "fsdp_plugin", None) - fsdp_version = getattr(fsdp_plugin, "fsdp_version", 1) if fsdp_plugin else 1 - if fsdp_version == 1: - self._sync_fsdp1_params_to_vllm(self.model) # use memory-efficient post-order traversal for FSDP - elif fsdp_version == 2: - self._sync_fsdp2_params_to_vllm(self.model) - else: - for name, param in self.model.named_parameters(): - name = self._fix_param_name_to_vllm(name) - with gather_if_zero3([param]): - if self.vllm_mode == "server" and self.accelerator.is_main_process: - self.vllm_client.update_named_param(name, param.data) - elif self.vllm_mode == "colocate": - - pass - - pass - - def _sync_fsdp1_params_to_vllm(self, module: nn.Module, prefix: str = "", visited=None): - """Memory-efficient post-order traversal of FSDP modules to extract full parameters and sync with vLLM.""" - # For FSDP1, we need to recurse into children and also use summon_full_params - if visited is None: - visited = set() - for child_name, child_module in module.named_children(): - child_prefix = f"{prefix}.{child_name}" if prefix else child_name - self._sync_fsdp1_params_to_vllm( - child_module, prefix=child_prefix, visited=visited - ) # recurse into the child - - if isinstance(module, FSDP): - with FSDP.summon_full_params(module, recurse=False, writeback=False): - for param_name, param in module.named_parameters(): - full_name = f"{prefix}.{param_name}" if prefix else param_name - full_name = self._fix_param_name_to_vllm(full_name, extra_prefixes=["_fsdp_wrapped_module."]) - - if full_name in visited: - continue # skip FSDP subtrees already traversed - visited.add(full_name) - - if self.vllm_mode == "server" and self.accelerator.is_main_process: - self.vllm_client.update_named_param(full_name, param.data) - elif self.vllm_mode == "colocate": - - pass - - pass - - def _fix_param_name_to_vllm(self, name, extra_prefixes: list[str] | None = None): - """Clean parameter names for vLLM compatibility""" - extra_prefixes = extra_prefixes or [] - prefixes = ["_checkpoint_wrapped_module."] + extra_prefixes - for prefix in prefixes: - name = name.replace(prefix, "") - return name - - def process_vision_row( - self, features: dict[str, list | torch.Tensor], processing_class=None - ) -> dict[str, list[int]]: - """ - Process a vision row for VLM models (adapted from DPO trainer) - """ - processor = processing_class or self.processing_class - processed_features = processor(images=[features["image"]], text=features["prompt"], add_special_tokens=False) - - prompt_input_ids = processed_features["input_ids"][0] - - # Create the output dict with required fields - output = { - "prompt_input_ids": prompt_input_ids, - "prompt_attention_mask": processed_features["attention_mask"][0], - } - - # Add vision-specific fields - 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] - - return output - - def _generate(self, model, prompts, images=None): - """Generate completions using the model""" - device = next(model.parameters()).device - pad_token_id = self._tokenizer.pad_token_id - - # Apply chat template and tokenize the input - inputs = [{"prompt": prompt} for prompt in prompts] - - # Add images if provided (VLM support) - if images is not None: - for i, image in enumerate(images): - inputs[i]["image"] = image - - # Apply chat template to get text prompts - prompts_text = [maybe_apply_chat_template(x, self.processing_class)["prompt"] for x in inputs] - - # Handle image token collapsing/removal - # The chat template sometimes 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. We need to handle this properly. - if self.image_token is not None and images is not None: - escaped_img_token = re.escape(self.image_token) - # Search for the image token in the chat template - if hasattr(self.processing_class, "chat_template") and self.processing_class.chat_template: - if re.search(escaped_img_token, self.processing_class.chat_template): - # Collapse repeated image tokens back into a single token - prompts_text = [ - re.sub(rf"({escaped_img_token})+", self.image_token, text) for text in prompts_text - ] - else: - # If the chat template doesn't use the image token, remove all instances - if self.vision_end_token_id is not None: - escaped_eoi_token = re.escape(self._tokenizer.decode([self.vision_end_token_id])) - prompts_text = [ - re.sub(rf"({escaped_img_token})+{escaped_eoi_token}", "", text) for text in prompts_text - ] - else: - # If vision_end_token_id is None, just remove the image tokens - prompts_text = [re.sub(rf"({escaped_img_token})+", "", text) for text in prompts_text] - - # Prepare kwargs for processing class - kwargs = {} - if images is not None: - kwargs = {"images": [[img] for img in images]} - - # Process inputs using the processing class (handles both VLM and LLM) - prompt_inputs = self.processing_class( - text=prompts_text, - return_tensors="pt", - padding=True, - padding_side="left", - add_special_tokens=False, - **kwargs, - ) - - prompt_inputs = {k: v.to(device) for k, v in prompt_inputs.items()} - # Convert vision inputs to model's dtype for proper computation - if "pixel_values" in prompt_inputs: - # Handle DataParallel wrapped models - model_dtype = getattr(model, "dtype", None) - if model_dtype is None and hasattr(model, "module"): - model_dtype = model.module.dtype - if model_dtype is not None: - prompt_inputs["pixel_values"] = prompt_inputs["pixel_values"].to(model_dtype) - - # Sample 2 completions per prompt of size `max_new_tokens` from the model - prompt_ids = prompt_inputs["input_ids"].repeat(2, 1) - prompt_mask = prompt_inputs["attention_mask"].repeat(2, 1) - - # Prepare vision inputs if available - vision_generation_kwargs = {} - if self.is_vision_model and images is not None: - if "pixel_values" in prompt_inputs: - vision_generation_kwargs["pixel_values"] = prompt_inputs["pixel_values"].repeat(2, 1, 1, 1) - if "pixel_attention_mask" in prompt_inputs: - vision_generation_kwargs["pixel_attention_mask"] = prompt_inputs["pixel_attention_mask"].repeat(2, 1) - if "image_sizes" in prompt_inputs: - vision_generation_kwargs["image_sizes"] = prompt_inputs["image_sizes"].repeat(2, 1) - if "image_grid_thw" in prompt_inputs: - vision_generation_kwargs["image_grid_thw"] = prompt_inputs["image_grid_thw"].repeat(2, 1) - - with ( - profiling_context(self, "transformers.generate"), - unwrap_model_for_generation( - model, - self.accelerator, - gather_deepspeed3_params=self.args.ds3_gather_for_generation, - generation_kwargs=self.generation_kwargs, # Override model.generation_config with generation_kwargs to fix transformers#42762 - ) as unwrapped_model, - torch.no_grad(), - FSDP.summon_full_params(self.model_wrapped, recurse=False) if self.is_fsdp_enabled else nullcontext(), - ): - # Setup cache implementation if specified - if self.args.cache_implementation is not None: - unwrapped_model.generation_config.cache_implementation = self.args.cache_implementation - - # Standard generation - output = unwrapped_model.generate( - input_ids=prompt_ids, - attention_mask=prompt_mask, - generation_config=self.generation_config, - **vision_generation_kwargs, - ) - - completion_ids = output[:, prompt_ids.size(1) :] - completion_ids, completion_mask = truncate_right(completion_ids, self.eos_token_ids, pad_token_id) - - return prompt_ids, prompt_mask, completion_ids, completion_mask - - def _calculate_rewards_from_functions(self, prompts, completions, completion_ids_list, **reward_kwargs): - """ - Calculate rewards using reward functions - """ - device = self.accelerator.device - rewards_per_func = torch.zeros(len(prompts), len(self.reward_funcs), device=device) - - # Add trainer state to reward kwargs for dynamic reward shaping - reward_kwargs["trainer_state"] = self.state - - for i, (reward_func, reward_processing_class) in enumerate( - zip(self.reward_funcs, self.reward_processing_classes, strict=True) - ): - if isinstance(reward_func, nn.Module): # Model-based reward function - # Handle conversational vs text input - if is_conversational({"prompt": prompts[0]}): - messages = [{"messages": p + c} for p, c in zip(prompts, completions, strict=True)] - texts = [apply_chat_template(x, reward_processing_class)["text"] for x in messages] - else: - texts = [p + c for p, c in zip(prompts, completions, strict=True)] - - # Tokenize and get reward scores - reward_inputs = reward_processing_class( - text=texts, return_tensors="pt", padding=True, padding_side="right", add_special_tokens=False - ) - reward_inputs = {k: v.to(device) for k, v in reward_inputs.items()} - - with torch.inference_mode(): - rewards_per_func[:, i] = reward_func(**reward_inputs).logits[:, 0] # Shape (B*G,) - else: - # Custom reward function - output_reward_func = reward_func( - prompts=prompts, completions=completions, completion_ids=completion_ids_list, **reward_kwargs - ) - # Convert None values to NaN - output_reward_func = [reward if reward is not None else torch.nan for reward in output_reward_func] - rewards_per_func[:, i] = torch.tensor(output_reward_func, dtype=torch.float32, device=device) - - # Weight and sum across all reward functions - if self.reward_weights is not None: - total_rewards = (rewards_per_func * self.reward_weights.to(device).unsqueeze(0)).nansum(dim=1) - else: - total_rewards = rewards_per_func.nansum(dim=1) - - return total_rewards - - def _forward(self, model, prompt_ids, prompt_mask, completion_ids, completion_mask, vision_inputs=None): - # Get the number of tokens to truncate from prompt - num_tokens_to_truncate = max(prompt_ids.size(1) + completion_ids.size(1) - self.max_length, 0) - - # Truncate left to avoid oom - prompt_ids = prompt_ids[:, num_tokens_to_truncate:] - prompt_mask = prompt_mask[:, num_tokens_to_truncate:] - - # Concat the prompt and completion - prompt_completion_ids = torch.cat((prompt_ids, completion_ids), dim=1) - prompt_completion_mask = torch.cat((prompt_mask, completion_mask), dim=1) - - # Prepare model kwargs with vision inputs if available - model_kwargs = {"attention_mask": prompt_completion_mask} - if vision_inputs is not None: - if "pixel_values" in vision_inputs: - model_kwargs["pixel_values"] = vision_inputs["pixel_values"] - if "pixel_attention_mask" in vision_inputs: - model_kwargs["pixel_attention_mask"] = vision_inputs["pixel_attention_mask"] - if "image_sizes" in vision_inputs: - model_kwargs["image_sizes"] = vision_inputs["image_sizes"] - if "image_grid_thw" in vision_inputs: - model_kwargs["image_grid_thw"] = vision_inputs["image_grid_thw"] - - # Get the logprobs of the completions from the model - _unsloth_left_pad = None - if not vision_inputs: - _unsloth_left_pad = (prompt_mask == 0).sum(dim = 1) - _unsloth_order = torch.argsort(prompt_completion_mask != 0, dim = 1, descending = True, stable = True) - prompt_completion_ids = prompt_completion_ids.gather(1, _unsloth_order) - prompt_completion_mask = prompt_completion_mask.gather(1, _unsloth_order) - model_kwargs["attention_mask"] = prompt_completion_mask - _unsloth_odpo_kwargs = model_kwargs - _unsloth_odpo = _unsloth_odpo_chunked_head(self, model, _unsloth_odpo_kwargs, "_forward") - output = _unsloth_dpo_model_forward( - model, prompt_completion_ids, _unsloth_odpo_kwargs, _unsloth_odpo is not None - ) - _unsloth_odpo_hidden = _unsloth_dpo_returns_hidden_states(self, _unsloth_odpo, output.logits) - - # There is 1 offset, because the model predicts the next token - prompt_len = prompt_ids.size(1) - start_idx = prompt_len - 1 if prompt_len > 0 else 0 - # Only slice off the last logit when we have a prompt, otherwise we need all logits - end_idx = -1 if prompt_len > 0 else None - if _unsloth_left_pad is not None: - _unsloth_index = (start_idx - _unsloth_left_pad).unsqueeze(1) + torch.arange( - completion_ids.size(1), device = completion_ids.device - ).unsqueeze(0) - _unsloth_index = _unsloth_index.clamp(0, output.logits.size(1) - 1) - _unsloth_rows, _unsloth_len = output.logits.shape[:2] - _unsloth_index = _unsloth_index + _unsloth_len * torch.arange( - _unsloth_rows, device = _unsloth_index.device - ).unsqueeze(1) - logits = output.logits.reshape(_unsloth_rows * _unsloth_len, -1).index_select( - 0, _unsloth_index.reshape(-1) - ).view(_unsloth_rows, -1, output.logits.size(-1)) - output = None - else: - logits = output.logits[:, start_idx:end_idx] - - # Take the completion tokens logprob - if _unsloth_odpo_hidden: - logprobs = _unsloth_odpo_hidden_logps(self, _unsloth_odpo, logits, completion_ids) - else: - logprobs = torch.take_along_dim(logits.log_softmax(dim=-1), completion_ids.unsqueeze(-1), dim=2).squeeze(-1) - return logprobs - - def training_step( - self, model: nn.Module, inputs: dict[str, torch.Tensor | Any], num_items_in_batch: int | None = None - ) -> torch.Tensor: - model.train() - - prompts = inputs["prompt"] - batch_size = len(prompts) - - # Handle images for VLM support - has_images = "image" in inputs - images = None - if has_images: - images = inputs["image"] - # Convert conversational prompts to include image tokens - for prompt in prompts: - if isinstance(prompt, list): - for message in prompt: - if not isinstance(message, dict): - continue - content = message.get("content") - role = message.get("role") - if isinstance(content, str): - if role == "user": - message["content"] = [{"type": "image"}, {"type": "text", "text": content}] - elif role == "system": - message["content"] = [{"type": "text", "text": content}] - - if self.args.use_vllm: - prompt_ids, prompt_mask, completion_ids, completion_mask = self._generate_vllm(prompts, images) - else: - prompt_ids, prompt_mask, completion_ids, completion_mask = self._generate(model, prompts, images) - - contain_eos_token = torch.any( - torch.isin(completion_ids, torch.tensor(self.eos_token_ids, device=completion_ids.device)), dim=-1 - ) - - # Extract vision inputs if available for VLM support - vision_inputs = None - if has_images and self.is_vision_model and not self.args.use_vllm: - # For vision models with transformers generation, we need to prepare vision inputs - # Process the images to get vision inputs that can be passed through the forward pass - vision_inputs = {} - kwargs = {"images": [[img] for img in images]} - processed = self.processing_class( - text=[""] * len(images), # Dummy text for vision processing - return_tensors="pt", - **kwargs, - ) - # Handle DataParallel wrapped models - model_device = getattr(model, "device", None) - model_dtype = getattr(model, "dtype", None) - if model_device is None and hasattr(model, "module"): - model_device = model.module.device - model_dtype = model.module.dtype - # Move vision tensors to device and convert to model dtype - # Need to duplicate for 2 completions per prompt - if "pixel_values" in processed: - vision_inputs["pixel_values"] = ( - processed["pixel_values"].to(model_device, dtype=model_dtype).repeat(2, 1, 1, 1) - ) - if "pixel_attention_mask" in processed: - vision_inputs["pixel_attention_mask"] = processed["pixel_attention_mask"].to(model_device).repeat(2, 1) - if "image_sizes" in processed: - vision_inputs["image_sizes"] = processed["image_sizes"].to(model_device).repeat(2, 1) - if "image_grid_thw" in processed: - vision_inputs["image_grid_thw"] = processed["image_grid_thw"].to(model_device).repeat(2, 1) - - logprobs = self._forward(model, prompt_ids, prompt_mask, completion_ids, completion_mask, vision_inputs) - with torch.no_grad(): - if self.ref_model is not None: - ref_logprobs = self._forward( - self.ref_model, prompt_ids, prompt_mask, completion_ids, completion_mask, vision_inputs - ) - else: # peft case: we just need to disable the adapter - with self.model.disable_adapter(): - ref_logprobs = self._forward( - self.model, prompt_ids, prompt_mask, completion_ids, completion_mask, vision_inputs - ) - - # Decode the completions, and format them if the input is conversational - device = logprobs.device - completions = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) - if is_conversational({"prompt": prompts[0]}): - completions = [[{"role": "assistant", "content": completion}] for completion in completions] - - # Get the reward from reward functions - # First create completion_ids_list for custom reward functions - completion_ids_list = [completion_ids[i].tolist() for i in range(completion_ids.shape[0])] - - # Extract additional fields from inputs for reward functions - reward_kwargs = {} - keys = [key for key in inputs if key not in ["prompt"]] - for key in keys: - if isinstance(inputs[key], (list, tuple)): - # Repeat input fields to match number of completions (2 per prompt) - reward_kwargs[key] = inputs[key] * 2 - else: - reward_kwargs[key] = inputs[key] - - # Calculate rewards using reward functions - rewards = self._calculate_rewards_from_functions( - prompts=2 * prompts, completions=completions, completion_ids_list=completion_ids_list, **reward_kwargs - ) - - # Apply missing EOS penalty if configured - if self.args.missing_eos_penalty is not None: - rewards[~contain_eos_token] -= self.args.missing_eos_penalty - - # Split rewards into chosen/rejected pairs - first_half, second_half = rewards.split(batch_size) - mask = first_half >= second_half - - batch_range = torch.arange(batch_size, device=device) - chosen_indices = batch_range + (~mask * batch_size) - rejected_indices = batch_range + (mask * batch_size) - - # Build tensor so that the first half is the chosen examples and the second half the rejected examples - cr_indices = torch.cat((chosen_indices, rejected_indices), dim=0) # cr = chosen and rejected - cr_logprobs = logprobs[cr_indices] - cr_ref_logprobs = ref_logprobs[cr_indices] - - # mask out the padding tokens - padding_mask = ~completion_mask.bool() - cr_padding_mask = padding_mask[cr_indices] - - cr_logprobs_sum = (cr_logprobs * ~cr_padding_mask).sum(1) - cr_ref_logprobs_sum = (cr_ref_logprobs * ~cr_padding_mask).sum(1) - - # Split the chosen and rejected examples - chosen_logprobs_sum, rejected_logprobs_sum = torch.split(cr_logprobs_sum, batch_size) - chosen_ref_logprobs_sum, rejected_ref_logprobs_sum = torch.split(cr_ref_logprobs_sum, batch_size) - pi_logratios = chosen_logprobs_sum - rejected_logprobs_sum - ref_logratios = chosen_ref_logprobs_sum - rejected_ref_logprobs_sum - - logits = pi_logratios - ref_logratios - - if self.args.loss_type == "sigmoid": - losses = -F.logsigmoid(self.beta * logits) - elif self.args.loss_type == "ipo": - losses = (logits - 1 / (2 * self.beta)) ** 2 - else: - raise NotImplementedError(f"invalid loss type {self.args.loss_type}") - - loss = losses.mean() - - # Log everything - if self.reward_funcs is not None: - # When using reward_funcs, we have rewards instead of scores - scores_margin = rewards[chosen_indices] - rewards[rejected_indices] - self.stats["objective/scores_margin"].append( - self.accelerator.gather_for_metrics(scores_margin.mean()).mean().item() - ) - self.stats["objective/scores"].append(self.accelerator.gather_for_metrics(rewards.mean()).mean().item()) - self.stats["val/contain_eos_token"].append(contain_eos_token.float().mean().item()) - self.stats["logps/chosen"].append(self.accelerator.gather_for_metrics(chosen_logprobs_sum).mean().item()) - self.stats["logps/rejected"].append(self.accelerator.gather_for_metrics(rejected_logprobs_sum).mean().item()) - - kl = logprobs - ref_logprobs - mean_kl = kl.sum(1).mean() - self.stats["objective/kl"].append(self.accelerator.gather_for_metrics(mean_kl).mean().item()) - non_score_reward = (-self.beta * kl).sum(1) - mean_non_score_reward = non_score_reward.mean() - self.stats["objective/non_score_reward"].append( - self.accelerator.gather_for_metrics(mean_non_score_reward).mean().item() - ) - if self.reward_funcs is not None: - # Calculate RLHF reward by combining rewards with non_score_reward - rlhf_reward = rewards + non_score_reward - self.stats["objective/rlhf_reward"].append(self.accelerator.gather_for_metrics(rlhf_reward).mean().item()) - - mean_entropy = -logprobs.sum(1).mean() - self.stats["objective/entropy"].append(self.accelerator.gather_for_metrics(mean_entropy).mean().item()) - chosen_rewards = self.beta * (chosen_logprobs_sum - chosen_ref_logprobs_sum) - gathered_chosen_rewards = self.accelerator.gather_for_metrics(chosen_rewards) - self.stats["rewards/chosen"].append(gathered_chosen_rewards.mean().item()) - rejected_rewards = self.beta * (rejected_logprobs_sum - rejected_ref_logprobs_sum) - gathered_rejected_rewards = self.accelerator.gather_for_metrics(rejected_rewards) - self.stats["rewards/rejected"].append(gathered_rejected_rewards.mean().item()) - margin = gathered_chosen_rewards - gathered_rejected_rewards - self.stats["rewards/margins"].append(margin.mean().item()) - accuracy = margin > 0 - self.stats["rewards/accuracies"].append(accuracy.float().mean().item()) - self.stats["beta"].append(self.beta) - - if ( - self.args.torch_empty_cache_steps is not None - and self.state.global_step % self.args.torch_empty_cache_steps == 0 - ): - empty_cache() - - kwargs = {} - - # For LOMO optimizers you need to explicitly use the learning rate - if self.args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]: - kwargs["learning_rate"] = self._get_learning_rate() - - if self.args.n_gpu > 1: - loss = loss.mean() # mean() to average on multi-gpu parallel training - - self.accelerator.backward(loss, **kwargs) - - return loss.detach() / self.args.gradient_accumulation_steps - - # Same as Trainer._maybe_log_save_evaluate but log our metrics - def _maybe_log_save_evaluate( - self, tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval, start_time, learning_rate=None - ): - if self.control.should_log and self.state.global_step > self._globalstep_last_logged: - logs: dict[str, float] = {} - - # all_gather + mean() to get average loss over all processes - if Version(transformers.__version__) >= Version("5.2.0"): - tr_loss_scalar = nested_gather(tr_loss, self.args.parallel_mode).mean().item() - else: - tr_loss_scalar = self._nested_gather(tr_loss).mean().item() - - # reset tr_loss to zero - tr_loss -= tr_loss - - logs["loss"] = round(tr_loss_scalar / (self.state.global_step - self._globalstep_last_logged), 4) - if grad_norm is not None: - logs["grad_norm"] = grad_norm.detach().item() if isinstance(grad_norm, torch.Tensor) else grad_norm - if learning_rate is not None: - logs["learning_rate"] = learning_rate - else: - logs["learning_rate"] = self._get_learning_rate() - - # Add our metrics - for key, val in self.stats.items(): - logs[key] = sum(val) / len(val) - self.stats = {key: [] for key in self.stats} # reset stats - - self._total_loss_scalar += tr_loss_scalar - self._globalstep_last_logged = self.state.global_step - self.store_flos() - self.log(logs, start_time) - - metrics = None - if self.control.should_evaluate: - metrics = self._evaluate(trial, ignore_keys_for_eval) - is_new_best_metric = self._determine_best_metric(metrics=metrics, trial=trial) - - if self.args.save_strategy == "best": - self.control.should_save = is_new_best_metric - - if self.control.should_save: - self._save_checkpoint(model, trial) - self.control = self.callback_handler.on_save(self.args, self.state, self.control) - - # Ensure the model card is saved along with the checkpoint - def _save_checkpoint(self, model, trial): - if self.args.hub_model_id is None: - model_name = Path(self.args.output_dir).name - else: - model_name = self.args.hub_model_id.split("/")[-1] - self.create_model_card(model_name=model_name) - super()._save_checkpoint(model, trial) -class UnslothOnlineDPOTrainer(_UnslothOnlineDPOTrainer): - """ - - Initialize OnlineDPOTrainer. - - Args: - model (`str | nn.Module | PreTrainedModel`): - Model to be trained. Can be either: - - - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a - path to a *directory* containing model weights saved using - [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded - using [`~transformers.AutoModelForCausalLM.from_pretrained`] with the keyword arguments in - `args.model_init_kwargs`. - - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. - ref_model ([`~transformers.PreTrainedModel`] or `torch.nn.Module` or `None`): - The reference model to use for training. If None is specified, the reference model will be created from the - model. - reward_funcs (`RewardFunc | list[RewardFunc]`): - Reward functions to be used for computing the rewards. To compute the rewards, we call all the reward - functions with the prompts and completions and sum the rewards. Can be either: - - - A single reward function: Can be a string (path to model), a [`~transformers.PreTrainedModel`], or a - custom callable function. - - A list of reward functions: Must all be of compatible types. - args ([`experimental.online_dpo.OnlineDPOConfig`]): - The online DPO config arguments to use for training. - data_collator ([`~transformers.DataCollator`]): - The data collator to use for training. If None is specified, the default data collator - ([`experimental.utils.DPODataCollatorWithPadding`]) will be used which will pad the sequences to the - maximum length of the sequences in the batch, given a dataset of paired sequences. - train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): - The dataset to use for training. - eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`] or `dict[str, Dataset | IterableDataset]`): - The dataset to use for evaluation. - processing_class ([`~transformers.PreTrainedTokenizerBase`] or [`~transformers.ProcessorMixin`], *optional*): - Processing class used to process the data. If provided, will be used to automatically process the inputs - for the model, and it will be saved along the model to make it easier to rerun an interrupted training or - reuse the fine-tuned model. - reward_processing_classes ([`~transformers.PreTrainedTokenizerBase`] or `list[PreTrainedTokenizerBase]`, *optional*): - Processing classes corresponding to the reward functions specified in `reward_funcs`. Can be either: - - - A single processing class: Used when `reward_funcs` contains only one reward function. - - A list of processing classes: Must match the order and length of the reward functions in `reward_funcs`. - - If set to `None`, the tokenizer for each model-based reward function is automatically loaded using - [`~transformers.AutoTokenizer.from_pretrained`]. - peft_config ([`~peft.PeftConfig`], *optional*): - PEFT configuration used to wrap the model. If `None`, the model is not wrapped. - compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): - The function to use to compute the metrics. Must take a `EvalPrediction` and return a dictionary string to - metric values. - callbacks (`list[transformers.TrainerCallback]`): - The callbacks to use for training. - optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`): - The optimizer and scheduler to use for training. - preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): - The function to use to preprocess the logits before computing the metrics. - - """ - def __init__( - self, - model, - ref_model = None, - reward_funcs = None, - args = None, - data_collator = None, - train_dataset = None, - eval_dataset = None, - processing_class = None, - reward_processing_classes = None, - peft_config = None, - compute_metrics = None, - callbacks = None, - preprocess_logits_for_metrics = None, - **kwargs - ): - if args is None: args = UnslothOnlineDPOConfig() - _unsloth_precision_requested = tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) - use_bf16 = getattr(args, 'bf16', False) - if type(use_bf16) is not bool: use_bf16 = False - use_fp16 = getattr(args, 'fp16', False) - if type(use_fp16) is not bool: use_fp16 = False - force_float32 = False - try: - from unsloth_zoo.device_type import device_is_bf16_supported as _bf16_supported - except Exception: - _bf16_supported = torch.cuda.is_bf16_supported - try: - from unsloth_zoo.device_type import DEVICE_TYPE as _unsloth_device_type - except Exception: - _unsloth_device_type = None - full_finetuning = getattr(model, '_unsloth_full_finetuning', None) - if full_finetuning is None: full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1' - model_forced_float32 = getattr(model, '_unsloth_forced_float32', None) - if model_forced_float32 is None: model_forced_float32 = os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1' - if model_forced_float32 and not (full_finetuning and _bf16_supported()): - print('Unsloth: Switching to float32 training since model cannot work with float16') - force_float32 = True - mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') - dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None) - if dtype is None: dtype = model.get_input_embeddings().weight.dtype - from unsloth_zoo.utils import _get_dtype - dtype = _get_dtype(dtype) - float16 = dtype == torch.float16 - bfloat16 = dtype == torch.bfloat16 - float32 = dtype == torch.float32 - user_float32 = bool(getattr(model, '_unsloth_user_float32', False)) - if full_finetuning: - if bfloat16 and use_fp16: use_fp16 = False - if float16 and use_bf16: use_bf16 = False - if not force_float32 and bfloat16 and use_fp16 and _unsloth_device_type == 'xpu' and not _bf16_supported(): - print('Unsloth: The model was loaded in bfloat16 on purpose, so training uses bfloat16 mixed precision instead of float16.') - use_fp16 = False - use_bf16 = True - args.fp16 = False - args.bf16 = True - if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`') - if not force_float32 and (bfloat16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`') - if force_float32: - # Forced float32 training - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32' and float32 and user_float32 and not _bf16_supported(): - print('Unsloth: Model is in float32 and this GPU has no bfloat16 support, so training stays in float32. Pass fp16 = True to force float16 mixed precision instead.') - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32': - # Mixed precision training. bf16 only if the GPU supports it; V100/T4 use fp16. - use_bf16_amp = (not float16) and (_bf16_supported() or (bfloat16 and _unsloth_device_type == 'xpu')) - args.fp16 = not use_bf16_amp - args.bf16 = use_bf16_amp - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16_amp else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16_amp else 'fp16' - # args.mixed_precision is a new argument which needs to be set now - elif mixed_precision_dtype == 'bfloat16': - # Both False since bfloat16 full finetuning doesn't do any autocasting. - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif use_bf16 or use_fp16: - # transformers <5 exported this itself from fp16/bf16; 5.x dropped the write, so an - # explicit flag left it unset and GRPO readers defaulted to 'fp16', wrapping a - # bfloat16 model in a float16 autocast. See unslothai/unsloth#4891. - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16 else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16 else 'fp16' - - if getattr(args, 'eval_dataset', None) is not None and getattr(args, 'eval_strategy', 'no') == 'no': - args.eval_strategy = 'steps' - if getattr(args, 'eval_steps', None) is None: args.eval_steps = 0.1 - ga_steps = getattr(args, 'gradient_accumulation_steps', None) - if ga_steps is not None and ga_steps > 1: - from transformers import __version__ as transformers_version - if Version(transformers_version) <= Version('4.45.2'): - print('**** Unsloth: Please use our fixed gradient_accumulation_steps by updating transformers, TRL and Unsloth!\n' - '`pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth transformers trl unsloth_zoo`') - if getattr(args, 'eval_strategy', 'no') != 'no': - eval_bsz = getattr(args, 'per_device_eval_batch_size', 8) - eval_generations = getattr(args, 'num_generations_eval', None) or getattr(args, 'num_generations', None) or 1 - if eval_bsz == 8 and args.per_device_train_batch_size < eval_bsz and (args.per_device_train_batch_size * args.world_size) % eval_generations == 0: args.per_device_eval_batch_size = args.per_device_train_batch_size - if getattr(args, 'eval_accumulation_steps', None) is None and ga_steps is not None: args.eval_accumulation_steps = ga_steps - fp16_full_eval = getattr(args, 'fp16_full_eval', False) - if type(fp16_full_eval) is not bool: fp16_full_eval = False - bf16_full_eval = getattr(args, 'bf16_full_eval', False) - if type(bf16_full_eval) is not bool: bf16_full_eval = False - if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True - if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False - if force_float32: - args.bf16_full_eval = False - args.fp16_full_eval = False - elif os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16': - args.bf16_full_eval = True - args.fp16_full_eval = False - elif not bf16_full_eval and not fp16_full_eval: - args.bf16_full_eval = args.bf16 - args.fp16_full_eval = args.fp16 - if tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) != _unsloth_precision_requested: - from unsloth.models.loader_utils import sync_deepspeed_mixed_precision - sync_deepspeed_mixed_precision(args) - if model is not None: - from unsloth.models.loader_utils import apply_deepspeed_zero3_model_setup - apply_deepspeed_zero3_model_setup(model) - if model is not None: - _warnings_issued = getattr(model, 'warnings_issued', None) - if _warnings_issued is None: - model.warnings_issued = {} - elif not isinstance(_warnings_issued, dict): - try: - model.warnings_issued = dict(_warnings_issued) - except Exception: - model.warnings_issued = {} - if 'max_seq_length' not in locals() and not hasattr(args, 'max_seq_length'): - pass - else: - model_max_seq_length = getattr(model, 'max_seq_length', None) - args_max_seq_length = getattr(args, 'max_seq_length', None) - if args_max_seq_length is None and model_max_seq_length is not None: - max_seq_length = model.max_seq_length - if hasattr(args, 'max_seq_length'): args.max_seq_length = max_seq_length - elif args_max_seq_length is not None and model_max_seq_length is not None: - if args_max_seq_length > model_max_seq_length: - print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but ' - 'the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.') - args.max_seq_length = model_max_seq_length - if model is not None and hasattr(model, 'for_training'): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right' - if 'processing_class' in locals(): - if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right' - if hasattr(processing_class, 'tokenizer') and hasattr(processing_class.tokenizer, 'padding_side'): processing_class.tokenizer.padding_side = 'right' - __tokenizer = processing_class if 'processing_class' in locals() else tokenizer - from unsloth_zoo.vision_utils import UnslothVisionDataCollator - if not isinstance(data_collator, UnslothVisionDataCollator): - if isinstance(data_collator, DataCollatorForSeq2Seq) and 'labels' not in _unsloth_dataset_column_names(train_dataset): - data_collator = TransformersDataCollatorForLanguageModeling( - __tokenizer, - mlm = False, - mlm_probability = 0.0, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling) and 'labels' in _unsloth_dataset_column_names(train_dataset): - data_collator = DataCollatorForSeq2Seq( - __tokenizer, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - else: - if hasattr(args, 'remove_unused_columns'): args.remove_unused_columns = False - if hasattr(args, 'dataset_text_field'): args.dataset_text_field = '' - if hasattr(args, 'dataset_kwargs'): args.dataset_kwargs = {'skip_prepare_dataset': True} - if not isinstance(data_collator, UnslothVisionDataCollator): - if not hasattr(__tokenizer, 'pad') and hasattr(__tokenizer, 'tokenizer'): - if isinstance(data_collator, DataCollatorForSeq2Seq): - data_collator = DataCollatorForSeq2Seq( - __tokenizer.tokenizer, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling): - data_collator = TransformersDataCollatorForLanguageModeling( - __tokenizer.tokenizer, - mlm = False, - mlm_probability = 0.0, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - other_metrics = [] - - from unsloth_zoo.logging_utils import PatchRLStatistics - PatchRLStatistics('online_dpo_trainer', other_metrics) - - # [TODO] Fix up DataParallel multiplying batch sizes - # [TODO] DDP works, but DP seems to not work? [TODO] - if getattr(args, "parallel_mode", None) == ParallelMode.NOT_DISTRIBUTED and args.n_gpu > 1: - if getattr(args, "_n_gpu", 1) != 1: - args._n_gpu = 1 - if "model" in locals(): - try: - from unsloth.models._gc_config_default import honour_gradient_checkpointing_off - honour_gradient_checkpointing_off(args, model) - except Exception: - pass - if "model" in locals() and hasattr(model, "for_training"): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if "model" in locals(): - try: - from unsloth.models._ddp_defaults import unsloth_default_ddp_find_unused_parameters - unsloth_default_ddp_find_unused_parameters(args, model) - except Exception: - pass - super().__init__( - model = model, - ref_model = ref_model, - reward_funcs = reward_funcs, - args = args, - data_collator = data_collator, - train_dataset = train_dataset, - eval_dataset = eval_dataset, - processing_class = processing_class, - reward_processing_classes = reward_processing_classes, - peft_config = peft_config, - compute_metrics = compute_metrics, - callbacks = callbacks, - preprocess_logits_for_metrics = preprocess_logits_for_metrics,**kwargs) - # A tensor parallel model (tp_native.tp_trainer_setup) stays on its training path: for_inference would - # route the decoder layers onto decode paths outside the tensor parallel regions. - if "model" in locals() and hasattr(model, "for_inference") and not getattr(self, "_unsloth_tensor_parallel", False): - model.for_inference() - if hasattr(self, 'neftune_hook_handle'): - self.neftune_hook_handle.remove() - if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle - if getattr(args, 'neftune_noise_alpha', None) is not None: - model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha - pass - if hasattr(self, 'accelerator'): - scaler = self.accelerator.scaler - current_model = model - while hasattr(current_model, 'model'): - current_model.accelerator_scaler = scaler - current_model = current_model.model - current_model.accelerator_scaler = scaler - pass - if hasattr(self, 'train'): - self.train = MethodType(prepare_for_training_mode(self.__class__.train), self) - pass - if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'): - _vllm_tok = self.llm.get_tokenizer() - _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None) - if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None: - _vllm_tok.chat_template = _pc.chat_template - pass - if getattr(self, 'aux_loss_enabled', False) and hasattr(getattr(self, 'model', None), 'config'): - _text_config = self.model.config - if hasattr(_text_config, 'get_text_config'): _text_config = _text_config.get_text_config() - _n_experts = [getattr(_text_config, _k) for _k in ('num_local_experts', 'num_experts', 'n_routed_experts', 'moe_num_experts') if isinstance(getattr(_text_config, _k, None), int)] - if _n_experts and all(_n <= 1 for _n in _n_experts): - self.aux_loss_enabled = False - _text_config.output_router_logits = False - pass - -pass diff --git a/unsloth_compiled_cache/UnslothRLOOTrainer.py b/unsloth_compiled_cache/UnslothRLOOTrainer.py deleted file mode 100644 index befb253e63b5e0f9abfae9352e7a0a964f523174..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/UnslothRLOOTrainer.py +++ /dev/null @@ -1,3196 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from unsloth_zoo.temporary_patches.common import torch_compile -from unsloth_zoo.temporary_patches.common import _maybe_compile -import functools -import transformers -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from trl.trainer.rloo_trainer import (Any, AutoModelForSequenceClassification, AutoProcessor, AutoTokenizer, BaseTunerLayer, BitsAndBytesConfig, DataLoader, Dataset, DatasetDict, DistributedBackend, GenerationConfig, IterableDataset, IterableDatasetDict, LoraConfig, Path, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, RLOOConfig, RLOOTrainer, RepeatSampler, RewardFunc, Sampler, SyncRefModelCallback, TrainerCallback, VLLMGeneration, Version, _BaseTrainer, add_fused_lm_head, apply_chat_template, asyncio, atexit, copy, create_model_from_path, defaultdict, deque, disable_dropout_in_model, disable_gradient_checkpointing, gather, gather_object, get_callable_name, get_config_model_id, get_peft_model, identity, is_async_callable, is_conversational, is_peft_available, is_peft_model, is_rich_available, logger, math, nanmax, nanmin, nanstd, nn, np, pad, pd, peft, prepare_deepspeed, prepare_fsdp, prepare_multimodal_messages, print_prompt_completions_sample, profiling_context, profiling_decorator, repeat_iterable_dataset, set_seed, shuffle_sequence_dict, shutdown_event_loop_in_daemon, split_pixel_values_by_grid, split_tensor_dict, start_event_loop_in_daemon, strip_images_from_messages, textwrap, time, torch, transformers, unsplit_pixel_values_by_grid, unwrap_model_for_generation, use_adapter, warnings, AutoModelForSequenceClassification, AutoProcessor, AutoTokenizer, BaseTunerLayer, BitsAndBytesConfig, Dataset, DatasetDict, DistributedBackend, GenerationConfig, IterableDataset, IterableDatasetDict, LoraConfig, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, RLOOConfig, RLOOTrainer, RepeatSampler, RewardFunc, Sampler, SyncRefModelCallback, TrainerCallback, VLLMGeneration, Version, add_fused_lm_head, atexit, copy, create_model_from_path, defaultdict, deque, disable_dropout_in_model, gather, get_callable_name, get_config_model_id, get_peft_model, identity, is_async_callable, is_peft_available, is_peft_model, logger, nn, np, pad, pd, peft, prepare_deepspeed, prepare_fsdp, repeat_iterable_dataset, set_seed, shutdown_event_loop_in_daemon, start_event_loop_in_daemon, time, torch, transformers, warnings, Any, np, profiling_decorator, shuffle_sequence_dict, split_pixel_values_by_grid, split_tensor_dict, torch, unsplit_pixel_values_by_grid, PeftModel, PreTrainedModel, is_peft_available, logger, peft, torch) - - -import os -import math -import logging -from typing import * -from dataclasses import dataclass, field -from packaging.version import Version -import torch -import numpy as np -from contextlib import nullcontext -from torch.nn import functional as F -import inspect -from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling -from transformers.training_args import ParallelMode -from unsloth_zoo.device_type import DEVICE_TYPE, DEVICE_TYPE_TORCH, device_synchronize - -# Wrap trainer with padding to right and enable training mode -import functools -from types import MethodType -try: - from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers -except: - def reset_unsloth_gradient_checkpointing_buffers(): pass -# Canonical reset lives in unsloth.models._utils so the SFT auto-packing wrapper and the plain -# Trainer loop can import the same helper; fall back to a no-op only if it can't be imported. -try: - from unsloth.models._utils import _unsloth_reset_stray_compile_cache -except Exception: - def _unsloth_reset_stray_compile_cache(self): pass -try: - from unsloth.models._utils import _unsloth_dataset_column_names -except Exception: - def _unsloth_dataset_column_names(dataset): return dataset.column_names -# Drops/renames config arguments the installed TRL no longer accepts, so a -# script pinned to an older TRL keeps working after an upgrade. Falls back to -# the historical raw passthrough so this can never break trainer construction. -try: - from unsloth.models.rl_config_compat import filter_config_init_kwargs as _unsloth_filter_config_init_kwargs - # A cache file generated here can be imported by an older Unsloth whose filter - # predates `mirrored_from`, so drop the argument rather than raise TypeError. - if "mirrored_from" not in inspect.signature(_unsloth_filter_config_init_kwargs).parameters: - _unsloth_filter_config_init_kwargs_old = _unsloth_filter_config_init_kwargs - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): - return _unsloth_filter_config_init_kwargs_old(config_class, kwargs) -except Exception: - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): return kwargs -def prepare_for_training_mode(f): - @functools.wraps(f) - def wrapper(self, *args, **kwargs): - # Drop any torch.compile graph cache poisoned by a stray pre-train forward. - try: - _unsloth_reset_stray_compile_cache(self) - except Exception: - pass - # Finish the previous W&B run if this is a subsequent train() call. - # We do this at the START of train() (not the end) so that - # evaluate() / log() still work after train() completes. - # HF's WandbCallback.setup() will call wandb.init() for the new run. - # See: https://github.com/unslothai/unsloth/issues/3954 - if getattr(self, '_unsloth_training_completed', False): - try: - import wandb - if wandb.run is not None: - wandb.finish() - # Reset HF's WandbCallback so it calls wandb.init() for the new run - for cb in self.callback_handler.callbacks: - if type(cb).__name__ == 'WandbCallback': - cb._initialized = False - break - except: - pass - # Enable training mode - _was_training = None - # Restore the GC mode the model was configured with at setup; fall back to - # the training args only when it wasn't recorded (issue #4735). Use hasattr, - # not a None sentinel, so a deliberately-recorded None is restored verbatim. - _model = getattr(self, 'model', None) - if hasattr(_model, '_unsloth_gradient_checkpointing'): - use_gc = _model._unsloth_gradient_checkpointing - else: - use_gc = getattr(self.args, 'gradient_checkpointing', True) - if hasattr(self, 'model') and hasattr(self.model, "training"): - _was_training = self.model.training - if hasattr(self, 'model') and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - output = f(self, *args, **kwargs) - # Restore previous mode when possible - if hasattr(self, 'model') and hasattr(self.model, "for_inference"): - if _was_training is False: - self.model.for_inference() - elif _was_training is True and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - # Reset gradient checkpointing buffers to free memory while staying ready for next run - try: - reset_unsloth_gradient_checkpointing_buffers() - except: - pass - # Mark that training completed so the next train() call can - # finish this W&B run before starting a new one - self._unsloth_training_completed = True - return output - return wrapper -pass - -torch_compile_options = { - "epilogue_fusion" : True, - "max_autotune" : False, - "shape_padding" : True, - "trace.enabled" : False, - "triton.cudagraphs" : False, -} - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_hidden_states_selective_log_softmax( - hidden_states: torch.Tensor, - lm_head: torch.Tensor, - index: torch.Tensor, - chunks: int = 4, - logit_scale_multiply: float = 0.0, - logit_scale_divide: float = 0.0, - logit_softcapping: float = 0.0, - temperature: float = 1.0, - # Rows per chunk cap. Read HERE, in the default, not in the body: the body - # is traced with fullgraph = True, and `os.environ` is an unsupported op - # there on torch 2.4 -- Dynamo raises - # torch._dynamo.exc.Unsupported: const method call bytes.decode - # from os._Environ.__getitem__, which the eager fallback does not catch - # (it only catches recompile-limit and disabled-hook breaks), so the very - # first call would die even with the variable unset. A default is evaluated - # once when the def runs, which is import time, outside any traced region. - # It is also a plain int argument, so Dynamo guards on it instead of - # constant-folding an unguarded read (2.7+ never notice a later change). - # A non-numeric value is ignored rather than raised on, so a typo cannot - # break the import. 0 keeps the previous chunk boundaries exactly. - max_rows_per_chunk: int = ( - int(os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip()) - if os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip().isdigit() - else 0 - ), - lm_head_bias: Optional[torch.Tensor] = None, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, - # Read in the default for the same Dynamo reason as max_rows_per_chunk. Pairs with vLLM's - # head_dtype=float32 that load_vllm requests under the same flag. - fp32_lm_head: bool = ( - os.environ.get("UNSLOTH_VLLM_FP32_LM_HEAD", "0").strip().lower() in ("1", "true", "on", "yes") - ), - # Older Unsloth scores old / ref log-probs here and the policy with the fused kernel: one scorer keeps ratio at 1. - fused_linear: bool = ( - os.environ.get("UNSLOTH_FUSED_LINEAR", "1") != "0" - and os.environ.get("UNSLOTH_GRPO_CHUNKED_LOGPROBS", "1") != "0" - ), -) -> torch.Tensor: - # All Unsloth Zoo code licensed under AGPL3 - from unsloth_zoo.fused_losses import chunked_logprob as _chunked_logprob - if (fused_linear and _chunked_logprob._FUSED_LINEAR and not fp32_lm_head - and support_top_k == 0 and support_top_p == 1.0 and support_min_p == 0.0 - and support_excluded_ids is None): - return _chunked_logprob.unsloth_chunked_logprobs( - hidden_states, lm_head, index, - temperature = temperature, - logit_scale_multiply = logit_scale_multiply, - logit_scale_divide = logit_scale_divide, - logit_softcapping = logit_softcapping, - lm_head_bias = lm_head_bias, - max_rows_per_chunk = max_rows_per_chunk if max_rows_per_chunk > 0 else None, - ).reshape((hidden_states.shape[0], hidden_states.shape[1])) - # Reshape on this tensor's own last dim: a no-op, so a wrong-width caller - # cannot have its row count silently rewritten and instead fails at the - # matmul below, which prints both operands. Do not swap in a bare - # torch._check: it reports only "Expected cond to be True", naming neither - # operand, and Dynamo rejects a message-carrying one. Callers dispatch on - # the width first -- see `compute_logprobs_chunk`, the packed path and - # `_pg_grad_forward`. - flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) - flat_index = index.reshape(-1) - - # Each chunk materialises rows x vocab logits and then a float32 copy of - # them, all on the device holding the output head. With a large vocabulary - # and a fixed chunk count that grows with the batch, so the peak scales with - # the batch rather than staying bounded. max_rows_per_chunk caps the rows - # per chunk instead, which is pure loop splitting: more, smaller chunks, - # same concatenated result. 0 (the default) keeps the previous chunk - # boundaries exactly. - if max_rows_per_chunk > 0: - n_rows = flat_hidden_states.shape[0] - chunks = max(chunks, -(-n_rows // max_rows_per_chunk)) - chunks = min(chunks, max(n_rows, 1)) - - # A one-row chunk's head gradient is a K=1 matmul Inductor mis-lowers under dynamic shapes (illegal - # memory access or NaN head gradients, torch 2.13), so keep chunks at 2+ rows when the head trains. - # Exactly `chunks` row ranges, so the unrolled loop length never depends on T: torch.chunk - # returns ceil(T / ceil(T / chunks)) pieces and recompiled per new count; tensor_split guards - # on T % chunks. At most rows // 2 chunks when the head trains keeps every piece at 2+ rows. - lm_head_grad = torch.is_grad_enabled() and lm_head.requires_grad - if lm_head_grad and flat_hidden_states.shape[0] < 2 * chunks: - chunks = max(flat_hidden_states.shape[0] // 2, 1) - - n_rows = flat_hidden_states.shape[0] - chunks = max(1, min(chunks, n_rows)) - bounds = [(i * n_rows) // chunks for i in range(chunks + 1)] - chunked_hidden_states = [flat_hidden_states[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - - # Cast once, outside the loop, so every chunk shares one fp32 head copy (also for autograd). - if fp32_lm_head: - lm_head = lm_head.to(torch.float32) - if lm_head_bias is not None: - lm_head_bias = lm_head_bias.to(torch.float32) - - all_per_token_logps = [] - - for chunk_hidden_states, chunk_index in zip(chunked_hidden_states, chunked_index): - # When the model is dispatched over several devices, the output head can - # sit on a different one from the hidden states, because accelerate - # places the tail of the model on the last device it fills. Co-locate on - # the head's device before the matmul, otherwise this raises - # Unhandled FakeTensor Device Propagation for aten.mm.default, - # found two different devices cuda:0, cuda:1 - # On a single device every .to() here is a no-op and the result is - # bit-identical to before. - chunk_hidden_states = chunk_hidden_states.to(device = lm_head.device, dtype = lm_head.dtype) - chunk_index = chunk_index.to(lm_head.device) - if fp32_lm_head: - # The GRPO loss runs this under bf16 / fp16 autocast, which would downcast the matmul. - with torch.autocast(device_type = chunk_hidden_states.device.type, enabled = False): - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - elif lm_head_bias is None: - chunk_logits = chunk_hidden_states @ lm_head.t() - else: - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - - if logit_scale_multiply != 0.0: - chunk_logits = chunk_logits * logit_scale_multiply - if logit_scale_divide != 0.0: - chunk_logits = chunk_logits / logit_scale_divide - if logit_softcapping != 0.0: - chunk_logits = logit_softcapping * torch.tanh(chunk_logits / logit_softcapping) - - chunk_logits = chunk_logits.to(torch.float32) - - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - per_token_logps = _truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - ) - else: - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - # Return to the caller's device so the concatenation below and every - # downstream consumer see the device they started on. - all_per_token_logps.append(per_token_logps.to(hidden_states.device)) - - all_per_token_logps = torch.concat(all_per_token_logps) - - all_per_token_logps = all_per_token_logps.reshape((hidden_states.shape[0], hidden_states.shape[1])) - return all_per_token_logps - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_selective_log_softmax( - logits, - index, - temperature: float = 1.0, - chunks: int = 4, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, -): - # Fixed `chunks` row ranges; see chunked_hidden_states_selective_log_softmax. - flat_logits = logits.reshape(-1, logits.shape[-1]) - flat_index = index.reshape(-1) - # Fewer rows than chunks: one row per range, as torch.chunk did, never empty ranges. - chunks = max(1, min(chunks, flat_logits.shape[0])) - bounds = [(i * flat_logits.shape[0]) // chunks for i in range(chunks + 1)] - chunked_logits = [flat_logits[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - all_per_token_logps = [] - # Per-chunk selective_log_softmax. - for chunk_logits, chunk_index in zip(chunked_logits, chunked_index): - chunk_logits = chunk_logits.to(torch.float32) - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - all_per_token_logps.append(_truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - )) - continue - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - all_per_token_logps.append(per_token_logps) - pass - all_per_token_logps = torch.concat(all_per_token_logps) - all_per_token_logps = all_per_token_logps.reshape((logits.shape[0], logits.shape[1])) - return all_per_token_logps - -def calculate_pad_tokens_in_prompt( - input_ids: torch.Tensor, - logits_to_keep: int, - pad_token_id: int -) -> torch.Tensor: - """Count left-padded tokens per sequence, e.g. [pad, pad, pad, cat] -> 3.""" - if logits_to_keep >= input_ids.shape[1]: - raise ValueError("logits_to_keep must be smaller than the sequence length.") - - prompt_section = input_ids[:, :-logits_to_keep] - - padding_mask = (prompt_section == pad_token_id) - - pad_token_counts = padding_mask.sum(dim=1) - - return pad_token_counts - -def create_completion_attention_mask( - completion_input_ids: torch.Tensor, - left_pad_tokens_per_prompt: torch.Tensor, - max_left_pad: int, - pad_token_id: int -) -> torch.Tensor: - """Build a completion mask that zeros leading prompt and trailing pad tokens. - - For [p,p,p,c,c,c,pad,pad,pad] (p=sliced prompt, c=completion, pad=padding) - this returns [0,0,0,1,1,1,0,0,0]. - """ - batch_size, completion_len = completion_input_ids.shape - device = completion_input_ids.device - - num_tokens_to_mask = max_left_pad - left_pad_tokens_per_prompt - - indices = torch.arange(completion_len, device=device).unsqueeze(0) - shift_mask = indices >= num_tokens_to_mask.unsqueeze(1) - - non_padding_mask = (completion_input_ids != pad_token_id) - - final_mask = shift_mask & non_padding_mask - - return final_mask - -def left_pack_padding(tensor: torch.Tensor, pad_id: int) -> torch.Tensor: - """Move all padding tokens in each sequence to the right.""" - mask = (tensor != pad_id) - # stable=True since the binary mask is unordered. - sorted_indices = torch.argsort(mask, dim=1, descending=True, stable=True) - packed_tensor = torch.gather(tensor, 1, sorted_indices) - return packed_tensor - -def align_logprobs_with_mask( - logprob_tensor: torch.Tensor, - attention_mask: torch.Tensor, - pad_value: float = 0.0 -) -> torch.Tensor: - """Align a log probability tensor with a given attention mask.""" - - device = logprob_tensor.device - batch_size, logprob_seq_len = logprob_tensor.shape - mask_seq_len = attention_mask.shape[1] - - padded_logprobs = torch.full( - attention_mask.shape, - fill_value=pad_value, - dtype=logprob_tensor.dtype, - device=device - ) - - left_pad_counts = torch.argmax(attention_mask, dim=1) - - cols = torch.arange(logprob_seq_len, device=device) - dest_indices = left_pad_counts.unsqueeze(1) + cols - - # Destination row indices, shape [batch_size, logprob_seq_len]. - row_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand_as(dest_indices) - - # Keep only in-bounds destinations, then scatter via advanced indexing. - valid_mask = dest_indices < mask_seq_len - valid_rows = row_indices[valid_mask] - valid_cols = dest_indices[valid_mask] - valid_vals = logprob_tensor[valid_mask] - padded_logprobs[valid_rows, valid_cols] = valid_vals - - return padded_logprobs - -def align_completion_tool_mask( - tool_mask: torch.Tensor, - completion_mask: torch.Tensor, -) -> torch.Tensor: - """Align a raw completion-length tool/env mask with Unsloth's repacked loss mask.""" - if tool_mask is None: - return completion_mask - if tool_mask.shape[0] != completion_mask.shape[0]: - raise ValueError("tool_mask batch size must match completion_mask batch size.") - - tool_mask = tool_mask.to(device=completion_mask.device) - if tool_mask.shape == completion_mask.shape: - aligned_tool_mask = tool_mask - else: - aligned_tool_mask = align_logprobs_with_mask( - tool_mask, - completion_mask, - pad_value=0, - ) - return completion_mask * aligned_tool_mask.to(dtype=completion_mask.dtype) - -def autotune_batch_and_chunks( - total_input_rows, - seq_len, - hidden_size, - vocab_size, - dtype_bytes=16, - multiplier=None -): - if multiplier is None: - final_m = max(4, seq_len // 4096) - else: - final_m = multiplier - - if torch.cuda.is_available(): - free_bytes, _ = torch.cuda.mem_get_info() - limit_gb = (free_bytes / (1024**3))*.80 - elif hasattr(torch, "xpu") and torch.xpu.is_available(): - # XPU: estimate free memory as total - reserved. - total_mem = torch.xpu.get_device_properties(0).total_memory - reserved_mem = torch.xpu.memory_reserved() - free_bytes = total_mem - reserved_mem - limit_gb = (free_bytes / (1024**3)) * 0.80 - else: - # Fallback: assume 8GB available. - limit_gb = 8.0 - - bytes_to_gb = 1024**3 - - b_vals = torch.arange(total_input_rows, 0, -1, device='cpu', dtype=torch.float32) - - hidden_gb = (b_vals * seq_len * hidden_size * dtype_bytes) / bytes_to_gb - - base_logits = ((b_vals/total_input_rows) * b_vals * seq_len * vocab_size * dtype_bytes) / bytes_to_gb - logits_gb = base_logits / final_m - - total_mem_gb = hidden_gb + logits_gb - - valid_mask = total_mem_gb <= limit_gb - valid_indices = torch.nonzero(valid_mask, as_tuple=False) - - if valid_indices.shape[0] == 0: - #This means your GPU will OOM - # Capped at the row count: unsloth's no-grad pass divides rows by this without max(1, ...). - return max(1, min(4, total_input_rows)), final_m - - best_idx = valid_indices[0].item() - final_b = int(b_vals[best_idx].item()) - - return final_b, final_m - -def sanitize_logprob(logprob): - """Local port of trl.scripts.vllm_serve.sanitize_logprob. - Filters NaN logprobs from vLLM outputs.""" - value = logprob.logprob - if math.isnan(value): - logging.getLogger(__name__).warning( - f"Generated NaN logprob, token logprob '{logprob}' will be ignored" - ) - return None - return value -@dataclass -class UnslothRLOOConfig(RLOOConfig): - """ - - Configuration class for the [`RLOOTrainer`]. - - This class includes only the parameters that are specific to RLOO training. For a full list of training arguments, - please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may - differ from those in [`~transformers.TrainingArguments`]. - - Using [`~transformers.HfArgumentParser`] we can turn this class into - [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the - command line. - - Parameters: - > Parameters that control the model and reference model - - model_init_kwargs (`str` or `dict[str, Any]`, *optional*): - Keyword arguments for [`~transformers.AutoModelForCausalLM.from_pretrained`], used when the `model` - argument of the [`RLOOTrainer`] is provided as a string. The `revision` value is also used when loading - processing classes. - trust_remote_code (`bool`, *optional*, defaults to `False`): - Whether to allow loading models and tokenizers that ship custom Python code from the Hub. Forwarded to - [`~transformers.AutoModelForCausalLM.from_pretrained`] and [`~transformers.AutoProcessor.from_pretrained`]. - Also applied to reward-model and reward-tokenizer loads. - router_aux_loss_coef (`float`, *optional*): - Coefficient of the load-balancing auxiliary loss for Mixture-of-Experts (MoE) models, added to the training - loss with this weight. When not set, the value declared by the model config is used. Set to `0.0` to - disable it. Fails when used with a non-MoE model. - disable_dropout (`bool`, *optional*, defaults to `False`): - Whether to disable dropout in the model. This is useful for training with a reference model, as it prevents - the model from generating different logprobs for the same input. - - > Parameters that control the data preprocessing - - remove_unused_columns (`bool`, *optional*, defaults to `False`): - Whether to only keep the column `"prompt"` in the dataset. If you use a custom reward function that - requires any column other than `"prompts"` and `"completions"`, you should keep this to `False`. - num_generations (`int`, *optional*, defaults to `2`): - Number of generations per prompt to sample. The effective batch size (num_processes * per_device_batch_size - * gradient_accumulation_steps) must be evenly divisible by this value. - num_generations_eval (`int`, *optional*): - Number of generations to sample during evaluation. This allows using fewer generations during evaluation to - save computation. If `None`, uses the value of `num_generations`. - max_completion_length (`int` or `None`, *optional*, defaults to `512`): - Maximum length of the generated completion. - ds3_gather_for_generation (`bool`, *optional*, defaults to `True`): - This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for generation, - improving generation speed. However, disabling this option allows training models that exceed the VRAM - capacity of a single GPU, albeit at the cost of slower generation. Disabling this option is not compatible - with vLLM generation. - shuffle_dataset (`bool`, *optional*, defaults to `True`): - Whether to shuffle the training dataset. - pad_to_multiple_of (`int`, *optional*): - If set, the prompts ids and completions ids will be padded to a multiple of this value. - - > Parameters that control generation - - generation_batch_size (`int`, *optional*): - Batch size to use for generation. If `None`, it defaults to the effective training batch size: - `per_device_train_batch_size * num_processes * steps_per_generation`. In other words, there is one - generation batch processed per optimization step. Mutually exclusive with `steps_per_generation`. - steps_per_generation (`int`, *optional*): - Number of steps per generation. If `None`, it defaults to `gradient_accumulation_steps`. Mutually exclusive - with `generation_batch_size`. - temperature (`float`, *optional*, defaults to `1.0`): - Temperature for sampling. The higher the temperature, the more random the completions. - top_p (`float`, *optional*, defaults to `1.0`): - Float that controls the cumulative probability of the top tokens to consider. Must be in (0, 1]. Set to - `1.0` to consider all tokens. - top_k (`int`, *optional*, defaults to `0`): - Number of highest probability vocabulary tokens to keep for top-k-filtering. If `0`, top-k-filtering is - disabled and all tokens are considered. - min_p (`float`, *optional*): - Minimum token probability, which will be scaled by the probability of the most likely token. It must be a - value between `0.0` and `1.0`. Typical values are in the `0.01-0.2` range. - generation_kwargs (`dict[str, Any]`, *optional*): - Additional keyword arguments to pass to [`~transformers.GenerationConfig`] (if using transformers) or - `SamplingParams` (if using vLLM) when sampling completions. This can be used to further customize the - generation behavior, such as setting `suppress_tokens`, `num_beams`, etc. If it contains keys that conflict - with the other generation parameters (like `min_p`, `top_p`, etc.), they will override them. - chat_template_kwargs (`dict[str, Any]`, *optional*): - Additional keyword arguments to pass to the `apply_chat_template` function when generating completions. - repetition_penalty (`float`, *optional*, defaults to `1.0`): - Float that penalizes new tokens based on whether they appear in the prompt and the generated text so far. - Values > `1.0` encourage the model to use new tokens, while values < `1.0` encourage the model to repeat - tokens. - cache_implementation (`str`, *optional*): - Implementation of the cache method for faster generation when `use_vllm` is set to `False`. - - > Parameters that control generation acceleration powered by vLLM - - use_vllm (`bool`, *optional*, defaults to `False`): - Whether to use vLLM for generating completions. If set to `True`, the trainer will use vLLM for generation - instead of the default model.generate(). Requires `vllm` to be installed. - vllm_mode (`str`, *optional*, defaults to `"colocate"`): - Mode to use for vLLM integration when `use_vllm` is set to `True`. Must be one of `"server"` or - `"colocate"`. - - - `"server"`: The trainer will send generation requests to a separate vLLM server. Make sure a vLLM server - is running (start with `vllm serve`). - - `"colocate"`: vLLM will run in the same process and share the training GPUs. This avoids the need for a - separate server but may cause resource contention with training. - vllm_model_impl (`str`, *optional*, defaults to `"vllm"`): - Model implementation to use for vLLM. Must be one of `"transformers"` or `"vllm"`. `"transformers"`: Use - the `transformers` backend for model implementation. `"vllm"`: Use the `vllm` library for model - implementation. - vllm_enable_sleep_mode (`bool`, *optional*, defaults to `False`): - Enable vLLM sleep mode to offload weights/cache during the optimizer step. Keeps GPU memory usage low, but - waking the engine adds host–device transfer latency. - vllm_structured_outputs_regex (`str`, *optional*): - Regex for vLLM structured outputs. If `None` (default), structured outputs is disabled. - - > Parameters that control the vLLM server (only used when `vllm_mode` is `"server"`) - - vllm_server_base_url (`str`, *optional*): - Base URL for the vLLM server (e.g., `"http://localhost:8000"`). If provided, `vllm_server_host` and - `vllm_server_port` are ignored. - vllm_server_host (`str`, *optional*, defaults to `"0.0.0.0"`): - Host of the vLLM server to connect to. Ignored if `vllm_server_base_url` is provided. - vllm_server_port (`int`, *optional*, defaults to `8000`): - Port of the vLLM server to connect to. Ignored if `vllm_server_base_url` is provided. - vllm_server_timeout (`float`, *optional*, defaults to `240.0`): - Total timeout duration in seconds to wait for the vLLM server to be up. If the server is not up after the - timeout, a `ConnectionError` is raised. - vllm_group_port (`int`, *optional*, defaults to `51216`): - Port number for the weight update group. This is used to communicate with the vLLM server. Unless the port - is occupied, there is no need to change it. - - > Parameters that control colocated vLLM execution (only used when `vllm_mode` is `"colocate"`) - - vllm_gpu_memory_utilization (`float`, *optional*, defaults to `0.3`): - Control the GPU memory utilization for vLLM. This setting only applies when `vllm_mode` is set to - `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when - launching the vLLM server via the `--vllm_gpu_memory_utilization` flag. - vllm_max_model_length (`int`, *optional*): - Context window for vLLM. Set it to at least the maximum prompt length in the dataset plus - `max_completion_length`; if omitted, it is inferred from the model config. - vllm_tensor_parallel_size (`int`, *optional*, defaults to `1`): - Control the tensor parallel size for vLLM. This setting only applies when `vllm_mode` is set to - `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when - launching the vLLM server via the `--vllm_tensor_parallel_size` flag. - - > Parameters that control the training - - beta (`float`, *optional*, defaults to `0.05`): - KL coefficient. If `0.0`, the reference model is not loaded, reducing memory usage and improving training - speed. - num_iterations (`int`, *optional*, defaults to `1`): - Number of iterations per batch (denoted as μ in the algorithm). - epsilon (`float`, *optional*, defaults to `0.2`): - Epsilon value for clipping. - epsilon_high (`float`, *optional*): - Upper-bound epsilon value for clipping. If not specified, it defaults to the same value as the lower-bound - specified in argument `epsilon`. Paper [DAPO](https://huggingface.co/papers/2503.14476) recommends `0.28`. - reward_weights (`list[float]`, *optional*): - Weights for each reward function. Must match the number of reward functions. If `None`, all rewards are - weighted equally with weight `1.0`. - normalize_advantages (`bool`, *optional*, defaults to `False`): - Whether to normalize advantages. Normalization is done per generation batch to have mean `0.0` and standard - deviation of `1.0`. - reward_clip_range (`tuple[float, float]`, *optional*): - Clip range for rewards as (min, max). If `None`, no clipping is applied. - mask_truncated_completions (`bool`, *optional*, defaults to `False`): - When enabled, truncated completions are excluded from the loss calculation, preventing them from being - incorrectly penalized and introducing noise during training. According to the - [DAPO](https://huggingface.co/papers/2503.14476) paper, this is a good practice for training stability. - sync_ref_model (`bool`, *optional*, defaults to `False`): - Whether to synchronize the reference model with the active model every `ref_model_sync_steps` steps, using - the `ref_model_mixup_alpha` parameter. This synchronization originates from the - [TR-DPO](https://huggingface.co/papers/2404.09656) paper. - ref_model_mixup_alpha (`float`, *optional*, defaults to `0.6`): - α parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which controls the mix - between the current policy and the previous reference policy during updates. The reference policy is - updated according to the equation: `π_ref = α * π_θ + (1 - α) * π_ref_prev`. To use this parameter, you - must set `sync_ref_model=True`. - ref_model_sync_steps (`int`, *optional*, defaults to `512`): - τ parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which determines how - frequently the current policy is synchronized with the reference policy. To use this parameter, you must - set `sync_ref_model=True`. - - > Parameters that control the logging - - log_completions (`bool`, *optional*, defaults to `False`): - Whether to log a sample of (prompt, completion) pairs every `logging_steps` steps. If `rich` is installed, - it prints the sample. If `wandb` and/or `trackio` logging is enabled, it logs it to `wandb` and/or - `trackio`. - log_multimodal (`bool`, *optional*, defaults to `True`): - Whether to log multimodal content (images, videos, etc.) together with completions. Disable this to reduce - log size when using high-resolution multimodal data. - num_completions_to_print (`int`, *optional*): - Number of completions to print with `rich`. If `None`, all completions are logged. - log_unique_prompts (`bool`, *optional*, defaults to `False`): - Whether to log unique prompts. If `True`, only unique prompts are logged. If `False`, all prompts are - logged. - - > Parameters that control generation acceleration powered by transformers continuous batching - - use_transformers_continuous_batching (`bool`, *optional*, defaults to `False`): - Whether to use transformers' continuous batching engine for generating completions. Requires - `transformers>=5.8.0`. - transformers_continuous_batching_config (`dict`, *optional*): - Keyword arguments for [`~transformers.generation.ContinuousBatchingConfig`]. - - > Deprecated parameters - - use_transformers_paged: - - - - Parameter `use_transformers_paged` is deprecated and will be removed in version v2.0.0. Use - `use_transformers_continuous_batching` instead. - - - - > [!NOTE] - > These parameters have default values different from [`~transformers.TrainingArguments`]: - > - `logging_steps`: Defaults to `10` instead of `500`. - > - `gradient_checkpointing`: Defaults to `True` instead of `False`. - > - `bf16`: Defaults to `True` if `fp16` is not set, instead of `False`. - > - `learning_rate`: Defaults to `1e-6` instead of `5e-5`. - - """ - _unsloth_accepts_adamw_flash = True - vllm_sampling_params: Optional[Any] = field( - default = None, - metadata = {'help': 'vLLM SamplingParams'}, - ) - unsloth_num_chunks : Optional[int] = field( - default = -1, - metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'}, - ) - unsloth_logit_chunk_multiplier : Optional[int] = field( - default = None, - metadata = {'help': 'Multiplier for chunked logit computations.'}, - ) - unsloth_grpo_mini_batch : Optional[int] = field( - default = None, - metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'}, - ) - - def __init__( - self, - output_dir = None, - per_device_train_batch_size = 4, - num_train_epochs = 3.0, - max_steps = -1, - learning_rate = 5e-05, - lr_scheduler_type = 'linear', - lr_scheduler_kwargs = None, - warmup_steps = 0.1, - optim = 'adamw_8bit', - optim_args = None, - weight_decay = 0.001, - adam_beta1 = 0.9, - adam_beta2 = 0.999, - adam_epsilon = 1e-08, - optim_target_modules = None, - gradient_accumulation_steps = 2, - average_tokens_across_devices = True, - max_grad_norm = 1.0, - label_smoothing_factor = 0.0, - bf16 = False, - fp16 = False, - bf16_full_eval = False, - fp16_full_eval = False, - tf32 = None, - gradient_checkpointing = True, - gradient_checkpointing_kwargs = None, - torch_compile = False, - torch_compile_backend = None, - torch_compile_mode = None, - use_liger_kernel = False, - liger_kernel_config = None, - use_cache = False, - neftune_noise_alpha = None, - torch_empty_cache_steps = 250, - auto_find_batch_size = False, - logging_strategy = 'steps', - logging_steps = 1, - logging_first_step = False, - log_on_each_node = True, - logging_nan_inf_filter = False, - include_num_input_tokens_seen = False, - log_level = 'passive', - log_level_replica = 'warning', - disable_tqdm = None, - report_to = 'none', - run_name = None, - project = 'huggingface', - trackio_space_id = 'trackio', - eval_strategy = 'no', - eval_steps = None, - eval_delay = 0, - per_device_eval_batch_size = 4, - prediction_loss_only = False, - eval_on_start = False, - eval_do_concat_batches = True, - eval_use_gather_object = False, - eval_accumulation_steps = 2, - batch_eval_metrics = False, - save_only_model = False, - save_strategy = 'steps', - save_steps = 500, - save_on_each_node = False, - save_total_limit = None, - enable_jit_checkpoint = False, - push_to_hub = False, - hub_token = None, - hub_private_repo = None, - hub_model_id = None, - hub_strategy = 'every_save', - hub_always_push = False, - hub_revision = None, - load_best_model_at_end = False, - metric_for_best_model = None, - greater_is_better = None, - ignore_data_skip = False, - restore_callback_states_from_checkpoint = False, - full_determinism = False, - seed = 3407, - data_seed = None, - use_cpu = False, - accelerator_config = None, - parallelism_config = None, - dataloader_drop_last = False, - dataloader_num_workers = 0, - dataloader_pin_memory = True, - dataloader_persistent_workers = False, - dataloader_prefetch_factor = None, - remove_unused_columns = False, - label_names = None, - train_sampling_strategy = 'random', - length_column_name = 'length', - ddp_find_unused_parameters = None, - ddp_bucket_cap_mb = None, - ddp_broadcast_buffers = None, - ddp_backend = None, - ddp_timeout = 1800, - fsdp = None, - fsdp_config = None, - deepspeed = None, - debug = '', - skip_memory_metrics = True, - do_train = False, - do_eval = False, - do_predict = False, - resume_from_checkpoint = None, - warmup_ratio = None, - logging_dir = None, - local_rank = -1, - model_init_kwargs = None, - trust_remote_code = False, - router_aux_loss_coef = None, - disable_dropout = False, - num_generations = 8, - num_generations_eval = None, - max_completion_length = 512, - ds3_gather_for_generation = True, - shuffle_dataset = True, - pad_to_multiple_of = None, - generation_batch_size = None, - steps_per_generation = None, - temperature = 1.0, - top_p = 1.0, - top_k = 0, - min_p = None, - generation_kwargs = {}, - chat_template_kwargs = None, - repetition_penalty = 1.0, - cache_implementation = None, - use_vllm = False, - vllm_mode = 'colocate', - vllm_model_impl = 'vllm', - vllm_enable_sleep_mode = False, - vllm_structured_outputs_regex = None, - vllm_server_base_url = None, - vllm_server_host = '0.0.0.0', - vllm_server_port = 8000, - vllm_server_timeout = 240.0, - vllm_group_port = 51216, - vllm_gpu_memory_utilization = 0.3, - vllm_max_model_length = None, - vllm_tensor_parallel_size = 1, - beta = 0.05, - num_iterations = 1, - epsilon = 0.2, - epsilon_high = None, - reward_weights = None, - normalize_advantages = False, - reward_clip_range = None, - mask_truncated_completions = False, - sync_ref_model = False, - ref_model_mixup_alpha = 0.6, - ref_model_sync_steps = 512, - log_completions = False, - log_multimodal = True, - num_completions_to_print = None, - log_unique_prompts = False, - use_transformers_continuous_batching = False, - transformers_continuous_batching_config = None, - use_transformers_paged = False, - vllm_sampling_params = None, - unsloth_num_chunks = -1, - unsloth_logit_chunk_multiplier = None, - unsloth_grpo_mini_batch = None, - - **kwargs, - ): - if learning_rate < 1e-7: print(f'Unsloth: Your learning rate of `{learning_rate}` is too small and less than 1e-7! Consider increasing it, otherwise gradient updates will be close to 0!') - if learning_rate > 1: print(f'Unsloth: Your learning rate of `{learning_rate}` is way too larger > 1! Consider decreasing it to 1e-1, otherwise gradient updates will explode!') - _unsloth_fsdp2 = os.environ.get('FSDP_VERSION', '').strip() == '2' - _unsloth_fsdp_cfg = fsdp_config - if fsdp and isinstance(_unsloth_fsdp_cfg, str) and os.path.isfile(_unsloth_fsdp_cfg): - try: - import json as _unsloth_json - with open(_unsloth_fsdp_cfg, encoding = 'utf-8') as _unsloth_f: - _unsloth_fsdp_cfg = _unsloth_json.load(_unsloth_f) - except Exception: - _unsloth_fsdp_cfg = None - if fsdp and _unsloth_fsdp_cfg is None: - _unsloth_fsdp_cfg = {} - if fsdp and isinstance(_unsloth_fsdp_cfg, dict): - _unsloth_fsdp2 = _unsloth_fsdp2 or str(_unsloth_fsdp_cfg.get('fsdp_version', _unsloth_fsdp_cfg.get('version', 1))) == '2' - _unsloth_flash_adamw = str(getattr(optim, 'value', optim)).strip().lower() == 'adamw_flash' - if _unsloth_flash_adamw: optim = 'adamw_torch' - if _unsloth_fsdp2 and str(getattr(optim, 'value', optim)) == 'adamw_8bit': - print("Unsloth: FSDP2 shards parameters as DTensors, which bitsandbytes' 8-bit optimizers cannot step. Switching optim = 'adamw_8bit' to 'adamw_torch_fused'.") - optim = 'adamw_torch_fused' - if num_train_epochs is None: - num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override - if output_dir is None and save_strategy == 'steps' and save_steps == 500: - output_dir = 'unsloth_training_checkpoints' - save_strategy = 'no' - if os.environ.get('UNSLOTH_ENABLE_FLEX_ATTENTION', '0') == '1': - from unsloth_zoo.flex_attention import HAS_FLEX_ATTENTION - if HAS_FLEX_ATTENTION and pad_to_multiple_of is None: - from unsloth_zoo.flex_attention import FLEX_ATTENTION_BLOCK_SIZE - pad_to_multiple_of = FLEX_ATTENTION_BLOCK_SIZE - - if steps_per_generation is None and generation_batch_size is None: - ga = gradient_accumulation_steps - world_size = int(os.environ.get('WORLD_SIZE', '1')) - if (ga * world_size * per_device_train_batch_size) % num_generations != 0: - print('Unsloth: We now expect `per_device_train_batch_size` * `gradient_accumulation_steps` * `world_size` to be a multiple of `num_generations`.\nWe will change the batch size of ' + str(per_device_train_batch_size) + ' to the `num_generations` of ' + str(num_generations)) - per_device_train_batch_size = num_generations - - if eval_steps is not None and eval_strategy != 'steps': - print(f'Unsloth: `eval_steps = {eval_steps}` is ignored because `eval_strategy` is {getattr(eval_strategy, "value", eval_strategy)!r}. Set `eval_strategy = "steps"` to evaluate every `eval_steps` steps.') - - if temperature <= 0: - raise ValueError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.') - elif temperature >= 10: - raise ValueError('Unsloth: Please set a positive non-zero temperature less than 10, since sampling will be quite erratic.') - - - # One dict so the filter sees the mirrored parameters AND `**kwargs`: - # filtering kwargs alone would double-bind any argument TRL renamed, - # since the new name is itself a mirrored parameter. - _unsloth_config_arguments = dict( - output_dir = output_dir, - per_device_train_batch_size = per_device_train_batch_size, - num_train_epochs = num_train_epochs, - max_steps = max_steps, - learning_rate = learning_rate, - lr_scheduler_type = lr_scheduler_type, - lr_scheduler_kwargs = lr_scheduler_kwargs, - warmup_steps = warmup_steps, - optim = optim, - optim_args = optim_args, - weight_decay = weight_decay, - adam_beta1 = adam_beta1, - adam_beta2 = adam_beta2, - adam_epsilon = adam_epsilon, - optim_target_modules = optim_target_modules, - gradient_accumulation_steps = gradient_accumulation_steps, - average_tokens_across_devices = average_tokens_across_devices, - max_grad_norm = max_grad_norm, - label_smoothing_factor = label_smoothing_factor, - bf16 = bf16, - fp16 = fp16, - bf16_full_eval = bf16_full_eval, - fp16_full_eval = fp16_full_eval, - tf32 = tf32, - gradient_checkpointing = gradient_checkpointing, - gradient_checkpointing_kwargs = gradient_checkpointing_kwargs, - torch_compile = torch_compile, - torch_compile_backend = torch_compile_backend, - torch_compile_mode = torch_compile_mode, - use_liger_kernel = use_liger_kernel, - liger_kernel_config = liger_kernel_config, - use_cache = use_cache, - neftune_noise_alpha = neftune_noise_alpha, - torch_empty_cache_steps = torch_empty_cache_steps, - auto_find_batch_size = auto_find_batch_size, - logging_strategy = logging_strategy, - logging_steps = logging_steps, - logging_first_step = logging_first_step, - log_on_each_node = log_on_each_node, - logging_nan_inf_filter = logging_nan_inf_filter, - include_num_input_tokens_seen = include_num_input_tokens_seen, - log_level = log_level, - log_level_replica = log_level_replica, - disable_tqdm = disable_tqdm, - report_to = report_to, - run_name = run_name, - project = project, - trackio_space_id = trackio_space_id, - eval_strategy = eval_strategy, - eval_steps = eval_steps, - eval_delay = eval_delay, - per_device_eval_batch_size = per_device_eval_batch_size, - prediction_loss_only = prediction_loss_only, - eval_on_start = eval_on_start, - eval_do_concat_batches = eval_do_concat_batches, - eval_use_gather_object = eval_use_gather_object, - eval_accumulation_steps = eval_accumulation_steps, - batch_eval_metrics = batch_eval_metrics, - save_only_model = save_only_model, - save_strategy = save_strategy, - save_steps = save_steps, - save_on_each_node = save_on_each_node, - save_total_limit = save_total_limit, - enable_jit_checkpoint = enable_jit_checkpoint, - push_to_hub = push_to_hub, - hub_token = hub_token, - hub_private_repo = hub_private_repo, - hub_model_id = hub_model_id, - hub_strategy = hub_strategy, - hub_always_push = hub_always_push, - hub_revision = hub_revision, - load_best_model_at_end = load_best_model_at_end, - metric_for_best_model = metric_for_best_model, - greater_is_better = greater_is_better, - ignore_data_skip = ignore_data_skip, - restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint, - full_determinism = full_determinism, - seed = seed, - data_seed = data_seed, - use_cpu = use_cpu, - accelerator_config = accelerator_config, - parallelism_config = parallelism_config, - dataloader_drop_last = dataloader_drop_last, - dataloader_num_workers = dataloader_num_workers, - dataloader_pin_memory = dataloader_pin_memory, - dataloader_persistent_workers = dataloader_persistent_workers, - dataloader_prefetch_factor = dataloader_prefetch_factor, - remove_unused_columns = remove_unused_columns, - label_names = label_names, - train_sampling_strategy = train_sampling_strategy, - length_column_name = length_column_name, - ddp_find_unused_parameters = ddp_find_unused_parameters, - ddp_bucket_cap_mb = ddp_bucket_cap_mb, - ddp_broadcast_buffers = ddp_broadcast_buffers, - ddp_backend = ddp_backend, - ddp_timeout = ddp_timeout, - fsdp = fsdp, - fsdp_config = fsdp_config, - deepspeed = deepspeed, - debug = debug, - skip_memory_metrics = skip_memory_metrics, - do_train = do_train, - do_eval = do_eval, - do_predict = do_predict, - resume_from_checkpoint = resume_from_checkpoint, - warmup_ratio = warmup_ratio, - logging_dir = logging_dir, - local_rank = local_rank, - model_init_kwargs = model_init_kwargs, - trust_remote_code = trust_remote_code, - router_aux_loss_coef = router_aux_loss_coef, - disable_dropout = disable_dropout, - num_generations = num_generations, - num_generations_eval = num_generations_eval, - max_completion_length = max_completion_length, - ds3_gather_for_generation = ds3_gather_for_generation, - shuffle_dataset = shuffle_dataset, - pad_to_multiple_of = pad_to_multiple_of, - generation_batch_size = generation_batch_size, - steps_per_generation = steps_per_generation, - temperature = temperature, - top_p = top_p, - top_k = top_k, - min_p = min_p, - generation_kwargs = generation_kwargs, - chat_template_kwargs = chat_template_kwargs, - repetition_penalty = repetition_penalty, - cache_implementation = cache_implementation, - use_vllm = use_vllm, - vllm_mode = vllm_mode, - vllm_model_impl = vllm_model_impl, - vllm_enable_sleep_mode = vllm_enable_sleep_mode, - vllm_structured_outputs_regex = vllm_structured_outputs_regex, - vllm_server_base_url = vllm_server_base_url, - vllm_server_host = vllm_server_host, - vllm_server_port = vllm_server_port, - vllm_server_timeout = vllm_server_timeout, - vllm_group_port = vllm_group_port, - vllm_gpu_memory_utilization = vllm_gpu_memory_utilization, - vllm_max_model_length = vllm_max_model_length, - vllm_tensor_parallel_size = vllm_tensor_parallel_size, - beta = beta, - num_iterations = num_iterations, - epsilon = epsilon, - epsilon_high = epsilon_high, - reward_weights = reward_weights, - normalize_advantages = normalize_advantages, - reward_clip_range = reward_clip_range, - mask_truncated_completions = mask_truncated_completions, - sync_ref_model = sync_ref_model, - ref_model_mixup_alpha = ref_model_mixup_alpha, - ref_model_sync_steps = ref_model_sync_steps, - log_completions = log_completions, - log_multimodal = log_multimodal, - num_completions_to_print = num_completions_to_print, - log_unique_prompts = log_unique_prompts, - use_transformers_continuous_batching = use_transformers_continuous_batching, - transformers_continuous_batching_config = transformers_continuous_batching_config, - use_transformers_paged = use_transformers_paged,**kwargs) - super().__init__(**_unsloth_filter_config_init_kwargs(RLOOConfig, _unsloth_config_arguments, mirrored_from = __class__)) - self.vllm_sampling_params = vllm_sampling_params - self.unsloth_num_chunks = unsloth_num_chunks - if unsloth_grpo_mini_batch is not None: - if self.generation_batch_size >= unsloth_grpo_mini_batch: - self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch - else: - raise ValueError( - f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, " - f"which is self.per_device_train_batch_size * gradient_accumulation_steps." - ) - self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier - - # Unsloth: keep the reentrant checkpoint path - if getattr(self, 'gradient_checkpointing', False): - _gc_kwargs = getattr(self, 'gradient_checkpointing_kwargs', None) or {} - if _gc_kwargs.get('context_fn') is None and not _gc_kwargs.get('debug', False): - _gc_kwargs['use_reentrant'] = True - self.gradient_checkpointing_kwargs = _gc_kwargs - from unsloth.optimizers.flash_adamw import is_flash_adamw_marked, mark_flash_adamw, unmark_flash_adamw - if _unsloth_flash_adamw or is_flash_adamw_marked(self): - mark_flash_adamw(self) - else: - unmark_flash_adamw(self) - -pass - -class _UnslothRLOOTrainer(_BaseTrainer): - """""" - - loss_is_scaled_for_ga = False - - _tag_names = ["trl", "rloo"] - _name = "RLOO" - _paper = { - "title": "Back to Basics: Revisiting REINFORCE-Style Optimization for Learning from Human Feedback in LLMs", - "id": "2402.14740", - # docstyle-ignore - "citation": textwrap.dedent("""\ - @inproceedings{ahmadian2024back, - title = {{Back to Basics: Revisiting REINFORCE-Style Optimization for Learning from Human Feedback in LLMs}}, - author = {Arash Ahmadian and Chris Cremer and Matthias Gall{\'{e}} and Marzieh Fadaee and Julia Kreutzer and Olivier Pietquin and Ahmet {\"{U}}st{\"{u}}n and Sara Hooker}, - year = 2024, - booktitle = {Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), {ACL} 2024, Bangkok, Thailand, August 11-16, 2024}, - pages = {12248--12267}, - publisher = {Association for Computational Linguistics}, - editor = {Lun{-}Wei Ku and Andre Martins and Vivek Srikumar}, - }"""), - } - - def __init__( - self, - model: "str | PreTrainedModel | PeftModel", - reward_funcs: RewardFunc | list[RewardFunc], - args: RLOOConfig | None = None, - train_dataset: Dataset | IterableDataset | None = None, - eval_dataset: Dataset - | IterableDataset - | DatasetDict - | IterableDatasetDict - | dict[str, Dataset | IterableDataset] - | None = None, - processing_class: PreTrainedTokenizerBase | ProcessorMixin | None = None, - reward_processing_classes: PreTrainedTokenizerBase | list[PreTrainedTokenizerBase] | None = None, - callbacks: list[TrainerCallback] | None = None, - optimizers: tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None] = (None, None), - quantization_config: "BitsAndBytesConfig | None" = None, - peft_config: "PeftConfig | None" = None, - ): - - if hasattr(model, 'vllm_engine') and hasattr(args, 'use_vllm'): - _unsloth_fsdp_cfg = getattr(args, 'fsdp_config', None) - if getattr(args, 'fsdp', None) and isinstance(_unsloth_fsdp_cfg, dict) and str(_unsloth_fsdp_cfg.get('version', _unsloth_fsdp_cfg.get('fsdp_version', 1))) == '2': - raise NotImplementedError('Unsloth: `fast_inference = True` (vLLM) is not supported with FSDP2.\nLoad the model with `fast_inference = False` to train with FSDP2, or launch without FSDP2 to use vLLM.') - if (getattr(args, 'use_vllm', False) == False): - args.use_vllm = True - if getattr(args, 'top_k', -1) is None or getattr(args, 'top_k', -1) == 0: - args.top_k = -1 - # Args - if args is None: - model_name = model if isinstance(model, str) else get_config_model_id(model.config) - model_name = model_name.split("/")[-1] - args = RLOOConfig(f"{model_name}-RLOO") - - # Model - # PEFT initializes the adapter weights randomly, so set_seed must be done before creating the model to ensure - # reproducibility. - set_seed(args.seed) - if isinstance(model, str): - model_init_kwargs = dict(args.model_init_kwargs or {}) # copy to avoid mutating model_init_kwargs - if quantization_config is not None: - if "quantization_config" in model_init_kwargs: - raise ValueError( - "You set `quantization_config` both as a trainer argument and in `args.model_init_kwargs`. " - "Please set it in only one place, preferably as a trainer argument." - ) - model_init_kwargs["quantization_config"] = quantization_config - # Distributed training requires device_map=None ["auto" fails] - if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: - model_init_kwargs["device_map"] = None - model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) - model_revision = model_init_kwargs.get("revision") - model = create_model_from_path(model, **model_init_kwargs) - else: - model_revision = None - if args.model_init_kwargs is not None: - logger.warning( - "You passed `model_init_kwargs` to the `RLOOConfig`, but your model is already instantiated. " - "The `model_init_kwargs` will be ignored." - ) - if quantization_config is not None: - logger.warning( - "You passed `quantization_config` to the trainer, but your model is already instantiated. The " - "`quantization_config` will be ignored." - ) - # Non-quantized models do not have the `is_loaded_in_{8,4}bit` attributes, whereas quantized models do - _is_quantized_model = getattr(model, "is_loaded_in_4bit", False) or getattr(model, "is_loaded_in_8bit", False) - - # Processing class - if processing_class is None: - processing_class = AutoProcessor.from_pretrained( - get_config_model_id(model.config), - revision=model_revision, - truncation_side="left", - padding_side="left", - trust_remote_code=args.trust_remote_code, - ) - - if args.use_transformers_continuous_batching and isinstance(processing_class, ProcessorMixin): - raise ValueError( - "`use_transformers_continuous_batching` does not support multimodal models. Use `use_vllm` instead." - ) - - # Handle pad token for processors or tokenizers - if isinstance(processing_class, ProcessorMixin): - self._tokenizer = processing_class.tokenizer - elif isinstance(processing_class, PreTrainedTokenizerBase): - self._tokenizer = processing_class - else: - raise TypeError("The `processing_class` must be either a `PreTrainedTokenizerBase` or a `ProcessorMixin`") - - if self._tokenizer.pad_token is None: - self._tokenizer.pad_token = self._tokenizer.eos_token - - # The model must agree with the tokenizer on the pad token from construction, so mirror it onto the model - # configs. - model.config.get_text_config().pad_token_id = self._tokenizer.pad_token_id - model.generation_config.pad_token_id = self._tokenizer.pad_token_id - - # Every eos id the model declares [e.g. an end-of-turn token] ends a completion, not only the tokenizer's - eos_token_ids = model.generation_config.eos_token_id - if eos_token_ids is None: - eos_token_ids = [] - elif isinstance(eos_token_ids, int): - eos_token_ids = [eos_token_ids] - if self._tokenizer.eos_token_id not in eos_token_ids: - eos_token_ids = [self._tokenizer.eos_token_id, *eos_token_ids] - self.eos_token_ids = eos_token_ids - - # PEFT - if False: - if not is_peft_available(): - raise ImportError( - "You passed `peft_config` but the `peft` library is not installed. " - "Install it with `pip install trl[peft]`." - ) - if not isinstance(peft_config, PeftConfig): - raise TypeError( - f"`peft_config` must be a `peft.PeftConfig` instance (e.g. `peft.LoraConfig`), " - f"got {type(peft_config).__name__}." - ) - if is_peft_model(model): - raise ValueError( - "You passed a `PeftModel` instance together with a `peft_config` to the trainer. Please first merge " - "and unload the existing adapter, save the resulting base model, and then pass that base model along " - "with the new `peft_config` to the trainer." - ) - # Create PEFT model - # ZeRO-3 + PEFT for non-quantized models: - # - PEFT's default autocast_adapter_dtype=True upcasts LoRA adapter params to fp32 even when the base model is bf16. - # - ZeRO-3's _allgather_params_coalesced allocates output buffers using the dtype of the first persistent parameter, - # so mixed-dtype persistent_parameters [bf16 base + fp32 LoRA] cause a TypeError on the first optimizer step. - # - Passing autocast_adapter_dtype=False keeps adapter params in the base model dtype [bf16], fixing the mismatch. - # - This is safe: the fp32 upcast is a QLoRA-specific concern [low-bit quantized base models], not needed for - # non-quantized bf16 training. - # - See: - # - TRL issue: https://github.com/huggingface/trl/issues/6089 - # - Upstream issue: https://github.com/deepspeedai/DeepSpeed/issues/8072 - get_peft_model_kwargs = {} - if args.deepspeed_plugin is not None and args.deepspeed_plugin.zero_stage == 3 and not _is_quantized_model: - get_peft_model_kwargs["autocast_adapter_dtype"] = False - model = get_peft_model(model, peft_config, **get_peft_model_kwargs) - - elif is_peft_model(model): - # If the model is a PEFT model with a pretrained adapter, we need to create a "ref" adapter that is a copy - # of the "default" adapter, so that we can use it as the reference model during RLOO training. Before PEFT - # 0.20.0, only one adapter per model was supported when the LoRA config uses `target_parameters` [see - # peft#3340, fixed in peft#3350], so in that case we skip the "ref" adapter and compute the reference log - # probs with adapters disabled, i.e. with the base model. The fix only allows adapters targeting the same - # parameters, which holds here since the "ref" adapter reuses the "default" config. `target_parameters` - # itself was only added in PEFT 0.17.0, so the version check is bounded on both sides. - default_config = model.peft_config["default"] - if ( - isinstance(default_config, LoraConfig) - and Version("0.17.0") <= Version(peft.__version__) < Version("0.20.0") - and default_config.target_parameters - ): - logger.warning( - "PEFT<0.20.0 can't add a frozen reference adapter alongside one that uses `target_parameters` " - "(peft#3340], so the reference log probs are computed from the base model [adapters disabled]. " - "Upgrade to `peft>=0.20.0` to train against a copy of your adapter instead. If you wrapped the " - "model only to apply LoRA, pass a `peft_config` to the trainer instead; if you wrapped it " - "deliberately (pretrained adapter or custom init), note that the base model matches your adapter " - "only when it's freshly zero-initialized. If it is, this warning is safe to ignore." - ) - else: - model.add_adapter("ref", default_config) - for name, param in model.named_parameters(): - if ".default." in name: - ref_name = name.replace(".default.", ".ref.") - ref_param = model.get_parameter(ref_name) - ref_param.data.copy_(param.data) - - # PEFT + DeepSpeed ZeRO-3 requires reentrant checkpointing. For more details, see - # https://github.com/huggingface/trl/issues/2514#issuecomment-2692152703. - # Can be removed once https://github.com/deepspeedai/DeepSpeed/pull/8130 is merged and released. - if ( - is_peft_model(model) - and args.deepspeed_plugin is not None - and args.deepspeed_plugin.zero_stage == 3 - and args.gradient_checkpointing - ): - args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} - use_reentrant = args.gradient_checkpointing_kwargs.get("use_reentrant") - if use_reentrant is False: - logger.warning( - "You are using PEFT with DeepSpeed ZeRO-3 and gradient checkpointing with `use_reentrant=False`. " - "`use_reentrant` is forced to `True` in this configuration to ensure correct training. To remove " - "this warning, unset `use_reentrant` in `gradient_checkpointing_kwargs` or set it to `True`." - ) - args.gradient_checkpointing_kwargs["use_reentrant"] = True - - # When using gradient checkpointing with PEFT, we need to enable input gradients. transformers.Trainer normally - # handles this, but a bug currently prevents it; see https://github.com/huggingface/transformers/issues/42489 - if is_peft_model(model) and args.gradient_checkpointing: - model.enable_input_require_grads() - - # When using QLoRA, the PEFT adapter weights are converted to bf16 to follow the recommendations from the - # original paper [see https://huggingface.co/papers/2305.14314, paragraph 3]. Normally, this can be done by - # passing `autocast_adapter_dtype=False` to `get_peft_model`, but this option is not yet supported for - # quantized models. See: https://github.com/huggingface/peft/issues/2889 - # The DoRA magnitude vector is excluded: unlike LoRA A/B, its optimizer updates can be smaller than bf16 can - # represent, silently freezing it, see #7268. - if False: - for name, param in model.named_parameters(): - if param.requires_grad and "lora_magnitude_vector" not in name: - param.data = param.data.to(torch.bfloat16) - - # Reward functions - if not isinstance(reward_funcs, list): - reward_funcs = [reward_funcs] - self.reward_func_names = [] - reward_model_revisions = [None] * len(reward_funcs) - for i, reward_func in enumerate(reward_funcs): - if isinstance(reward_func, str): - model_init_kwargs = args.model_init_kwargs or {} - # Distributed training requires device_map=None ["auto" fails] - if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: - model_init_kwargs["device_map"] = None - model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) - reward_model_revisions[i] = model_init_kwargs.get("revision") - reward_funcs[i] = AutoModelForSequenceClassification.from_pretrained( - reward_func, num_labels=1, **model_init_kwargs - ) - if isinstance(reward_funcs[i], nn.Module): # Use Module over PretrainedModel for compat w/ compiled models - self.reward_func_names.append(get_config_model_id(reward_funcs[i].config).split("/")[-1]) - else: - self.reward_func_names.append(get_callable_name(reward_funcs[i])) - self.reward_funcs = reward_funcs - - self._has_async_funcs = any(is_async_callable(func) for func in self.reward_funcs) - if self._has_async_funcs: - self.async_loop_thread, self.async_loop, self.async_loop_ready_event = start_event_loop_in_daemon( - name="RLOOTrainer-AsyncRewardLoop" - ) - # wait until the event loop is running in the daemon thread - self.async_loop_ready_event.wait() - atexit.register(shutdown_event_loop_in_daemon, self.async_loop_thread, self.async_loop) - - # Reward weights - if args.reward_weights is not None: - if len(args.reward_weights) != len(reward_funcs): - raise ValueError( - f"Number of reward weights ({len(args.reward_weights)}) must match number of reward " - f"functions ({len(reward_funcs)})" - ) - self.reward_weights = torch.tensor(args.reward_weights, dtype=torch.float32) - else: - self.reward_weights = torch.ones(len(reward_funcs), dtype=torch.float32) - - # Reward processing class - if reward_processing_classes is None: - reward_processing_classes = [None] * len(reward_funcs) - elif not isinstance(reward_processing_classes, list): - reward_processing_classes = [reward_processing_classes] - if len(reward_processing_classes) != len(reward_funcs): - raise ValueError( - f"The number of reward processing classes ({len(reward_processing_classes)}) must match the number of " - f"reward functions ({len(reward_funcs)})." - ) - - for i, (reward_processing_class, reward_func) in enumerate( - zip(reward_processing_classes, reward_funcs, strict=True) - ): - if isinstance(reward_func, PreTrainedModel): - if reward_processing_class is None: - reward_processing_class = AutoTokenizer.from_pretrained( - get_config_model_id(reward_func.config), - revision=reward_model_revisions[i], - trust_remote_code=args.trust_remote_code, - ) - if reward_processing_class.pad_token_id is None: - reward_processing_class.pad_token = reward_processing_class.eos_token - # The reward model computes the reward for the latest non-padded token in the input sequence. - # So it's important to set the pad token ID to the padding token ID of the processing class. - reward_func.config.get_text_config().pad_token_id = reward_processing_class.pad_token_id - reward_processing_classes[i] = reward_processing_class - - self.reward_processing_classes = reward_processing_classes - - # Training arguments - self.max_completion_length = args.max_completion_length - self.num_generations = args.num_generations - self.num_generations_eval = args.num_generations_eval or self.num_generations - self.chat_template_kwargs = args.chat_template_kwargs or {} - self.temperature = args.temperature - self.top_p = args.top_p - self.top_k = args.top_k - self.min_p = args.min_p - self.repetition_penalty = args.repetition_penalty - self.use_transformers_continuous_batching = args.use_transformers_continuous_batching - if self.use_transformers_continuous_batching: - if not Version(transformers.__version__) >= Version("5.8.0"): - raise ImportError( - "Using `use_transformers_continuous_batching` requires transformers>=5.8.0. " - "Please upgrade with `pip install --upgrade transformers`." - ) - from transformers.generation import ContinuousBatchingConfig - - cb_kwargs = dict(args.transformers_continuous_batching_config or {}) - # The transformers default [0.9] leaves almost no VRAM for the training backward pass; - # use a training-aware default unless the user has set it explicitly. - cb_kwargs.setdefault("max_memory_percent", 0.5) - self.continuous_batching_config = ContinuousBatchingConfig(**cb_kwargs) - else: - self.continuous_batching_config = None - self.pad_to_multiple_of = args.pad_to_multiple_of - self.use_vllm = args.use_vllm - self.vllm_mode = args.vllm_mode - self.vllm_gpu_memory_utilization = args.vllm_gpu_memory_utilization # only applies to colocation mode - self.vllm_tensor_parallel_size = args.vllm_tensor_parallel_size # only applies to colocation mode - self.normalize_advantages = args.normalize_advantages - self.mask_truncated_completions = args.mask_truncated_completions - self.reward_clip_range = args.reward_clip_range - - # Datasets - self.shuffle_dataset = args.shuffle_dataset - - if train_dataset is None: - raise ValueError("`train_dataset` is required") - elif not isinstance(train_dataset, (Dataset, IterableDataset, list, tuple)): - raise TypeError( - f"`train_dataset` must be a `Dataset` or `IterableDataset`, got `{type(train_dataset).__name__}`." - ) - - # Iterable datasets can't be indexed, so the RepeatSampler can't be attached to them. Instead, the sampler's - # ordering is reproduced by streaming [see `get_train_dataloader`/`get_eval_dataloader` and - # `repeat_iterable_dataset`]. This requires `dispatch_batches=False`: with the default dispatch path, batches - # are collated on the main process and Accelerate tries to concatenate the string `prompt` column, which fails; - # `dispatch_batches=False` also lets each process shard the stream into contiguous slices, as the sampler does. - # See https://github.com/huggingface/trl/issues/3213 - uses_iterable_dataset = ( - isinstance(train_dataset, IterableDataset) - or isinstance(eval_dataset, IterableDataset) - or ( - isinstance(eval_dataset, dict) and any(isinstance(ds, IterableDataset) for ds in eval_dataset.values()) - ) - ) - if uses_iterable_dataset: - if args.accelerator_config.dispatch_batches: - raise ValueError( - "Iterable datasets require `dispatch_batches=False`, but it is set to `True` in " - "`accelerator_config`. Please set it to `False`." - ) - args.accelerator_config.dispatch_batches = False - # An iterable train set bakes prompt repeats into the stream, so it must be read by a single worker: multiple - # workers would shard and interleave it, splitting the num_generations groups. Map-style train keeps its workers. - if isinstance(train_dataset, IterableDataset) and args.dataloader_num_workers != 0: - logger.warning( - f"Iterable datasets require `dataloader_num_workers=0` to preserve prompt grouping; overriding the " - f"provided value ({args.dataloader_num_workers})." - ) - args.dataloader_num_workers = 0 - - # Multi-step - self.num_iterations = args.num_iterations - self.epsilon_low = args.epsilon - self.epsilon_high = args.epsilon_high if args.epsilon_high is not None else args.epsilon - - # MoE load-balancing auxiliary loss. `output_router_logits` in the config means the model returns its router - # logits; `router_aux_loss_coef` is the architecture's own coefficient, 0.0 when the config declares none. - text_config = model.config.get_text_config() - coef = args.router_aux_loss_coef - self.router_aux_loss_coef = getattr(text_config, "router_aux_loss_coef", 0.0) if coef is None else coef - if coef and not hasattr(text_config, "output_router_logits"): - raise ValueError( - f"`router_aux_loss_coef` is set to {coef} but {type(model).__name__} is not a Mixture-of-Experts model " - f"that returns its router logits, so there is no auxiliary loss to weight." - ) - self.aux_loss_enabled = hasattr(text_config, "output_router_logits") and self.router_aux_loss_coef != 0.0 - if is_peft_model(model) and isinstance(model.get_output_embeddings(), BaseTunerLayer): - # The log-probabilities are computed by multiplying the hidden states by `lm_head.weight` directly, so an - # adapter on the LM head would be ignored and never trained. - raise ValueError( - "Applying a PEFT adapter to `lm_head` is not supported: the log-probabilities are computed from " - "`lm_head.weight` directly, so the adapter would never be trained. Remove `'lm_head'` from your " - "`target_modules`." - ) - # Tracks the number of iterations [forward + backward passes], including those within a grad accum cycle - self._step = 0 - # Buffer the batch to reuse generated outputs across multiple updates. For more details, see - # `_get_train_sampler` and `_prepare_inputs`. - self._buffered_inputs = None - - # Transformers explicitly set use_reentrant=True in the past to silence a PyTorch warning, but the default was - # never updated once PyTorch switched to recommending use_reentrant=False. Until that change lands upstream - # [see https://github.com/huggingface/transformers/pull/43203] and is released [most likely in 5.0.0], we - # default to the recommended non-reentrant behavior here, while preserving any user-provided value. - if args.gradient_checkpointing and Version(transformers.__version__) < Version("5.0.0"): - args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} - args.gradient_checkpointing_kwargs.setdefault("use_reentrant", False) - - super().__init__( - model=model, - args=args, - data_collator=identity, # No data collation is needed in RLOO - train_dataset=train_dataset, - eval_dataset=eval_dataset, - processing_class=processing_class, - callbacks=callbacks, - optimizers=optimizers, - ) - - # With several GPUs visible and no distributed launcher, `Trainer` wraps the model in `nn.DataParallel`, whose - # replicas would all run the forward bound to the original model. Checked after `Trainer.__init__`, which sets - # `n_gpu` to 1 for a model split across devices with `device_map`. - if args.n_gpu > 1: - raise ValueError( - f"{type(self).__name__} does not support `nn.DataParallel`, which `Trainer` uses when several GPUs are " - "visible to a single process. Launch the script with `accelerate launch` or `torchrun`, or make a " - "single GPU visible with `CUDA_VISIBLE_DEVICES`." - ) - - # Reference model - self.beta = args.beta - if self.beta == 0.0: - # If beta is 0.0, the reference model is not needed - self.ref_model = None - elif is_peft_model(model): - # If PEFT is used, the reference model is not needed since the adapter can be disabled - # to revert to the initial model. - self.ref_model = None - else: - # For deepspeed, fsdp or non-distributed models, create a reference model from scratch - model_init_kwargs = dict(args.model_init_kwargs or {}) # copy to avoid mutating model_init_kwargs - if quantization_config is not None: - model_init_kwargs["quantization_config"] = quantization_config - # Distributed training requires device_map=None ["auto" fails] - if self.args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: - model_init_kwargs["device_map"] = None - model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) - self.ref_model = create_model_from_path(get_config_model_id(self.model.config), **model_init_kwargs) - - # Disable dropout in the models - if args.disable_dropout: - disable_dropout_in_model(model) - if self.ref_model is not None: - disable_dropout_in_model(self.ref_model) - - # Liger's fused linear cross-entropy replaces `model.forward` when training starts, which would drop the fused - # LM head, so only its layer kernels are applied - if args.use_liger_kernel: - warnings.warn( - "`use_liger_kernel=True` is deprecated and will be removed in v2.0.0. Use the Hub kernels instead, " - 'with `model_init_kwargs={"use_kernels": True}`.', - FutureWarning, - stacklevel=2, - ) - liger_kernel_config = args.liger_kernel_config or {} - if liger_kernel_config.get("fused_linear_cross_entropy"): - raise ValueError( - '`liger_kernel_config={"fused_linear_cross_entropy": True}` is not supported: it replaces the ' - "model's forward, which this trainer patches with a fused LM head." - ) - args.liger_kernel_config = {**liger_kernel_config, "fused_linear_cross_entropy": False} - - # Compute the per-token log-probabilities in chunks, without materializing the full logits - # `entropy` feeds the `entropy` metric - add_fused_lm_head( - self.model.get_base_model() if is_peft_model(self.model) else self.model, - temperature=self.temperature, - outputs=("log_probs", "entropy"), - ) - if self.ref_model is not None: - add_fused_lm_head(self.ref_model, temperature=self.temperature) - - # Initialize the metrics - self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} - self._total_train_tokens = 0 - self._current_train_step_time = 0.0 - self.log_completions = args.log_completions - self.log_multimodal = args.log_multimodal - self.log_unique_prompts = args.log_unique_prompts - self.num_completions_to_print = args.num_completions_to_print - # Keep logs sized to the generation batch to record only outputs from the latest model update. - self._logs = { - "images": deque(maxlen=args.generation_batch_size), - "prompt": deque(maxlen=args.generation_batch_size), - "completion": deque(maxlen=args.generation_batch_size), - "rewards": defaultdict(lambda: deque(maxlen=args.generation_batch_size)), - "advantages": deque(maxlen=args.generation_batch_size), - "extra": defaultdict(lambda: deque(maxlen=args.generation_batch_size)), - } - # Buffers for user-logged data from reward functions, flushed after gathering - self._pending_extra_logs = defaultdict(list) - self._pending_metrics = defaultdict(list) - - # Ensure each process receives a unique seed to prevent duplicate completions when generating with - # transformers if num_generations exceeds per_device_train_batch_size. We could skip it if we use vLLM, but - # it's safer to set it in all cases. - set_seed(args.seed, device_specific=True) - - if self.use_vllm: - self.vllm_generation = VLLMGeneration( - model=self.model, - accelerator=self.accelerator, - processing_class=self.processing_class, - mode=args.vllm_mode, - structured_outputs_regex=args.vllm_structured_outputs_regex, - server_base_url=args.vllm_server_base_url, - server_host=args.vllm_server_host, - server_port=args.vllm_server_port, - group_port=args.vllm_group_port, - server_timeout=args.vllm_server_timeout, - tensor_parallel_size=args.vllm_tensor_parallel_size, - gpu_memory_utilization=args.vllm_gpu_memory_utilization, - max_model_length=args.vllm_max_model_length, - max_num_seqs=args.per_device_train_batch_size - * args.vllm_tensor_parallel_size - * args.steps_per_generation, - enable_sleep_mode=args.vllm_enable_sleep_mode, - model_impl=args.vllm_model_impl, - trust_remote_code=args.trust_remote_code, - repetition_penalty=self.repetition_penalty, - temperature=self.temperature, - top_p=self.top_p, - top_k=self.top_k, - min_p=self.min_p, - max_completion_length=self.max_completion_length, - logprobs=None, - generation_kwargs=args.generation_kwargs, - ) - self._last_loaded_step = -1 - else: - generation_kwargs = { - "max_new_tokens": self.max_completion_length, - "do_sample": True, - "pad_token_id": self._tokenizer.pad_token_id, - "bos_token_id": self._tokenizer.bos_token_id, - "eos_token_id": self.eos_token_ids, - "temperature": self.temperature, - "top_p": self.top_p, - "top_k": self.top_k, - "min_p": self.min_p, - "repetition_penalty": self.repetition_penalty, - "cache_implementation": args.cache_implementation, - } - if args.generation_kwargs is not None: - generation_kwargs.update(args.generation_kwargs) - self.generation_config = GenerationConfig(**generation_kwargs, disable_compile=True) - # Keep training-specific generation kwargs to overwrite model's original generation config - self.generation_kwargs = generation_kwargs - - self._dist = DistributedBackend(self.accelerator) - - # Add tags to the model - self.model.add_model_tags(self._tag_names) - - if self.ref_model is not None: - if self.is_deepspeed_enabled: - self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator) - elif self.is_fsdp_enabled: - self.ref_model = prepare_fsdp(self.ref_model, self.accelerator) - else: - self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True) - - if args.sync_ref_model: - if self.beta == 0.0: - raise ValueError( - "You passed `sync_ref_model=True` while `beta=0.0`, which means the reference model is not used " - "during training. Consequently, RLOOTrainer does not create a `ref_model` instance, and there is " - "nothing to synchronize. Please set `sync_ref_model=False`, or set `beta` to a non-zero value." - ) - if is_peft_model(model): - raise NotImplementedError( - "You passed `sync_ref_model=True` while using a PEFT model, which is currently not supported. " - "With PEFT, RLOOTrainer does not keep a separate reference model in memory; instead, it recovers " - "reference behavior by temporarily disabling the adapter. As a result, there is no standalone " - "`ref_model` instance to synchronize. Use `sync_ref_model=False`, or opt for full fine-tuning if " - "you need a synced reference model. If you need `sync_ref_model` to work with PEFT, please open a " - "feature request at https://github.com/huggingface/trl/issues." - ) - self.add_callback(SyncRefModelCallback(ref_model=self.ref_model, accelerator=self.accelerator)) - - for i, reward_func in enumerate(self.reward_funcs): - if isinstance(reward_func, PreTrainedModel): - if self.is_deepspeed_enabled: - self.reward_funcs[i] = prepare_deepspeed(reward_func, self.accelerator) - else: - # set device placement to True to make `prepare_model` move `reward_func` to device when using fsdp - self.reward_funcs[i] = self.accelerator.prepare_model( - reward_func, evaluation_mode=True, device_placement=True - ) - - def _set_signature_columns_if_needed(self): - # If `self.args.remove_unused_columns` is True, non-signature columns are removed. - # By default, this method sets `self._signature_columns` to the model's expected inputs (usually, "input_ids" - # and "attention_mask"). In RLOOTrainer, we preprocess data, so using the model's signature columns doesn't - # work. Instead, we set them to the columns expected by the `training_step` method, hence the override. - if self._signature_columns is None: - self._signature_columns = ["prompt", "image", "images"] - - # This method overrides `Trainer.get_train_dataloader` to support our custom batching strategy. - # Instead of returning a standard per-step batch (i.e., `per_device_batch_size), our dataloader loads an - # *generation* batch (i.e., `per_device_batch_size × steps_per_generation`). This allows us to generate completions - # once every steps_per_generation step—rather than once per accumulation step—which is significantly more - # efficient. Thus, `_prepare_inputs` is called with this *generation* batch, and it handles the splitting internally. - # Maintenance note: this method is a copy-paste of the original `Trainer.get_train_dataloader` with two changes: the - # batch size is multiplied by `steps_per_generation`, and iterable datasets are wrapped (see `repeat_iterable_dataset`). - def get_train_dataloader(self): - dataset = self.train_dataset - if isinstance(dataset, IterableDataset): - # Iterable datasets can't be indexed, so RepeatSampler can't be attached. Reproduce its ordering by - # transforming the stream instead (see `repeat_iterable_dataset`). The full permutation done by - # RepeatSampler becomes a buffered shuffle here. - if self.shuffle_dataset: - dataset = dataset.shuffle(seed=self.args.seed) - dataset = repeat_iterable_dataset( - dataset, - mini_repeat_count=self.num_generations, - batch_size=self.args.generation_batch_size // self.num_generations, - repeat_count=self.num_iterations * self.args.steps_per_generation, - ) - return self._get_dataloader( - dataset=dataset, - description="Training", - batch_size=self._train_batch_size * self.args.steps_per_generation, # < this is the change - sampler_fn=self._get_train_sampler, - is_training=True, - ) - - def _get_train_sampler(self, dataset: Dataset | None = None) -> Sampler: - # Returns a sampler that - # 1. ensures each prompt is repeated across multiple processes. This guarantees that identical prompts are - # distributed to different GPUs, allowing rewards to be computed and normalized correctly within each prompt - # group. Using the same seed across processes ensures consistent prompt assignment, preventing discrepancies - # in group formation. - # 2. repeats the batch multiple times to allow reusing generations across multiple updates. Refer to - # _prepare_inputs to see how the generations are stored and reused. - - # In the following figure, the values are the prompt indices. Each row shows the per-step batch - # returned by `_prepare_inputs`; rows within a `steps_per_generation` block are slices of the same - # generated batch. When `num_iterations > 1`, that block is reused for multiple optimization passes - # before regenerating. - # - # | GPU 0 | GPU 1 | - # - # global_step step <-───> num_generations=2 - # <-───────> per_device_train_batch_size=3 - # grad_accum ▲ ▲ 0 0 0 0 1 1 2 2 <- Generate for the first `steps_per_generation` (prompts 0 to 11); store the completions; use the first slice to compute the loss - # =2 ▼ | 0 1 3 3 4 4 5 5 <- Take the stored generations and use the second slice to compute the loss - # | - # | 1 2 6 6 7 7 8 8 <- Take the stored generations and use the third slice to compute the loss - # steps_per_gen=4 ▼ 1 3 9 9 10 10 11 11 <- Take the stored generations and use the fourth slice to compute the loss - # - # 2 4 12 12 13 13 14 14 <- Generate for the second `steps_per_generation` (prompts 12 to 23); store the completions; use the first slice to compute the loss - # 2 5 15 15 16 16 17 17 <- Take the stored generations and use the second slice to compute the loss - # ... - if dataset is None: - dataset = self.train_dataset - return RepeatSampler( - data_source=dataset, - mini_repeat_count=self.num_generations, - batch_size=self.args.generation_batch_size // self.num_generations, - repeat_count=self.num_iterations * self.args.steps_per_generation, - shuffle=self.shuffle_dataset, - seed=self.args.seed, - ) - - def _get_eval_sampler(self, eval_dataset) -> Sampler: - # See _get_train_sampler for an explanation of the sampler. - return RepeatSampler( - data_source=eval_dataset, - mini_repeat_count=self.num_generations_eval, - seed=self.args.seed, - ) - - # This method overrides `Trainer.get_eval_dataloader` to wrap iterable eval datasets, reproducing the - # RepeatSampler ordering that can't be attached to them (see `get_train_dataloader`). Map-style datasets keep the - # default path via `_get_eval_sampler`, which shuffles with `seed`, so the iterable wrap shuffles too (buffered) - # to walk prompts in a matching order. - # Maintenance note: this method is a copy-paste of the original `Trainer.get_eval_dataloader`, with the iterable - # wrapping as the only addition. - def get_eval_dataloader(self, eval_dataset: str | Dataset | IterableDataset | None = None) -> DataLoader: - if eval_dataset is None and self.eval_dataset is None: - raise ValueError("Trainer: evaluation requires an eval_dataset.") - - # If we have persistent workers, don't do a fork bomb especially as eval datasets - # don't change during training - dataloader_key = eval_dataset if isinstance(eval_dataset, str) else "eval" - if ( - hasattr(self, "_eval_dataloaders") - and dataloader_key in self._eval_dataloaders - and self.args.dataloader_persistent_workers - ): - return self._eval_dataloaders[dataloader_key] - - eval_dataset = ( - self.eval_dataset[eval_dataset] - if isinstance(eval_dataset, str) - else eval_dataset - if eval_dataset is not None - else self.eval_dataset - ) - - if isinstance(eval_dataset, IterableDataset): - # Apply the `__init__` iterable config here too - if self.args.accelerator_config.dispatch_batches: - raise ValueError( - "Iterable datasets require `dispatch_batches=False`, but it is set to `True` in " - "`accelerator_config`. Please set it to `False`." - ) - self.accelerator.dataloader_config.dispatch_batches = False - eval_dataset = eval_dataset.shuffle(seed=self.args.seed) - eval_dataset = repeat_iterable_dataset(eval_dataset, mini_repeat_count=self.num_generations_eval) - # Force a single worker for this loader only, without persisting the change - num_workers = self.args.dataloader_num_workers - self.args.dataloader_num_workers = 0 - - try: - return self._get_dataloader( - dataset=eval_dataset, - description="Evaluation", - batch_size=self.args.eval_batch_size, - sampler_fn=self._get_eval_sampler, - dataloader_key=dataloader_key, - ) - finally: - if isinstance(eval_dataset, IterableDataset): - self.args.dataloader_num_workers = num_workers - - @profiling_decorator - def _get_per_token_logps_and_entropies( - self, - model, - input_ids, - attention_mask, - logits_to_keep, - batch_size=None, - compute_entropy=False, - compute_aux_loss=False, - pixel_values=None, - image_grid_thw=None, - num_images=None, - pixel_attention_mask=None, - spatial_shapes=None, - num_tiles=None, - image_sizes=None, - token_type_ids=None, - mm_token_type_ids=None, - image_position_ids=None, - ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: - """Compute log-probs, (optionally) entropies, and (optionally) the MoE load-balancing aux loss.""" - batch_size = batch_size or input_ids.size(0) # Chunk inputs into smaller batches to reduce memory peak - all_logps = [] - all_entropies = [] - all_aux_losses = [] - for start in range(0, input_ids.size(0), batch_size): - end = min(start + batch_size, input_ids.size(0)) # the last chunk can be smaller than batch_size - input_ids_batch = input_ids[start:end] - attention_mask_batch = attention_mask[start:end] - - # Build model inputs - model_inputs = {"input_ids": input_ids_batch, "attention_mask": attention_mask_batch} - if num_images is not None: - cum_imgs = torch.tensor([0] + num_images).cumsum(0) - img_start, img_end = cum_imgs[start], cum_imgs[end] - if image_grid_thw is not None and pixel_values is not None: - rows_per_image = image_grid_thw.prod(dim=-1) - rows_per_sample = torch.split(rows_per_image, num_images) - rows_per_sample = torch.stack([s.sum() for s in rows_per_sample]) - cum_rows = torch.cat([torch.tensor([0], device=rows_per_sample.device), rows_per_sample.cumsum(0)]) - row_start, row_end = cum_rows[start].item(), cum_rows[end].item() - model_inputs["pixel_values"] = pixel_values[row_start:row_end] - model_inputs["image_grid_thw"] = image_grid_thw[img_start:img_end] - elif image_position_ids is not None and pixel_values is not None: - model_inputs["pixel_values"] = pixel_values[img_start:img_end] - model_inputs["image_position_ids"] = image_position_ids[img_start:img_end] - elif spatial_shapes is not None and pixel_values is not None: - # LFM2-VL tensors are tile-indexed. - cum_tiles = torch.tensor([0] + num_tiles).cumsum(0) - tile_start, tile_end = cum_tiles[start], cum_tiles[end] - model_inputs["pixel_values"] = pixel_values[tile_start:tile_end] - model_inputs["pixel_attention_mask"] = pixel_attention_mask[tile_start:tile_end] - model_inputs["spatial_shapes"] = spatial_shapes[tile_start:tile_end] - elif num_tiles is not None and pixel_values is not None: - # InternVL tensors are tile-indexed. - cum_tiles = torch.tensor([0] + num_tiles).cumsum(0) - tile_start, tile_end = cum_tiles[start], cum_tiles[end] - model_inputs["pixel_values"] = pixel_values[tile_start:tile_end] - elif pixel_values is not None and pixel_values.size(0) == sum(num_images): - model_inputs["pixel_values"] = pixel_values[img_start:img_end] - elif pixel_values is not None: - model_inputs["pixel_values"] = pixel_values[start:end] - if pixel_attention_mask is not None and spatial_shapes is None: - model_inputs["pixel_attention_mask"] = pixel_attention_mask[start:end] - if image_sizes is not None: - model_inputs["image_sizes"] = image_sizes[img_start:img_end] - if token_type_ids is not None: - model_inputs["token_type_ids"] = token_type_ids[start:end] - if mm_token_type_ids is not None: - model_inputs["mm_token_type_ids"] = mm_token_type_ids[start:end] - - # MoE models: request router logits so the model returns the load-balancing loss - if compute_aux_loss: - model_inputs["output_router_logits"] = True - - # Only the completion tokens are scored - labels = input_ids_batch.masked_fill(attention_mask_batch == 0, -100) - labels[:, :-logits_to_keep] = -100 - with self.accelerator.autocast(): - outputs = model(**model_inputs, labels=labels, fused_lm_head=True) - all_logps.append(outputs.log_probs[:, -logits_to_keep:]) - if compute_entropy: - all_entropies.append(outputs.entropy[:, -logits_to_keep:]) - if compute_aux_loss: - all_aux_losses.append(outputs.aux_loss) - - logps = torch.cat(all_logps, dim=0) - entropies = torch.cat(all_entropies, dim=0) if compute_entropy else None - aux_loss = torch.stack(all_aux_losses).mean() if compute_aux_loss else None - return logps, entropies, aux_loss - - def training_step(self, model, inputs, num_items_in_batch): - time_before = time.perf_counter() - output = super().training_step(model, inputs, num_items_in_batch) - self._step += 1 - time_after = time.perf_counter() - self._current_train_step_time += time_after - time_before - if self._step % self.current_gradient_accumulation_steps == 0: - self._metrics["train"]["step_time"].append(self._current_train_step_time) - self._current_train_step_time = 0.0 - return output - - @profiling_decorator - def _prepare_inputs(self, generation_batch: dict[str, torch.Tensor | Any]) -> dict[str, torch.Tensor | Any]: - # Prepares inputs for model training/evaluation by managing completion generation and batch handling. - # During training: - # - Receives the local generation batch (Per-GPU batch size × steps per generation) - # from the modified training dataloader instead of the standard local batch - # - Generates completions once for the entire generation batch and splits it into batches of size - # `per_device_train_batch_size` - # - Buffers these completions and returns the appropriate slice for the current accumulation step - # - Optimizes by regenerating completions only periodically (every steps_per_generation * num_iterations) - # During evaluation: - # - The input is treated as a standard local batch (no accumulation, no multiple iterations) - # - Completions are generated for each batch without buffering or reuse - # Returns a single local batch in both cases. - - mode = "train" if self.model.training else "eval" - if mode == "train": - generate_every = self.args.steps_per_generation * self.num_iterations - if self._step % generate_every == 0 or self._buffered_inputs is None: - # self._buffered_inputs=None can occur when resuming from a checkpoint - generation_batch = self._generate_and_score_completions(generation_batch) - generation_batch = split_pixel_values_by_grid(generation_batch) - - try: generation_batch = shuffle_sequence_dict(generation_batch) - - except: pass - generation_batches = split_tensor_dict(generation_batch, self.args.steps_per_generation) - self._buffered_inputs = [unsplit_pixel_values_by_grid(batch) for batch in generation_batches] - inputs = self._buffered_inputs[self._step % self.args.steps_per_generation] - else: - # In evaluation, there is neither batch grouping for generation, nor multiple iterations, hence - # local generation batch == local eval batch - inputs = self._generate_and_score_completions(generation_batch) - return inputs - - def _log_completion_extra(self, column: str, values: list): - """ - Log extra columns to the completions table. Called from reward functions via the `log_extra` kwarg. - - Args: - column (`str`): - Name of the column to add. - values (`list`): - Values for the column, one per sample in the batch. - """ - self._pending_extra_logs[column].extend(values) - - def _log_metric(self, name: str, value: float): - """ - Log a scalar metric from a reward function. Called via the `log_metric` kwarg. Values are averaged over each - logging step and reported alongside built-in metrics like `kl` and `entropy`. - - Args: - name (`str`): - Name of the metric. - value (`float`): - Scalar value for this batch. - """ - self._pending_metrics[name].append(value) - - @profiling_decorator - def _calculate_rewards(self, inputs, prompts, completions, completion_ids_list): - device = self.accelerator.device - rewards_per_func = torch.zeros(len(prompts), len(self.reward_funcs), device=device) - - # Repeat all input columns (but "prompt", "completion", and "completion_ids") to match the num of generations - keys = [key for key in inputs[0] if key not in ["prompt", "completion", "completion_ids"]] - reward_kwargs = {key: [example[key] for example in inputs] for key in keys} - - # This allows for dynamic reward shaping based on training progress. - reward_kwargs["trainer_state"] = self.state - - # Allow reward functions to log extra columns to the completions table. - reward_kwargs["log_extra"] = self._log_completion_extra - - # Allow reward functions to log additional scalar metrics. - reward_kwargs["log_metric"] = self._log_metric - - async_funcs_info = [] # async custom functions for asyncio.gather - - for i, (reward_func, reward_processing_class, reward_func_name) in enumerate( - zip(self.reward_funcs, self.reward_processing_classes, self.reward_func_names, strict=True) - ): - if isinstance(reward_func, nn.Module): # Module (no PretrainedModel) for compat with compiled models - with profiling_context(self, reward_func_name): - if is_conversational({"prompt": prompts[0]}): - messages = [{"messages": p + c} for p, c in zip(prompts, completions, strict=True)] - texts = [ - apply_chat_template(x, reward_processing_class, **self.chat_template_kwargs)["text"] - for x in messages - ] - else: - texts = [p + c for p, c in zip(prompts, completions, strict=True)] - reward_inputs = reward_processing_class( - text=texts, return_tensors="pt", padding=True, padding_side="right", add_special_tokens=False - ) - reward_inputs = super()._prepare_inputs(reward_inputs) - with torch.inference_mode(): - rewards_per_func[:, i] = reward_func(**reward_inputs).logits[:, 0] # Shape (B*G,) - elif is_async_callable(reward_func): # Separate async reward funcs to run them in parallel later - async_funcs_info.append((i, reward_func, reward_func_name)) - else: - # Run synchronous reward function - with profiling_context(self, reward_func_name): - output_reward_func = reward_func( - prompts=prompts, completions=completions, completion_ids=completion_ids_list, **reward_kwargs - ) - # Convert None values to NaN - output_reward_func = [reward if reward is not None else torch.nan for reward in output_reward_func] - if len(output_reward_func) != len(prompts): - raise ValueError( - f"The reward function '{reward_func_name}' returned {len(output_reward_func)} rewards, but " - f"{len(prompts)} were expected (one per prompt-completion pair). Make sure the reward " - f"function returns exactly one reward per completion." - ) - rewards_per_func[:, i] = torch.tensor(output_reward_func, dtype=torch.float32, device=device) - - # Execute async custom functions in parallel using asyncio.gather - if async_funcs_info: - - async def _invoke_async(index, func, func_name): - with profiling_context(self, func_name): - output = await func( - prompts=prompts, completions=completions, completion_ids=completion_ids_list, **reward_kwargs - ) - output = [r if r is not None else torch.nan for r in output] - if len(output) != len(prompts): - raise ValueError( - f"The reward function '{func_name}' returned {len(output)} rewards, but {len(prompts)} " - f"were expected (one per prompt-completion pair). Make sure the reward function returns " - f"exactly one reward per completion." - ) - return index, output - - async def _run_async_funcs(): - coros = [_invoke_async(i, func, func_name) for (i, func, func_name) in async_funcs_info] - return await asyncio.gather(*coros) - - async_results = asyncio.run_coroutine_threadsafe(_run_async_funcs(), self.async_loop).result() - for idx, output_reward_func in async_results: - rewards_per_func[:, idx] = torch.tensor(output_reward_func, dtype=torch.float32, device=device) - - # If all reward functions return None for a given row, issue a detailed warning - if torch.isnan(rewards_per_func).all(dim=1).any(): - nan_row_idx = torch.isnan(rewards_per_func).all(dim=1).nonzero(as_tuple=True)[0][0] - row_reward_kwargs = { - key: value[nan_row_idx] - for key, value in reward_kwargs.items() - if key not in ("trainer_state", "log_extra", "log_metric") - } - row_reward_kwargs["prompt"] = prompts[nan_row_idx] - row_reward_kwargs["completion"] = completions[nan_row_idx] - logger.warning( - f"All reward functions returned None for the following kwargs:\n{row_reward_kwargs}\n" - "Please ensure that at least one reward function returns a valid reward." - ) - - # Gather the reward per function: this part is crucial, because the rewards are normalized per group and the - # completions may be distributed across processes - rewards_per_func = gather(rewards_per_func) - return rewards_per_func - - def _tokenize_prompts(self, prompts: list): - """Tokenize prompts and extract images/multimodal fields for generation.""" - if is_conversational({"prompt": prompts[0]}): - # Extract images from messages for VLM support - images = [] - has_images = False - for prompt in prompts: - prompt_images = [] - for message in prompt: - if isinstance(message["content"], list): - for part in message["content"]: - if part["type"] == "image": - prompt_images.append(part["image"]) - has_images = True - images.append(prompt_images if prompt_images else None) - images = images if has_images else None - - # Workaround for a bug in transformers 5.3.0 where some processors (e.g. Qwen2.5-VL) crash on - # batched unpadded input (transformers#44514). - # Fixed in transformers 5.4.0 (transformers#44563). - needs_padding_workaround = Version("5.3.0") <= Version(transformers.__version__) < Version("5.4.0") - tokenized = self.processing_class.apply_chat_template( - conversation=prompts, - add_generation_prompt=True, - tokenize=True, - return_dict=True, - **({"padding": True} if needs_padding_workaround else {}), - **self.chat_template_kwargs, - ) - if needs_padding_workaround: - # Unpad input_ids: remove padding tokens using attention_mask to get per-sequence lists - prompt_ids = [ - [tok for tok, m in zip(ids, mask, strict=True) if m] - for ids, mask in zip(tokenized["input_ids"], tokenized["attention_mask"], strict=True) - ] - else: - prompt_ids = tokenized["input_ids"] - # For VLMs, the processor returns extra multimodal fields (pixel_values, image_grid_thw, etc.) - multimodal_fields = {k: v for k, v in tokenized.items() if k not in ("input_ids", "attention_mask")} - else: - prompt_ids = self.processing_class(text=prompts)["input_ids"] - images = None - multimodal_fields = {} - return prompt_ids, images, multimodal_fields - - def _generate_single_turn(self, prompt_ids, images, multimodal_fields): - device = self.accelerator.device - mode = "train" if self.model.training else "eval" - - # Generate completions using either vLLM or regular generation - if self.use_vllm: - # Sync weights if training step changed - if self.state.global_step != self._last_loaded_step: - with profiling_context(self, "sync_weights"): - self.vllm_generation.sync_weights() - self._last_loaded_step = self.state.global_step - - # Generate using vLLM (note: RLOO doesn't use logprobs from generation, so we ignore them) - num_generations = self.num_generations if mode == "train" else self.num_generations_eval - _, completion_ids, _, _ = self.vllm_generation.generate( - prompts=prompt_ids, - images=images, - num_generations=num_generations, - profiler=profiling_context(self, "vLLM.generate"), - ) - - elif self.use_transformers_continuous_batching: - with ( - profiling_context(self, "transformers.generate_batch"), - unwrap_model_for_generation( - self.model_wrapped, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation - ) as unwrapped_model, - torch.no_grad(), - self._dist.summon_full_params(self.model_wrapped, recurse=False), - ): - # Cast to the appropriate dtype based on training configuration - if self.args.bf16: - unwrapped_model.to(torch.bfloat16) - elif self.args.fp16: - unwrapped_model.to(torch.float16) - all_outputs = unwrapped_model.generate_batch( - prompt_ids, - generation_config=self.generation_config, - continuous_batching_config=self.continuous_batching_config, - progress_bar=False, - ) - unwrapped_model.train() - completion_ids = [output.generated_tokens for output in all_outputs.values()] - - else: - # Regular generation path: left-pad token IDs into tensors - prompt_tensors = [torch.tensor(ids) for ids in prompt_ids] - padded_ids = pad(prompt_tensors, padding_value=self._tokenizer.pad_token_id, padding_side="left") - attention_mask = pad([torch.ones_like(t) for t in prompt_tensors], padding_value=0, padding_side="left") - generate_inputs = {"input_ids": padded_ids, "attention_mask": attention_mask} - # For VLMs, include multimodal fields as tensors (pixel_values, image_grid_thw, etc.) - for k, v in multimodal_fields.items(): - if isinstance(v, torch.Tensor): - generate_inputs[k] = v - elif isinstance(v, list) and v and isinstance(v[0], list): - # Per-token field (e.g., token_type_ids): left-pad like input_ids - generate_inputs[k] = pad([torch.tensor(x) for x in v], padding_value=0, padding_side="left") - else: - generate_inputs[k] = torch.tensor(np.array(v)) - generate_inputs = super()._prepare_inputs(generate_inputs) - - with ( - profiling_context(self, "transformers.generate"), - unwrap_model_for_generation( - self.model_wrapped, - self.accelerator, - gather_deepspeed3_params=self.args.ds3_gather_for_generation, - generation_kwargs=self.generation_kwargs, # Override model.generation_config with generation_kwargs to fix transformers#42762 - ) as unwrapped_model, - torch.no_grad(), - self._dist.summon_full_params(self.model_wrapped, recurse=False), - ): - prompt_completion_ids = unwrapped_model.generate( - **generate_inputs, generation_config=self.generation_config - ) - # Compute prompt length and extract completion ids - prompt_length = generate_inputs["input_ids"].size(1) - completion_ids = prompt_completion_ids[:, prompt_length:] - - # Mask everything after the first EOS token - is_eos = torch.isin(completion_ids, torch.tensor(self.eos_token_ids, device=device)) - eos_idx = torch.full((is_eos.size(0),), is_eos.size(1), dtype=torch.long, device=device) - eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)] - sequence_indices = torch.arange(is_eos.size(1), device=device).expand(is_eos.size(0), -1) - completion_mask = (sequence_indices <= eos_idx.unsqueeze(1)).int() - completion_ids = [ - c[m].tolist() for c, m in zip(completion_ids.cpu(), completion_mask.bool().cpu(), strict=True) - ] - - return completion_ids - - def _generate(self, prompts: list): - device = self.accelerator.device - mode = "train" if self.model.training else "eval" - - # Copy the prompts to avoid modifying the original list - prompts = copy.deepcopy(prompts) - - prompt_ids, images, multimodal_fields = self._tokenize_prompts(prompts) - completion_ids = self._generate_single_turn(prompt_ids, images, multimodal_fields) - - # Decode completions. It's important to use `parse_response` when possible, because it handles tool calls. - if is_conversational({"prompt": prompts[0]}): - contents = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) - completions = [[{"role": "assistant", "content": content}] for content in contents] - else: - completions = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) - - # Get completion length per sequence, used for logging - prompt_lengths = torch.tensor([len(ids) for ids in prompt_ids], device=device) - completion_lengths = torch.tensor([len(ids) for ids in completion_ids], device=device) - agg_prompt_lengths = self.accelerator.gather(prompt_lengths) - agg_completion_lengths = self.accelerator.gather(completion_lengths) - # Fail clearly if the generation backend returned no completions (avoids a cryptic min() error below). - if agg_completion_lengths.numel() == 0: - raise RuntimeError( - "No completions were generated. This usually means the generation backend failed to return any " - "results; see the generation logs above for the underlying error." - ) - total_prompt_tokens = agg_prompt_lengths.sum() - - # Log the metrics - if mode == "train": - self.state.num_input_tokens_seen += (total_prompt_tokens + agg_completion_lengths.sum()).item() - self._metrics[mode]["num_tokens"] = [self.state.num_input_tokens_seen] - - # Log completion lengths, mean, min, max - self._metrics[mode]["completions/mean_length"].append(agg_completion_lengths.float().mean().item()) - self._metrics[mode]["completions/min_length"].append(agg_completion_lengths.float().min().item()) - self._metrics[mode]["completions/max_length"].append(agg_completion_lengths.float().max().item()) - - # Identify sequences that terminated with EOS and log their lengths - eos_and_pad = [*self.eos_token_ids, self._tokenizer.pad_token_id] - is_truncated = torch.tensor([ids[-1] not in eos_and_pad for ids in completion_ids], device=device) - agg_is_truncated = self.accelerator.gather(is_truncated) - self._metrics[mode]["completions/clipped_ratio"].append(agg_is_truncated.float().mean().item()) - term_completion_lengths = agg_completion_lengths[~agg_is_truncated] - if len(term_completion_lengths) == 0: # edge case where no terminated sequences are found - term_completion_lengths = torch.zeros(1, device=device) - self._metrics[mode]["completions/mean_terminated_length"].append(term_completion_lengths.float().mean().item()) - self._metrics[mode]["completions/min_terminated_length"].append(term_completion_lengths.float().min().item()) - self._metrics[mode]["completions/max_terminated_length"].append(term_completion_lengths.float().max().item()) - - return prompt_ids, completion_ids, completions - - def _generate_and_score_completions( - self, inputs: list[dict[str, torch.Tensor | Any]] - ) -> dict[str, torch.Tensor | Any]: - device = self.accelerator.device - mode = "train" if self.model.training else "eval" - - prompts = [x["prompt"] for x in inputs] - - if "images" in inputs[0]: - images = [example.get("images") for example in inputs] - elif "image" in inputs[0]: - images = [[example.get("image")] if example.get("image") is not None else None for example in inputs] - else: - images = None - # Transformers requires at least one image in the batch, otherwise it throws an error - if images is not None and all(img_list == [] for img_list in images): - images = None - - # If the prompts are conversational and the inputs contain images, we need to convert the prompts from - # [{"role": "user", "content": "What color is the sky?"}] to - # [{"role": "user", "content": [{"type": "image", "image": }, {"type": "text", "text": "What color is the sky?"}]}] - if images is not None: - if not is_conversational(inputs[0]): - raise ValueError( - "Multimodal training requires conversational prompts. It looks like the dataset contains " - "non-conversational inputs, likely because a chat template was applied before passing the dataset " - "to the trainer. Please provide the raw conversational prompts and let the trainer apply the chat " - "template internally." - ) - prompts = [ - prepare_multimodal_messages(prompt, images=image_list) - for prompt, image_list in zip(prompts, images, strict=True) - ] - - prompt_ids_list, completion_ids_list, completions = self._generate(prompts) - - # Convert lists of token IDs to padded tensors - prompt_ids = [torch.tensor(ids) for ids in prompt_ids_list] - prompt_mask = [torch.ones_like(ids, dtype=torch.long) for ids in prompt_ids] - prompt_ids = pad( - prompt_ids, - padding_value=self._tokenizer.pad_token_id, - padding_side="left", - pad_to_multiple_of=self.pad_to_multiple_of, - ).to(device=device) - prompt_mask = pad( - prompt_mask, padding_value=0, padding_side="left", pad_to_multiple_of=self.pad_to_multiple_of - ).to(device=device) - completion_ids = [torch.tensor(ids) for ids in completion_ids_list] - completion_mask = [torch.ones_like(ids, dtype=torch.long) for ids in completion_ids] - completion_ids = pad( - completion_ids, - padding_value=self._tokenizer.pad_token_id, - padding_side="right", - pad_to_multiple_of=self.pad_to_multiple_of, - ).to(device=device) - completion_mask = pad( - completion_mask, padding_value=0, padding_side="right", pad_to_multiple_of=self.pad_to_multiple_of - ).to(device=device) - - # If mask_truncated_completions is enabled, zero out truncated completions in completion_mask - if self.mask_truncated_completions: - eos_and_pad = [*self.eos_token_ids, self._tokenizer.pad_token_id] - is_truncated = torch.tensor([ids[-1] not in eos_and_pad for ids in completion_ids_list], device=device) - # Mask completion_mask for attention masking - completion_mask = completion_mask * (~is_truncated).unsqueeze(1).int() - - # Concatenate prompt_mask with completion_mask for logit computation - prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1) # (B, P+C) - attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) # (B, P+C) - - logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens - batch_size = self.args.per_device_train_batch_size if mode == "train" else self.args.per_device_eval_batch_size - - num_images = [len(img_list) if img_list else 0 for img_list in images] if images is not None else None - - # Get forward_kwargs for models with multimodal inputs. - if images is not None: - prompts_text = [ - apply_chat_template({"prompt": prompt}, self.processing_class, **self.chat_template_kwargs)["prompt"] - for prompt in prompts - ] - prompt_inputs = self.processing_class(images=images, text=prompts_text, padding=True, return_tensors="pt") - prompt_inputs = super()._prepare_inputs(prompt_inputs) - forward_kwargs = {k: v for k, v in prompt_inputs.items() if k not in ["input_ids", "attention_mask"]} - else: - forward_kwargs = {} - - # Recover LFM2-VL tile counts; the full processor drops row/column metadata. - num_tiles = None - if images is not None and "spatial_shapes" in forward_kwargs: - image_info = self.processing_class.image_processor( - images=images, return_tensors="pt", return_row_col_info=True - ) - tiles_per_image = image_info["image_rows"] * image_info["image_cols"] - if self.processing_class.image_processor.use_thumbnail: - tiles_per_image = tiles_per_image + (tiles_per_image > 1).to(tiles_per_image.dtype) - num_tiles = [group.sum().item() for group in torch.split(tiles_per_image, num_images)] - # Same for InternVL, whose pixel_values is tile-indexed ([total_tiles, channels, height, width]). - elif ( - images is not None - and forward_kwargs["pixel_values"].ndim == 4 - and forward_kwargs["pixel_values"].size(0) != sum(num_images) - ): - num_patches = self.processing_class.image_processor( - images=images, crop_to_patches=True, return_tensors="pt" - )["num_patches"] - num_tiles = [group.sum().item() for group in torch.split(num_patches, num_images)] - - # If token_type_ids are used, extend them with zeros for the completion part - if "token_type_ids" in forward_kwargs: - token_type_ids = forward_kwargs["token_type_ids"] - if self.pad_to_multiple_of is not None: - # Needed only with pad_to_multiple_of: otherwise prompt_ids and token_type_ids must have equal len - padding_size = prompt_ids.size(1) - token_type_ids.size(1) - if padding_size > 0: - token_type_ids = torch.cat( - [token_type_ids.new_zeros((token_type_ids.size(0), padding_size)), token_type_ids], dim=1 - ) - forward_kwargs["token_type_ids"] = torch.cat( - [token_type_ids, token_type_ids.new_zeros(completion_ids.shape)], dim=1 - ) - # If mm_token_type_ids are used, extend them with zeros for the completion part - if "mm_token_type_ids" in forward_kwargs: - mm_token_type_ids = forward_kwargs["mm_token_type_ids"] - if self.pad_to_multiple_of is not None: - # Needed only with pad_to_multiple_of: otherwise prompt_ids and mm_token_type_ids must have equal len - padding_size = prompt_ids.size(1) - mm_token_type_ids.size(1) - if padding_size > 0: - mm_token_type_ids = torch.cat( - [mm_token_type_ids.new_zeros((mm_token_type_ids.size(0), padding_size)), mm_token_type_ids], - dim=1, - ) - forward_kwargs["mm_token_type_ids"] = torch.cat( - [mm_token_type_ids, mm_token_type_ids.new_zeros(completion_ids.shape)], dim=1 - ) - - # When gradient checkpointing is enabled with use_reentrant=True (non default), calling the model inside a - # torch.no_grad() block triggers a harmless PyTorch warning ("None of the inputs have requires_grad=True"). - # Temporarily disable checkpointing to avoid this warning during inference. - with torch.no_grad(), disable_gradient_checkpointing(self.model, self.args.gradient_checkpointing_kwargs): - # Compute the per-token log probabilities for the current model - old_per_token_logps, _, _ = self._get_per_token_logps_and_entropies( - self.model, - prompt_completion_ids, - attention_mask, - logits_to_keep, - batch_size, - num_images=num_images, - num_tiles=num_tiles, - **forward_kwargs, # may contain pixel_values, image_grid_thw, pixel_attention_mask, spatial_shapes, image_sizes, image_position_ids - ) - old_logps = (old_per_token_logps * completion_mask).sum(1) # mask out padding and tokens after EOS - - # Compute the per-token log probabilities for the reference model - if self.beta != 0.0: - if self.ref_model is not None: - ref_per_token_logps, _, _ = self._get_per_token_logps_and_entropies( - self.ref_model, - prompt_completion_ids, - attention_mask, - logits_to_keep, - batch_size=batch_size, - num_images=num_images, - num_tiles=num_tiles, - **forward_kwargs, # may contain pixel_values, image_grid_thw, pixel_attention_mask, spatial_shapes, image_sizes, image_position_ids - ) - else: - # When training a PEFT adapter, how we obtain the reference depends on the setup: - # - New adapter: disabling adapters yields the base model. - # - Re-training an existing adapter: an initial copy is loaded under the name "ref". - model = self.accelerator.unwrap_model(self.model) - with use_adapter(model, adapter_name="ref" if "ref" in model.peft_config else None): - ref_per_token_logps, _, _ = self._get_per_token_logps_and_entropies( - self.model, - prompt_completion_ids, - attention_mask, - logits_to_keep, - batch_size=batch_size, - num_images=num_images, - num_tiles=num_tiles, - **forward_kwargs, # may contain pixel_values, image_grid_thw, pixel_attention_mask, spatial_shapes, image_sizes, image_position_ids - ) - else: - ref_per_token_logps = None - - # Calculate rewards for each reward function. rewards_per_func aggregates rewards across all processes. This is - # important because rewards will be normalized per group, and completions are distributed. We will later slice - # rewards_per_func to extract each process's subset. - rewards_per_func = self._calculate_rewards(inputs, prompts, completions, completion_ids_list) - num_generations = self.num_generations if mode == "train" else self.num_generations_eval - - # A completion for which every reward function returned None is unscorable. nansum would collapse it to 0, - # which both biases the leave-one-out baseline and hands the completion a spurious advantage. Mark these rows - # NaN so they're excluded from the (nan-aware) baseline below; their advantage is forced to 0 afterwards. - unscorable_mask = torch.isnan(rewards_per_func).all(dim=1) - - # Apply weights to each reward function's output and sum - rewards = (rewards_per_func * self.reward_weights.to(device).unsqueeze(0)).nansum(dim=1) - rewards[unscorable_mask] = torch.nan - - # Apply reward clipping if specified - if self.reward_clip_range: - rewards = rewards.clamp(min=self.reward_clip_range[0], max=self.reward_clip_range[1]) - - # Include the KL penalty in the reward - if self.beta != 0.0: - # RLOO uses the first-order log ratio for the per-token KL estimate, following the original RLOO paper - # (Ahmadian et al., 2024, https://huggingface.co/papers/2405.14782). Unlike GRPOTrainer's Schulman - # approximation (always >= 0), this can be negative per token. The divergence is intentional: RLOO applies - # KL as a reward penalty (summed across tokens per sequence), while GRPO adds it to the per-token loss. - per_token_kl = old_per_token_logps - ref_per_token_logps - # Apply sequence-level KL penalty to rewards (sum KL across tokens first, then apply to each sequence) - kl = (per_token_kl * completion_mask).sum(-1) - kl = gather(kl) # rewards are gathered, so kl must be too - rewards = rewards - self.beta * kl - - grouped_rewards = rewards.view(-1, num_generations) - mean_grouped_rewards = torch.nanmean(grouped_rewards, dim=1) - if num_generations > 1: - std_rewards = nanstd(grouped_rewards, dim=1) - else: # doesn't occur during training, but could occur in eval when num_generations_eval=1 - std_rewards = torch.zeros_like(mean_grouped_rewards) - - # RLOO advantages computation. The leave-one-out baseline averages over scorable siblings only: nansum drops - # unscorable rewards and the divisor is (scorable count − 1). A group with a single scorable completion yields - # 0/0 = NaN, and unscorable rows stay NaN; both are zeroed by nan_to_num below. - scorable_counts = (~torch.isnan(grouped_rewards)).sum(dim=1, keepdim=True) # (num_prompts, 1) - grouped_sum = torch.nansum(grouped_rewards, dim=1, keepdim=True) # (num_prompts, 1) - if num_generations > 1: - baselines = (grouped_sum - grouped_rewards) / (scorable_counts - 1) # (num_prompts, num_generations) - baselines = baselines.view(-1) # Flatten back to match rewards shape - advantages = rewards - baselines - else: # this case doesn't occur during training, but could in eval when num_generations_eval=1 - advantages = torch.zeros_like(rewards) - - # Normalize advantages over the scorable subset only (unscorable advantages are still NaN here). - if self.normalize_advantages: - advantages = (advantages - torch.nanmean(advantages)) / (nanstd(advantages) + 1e-4) - - # Unscorable completions carry no learning signal: zero their advantage to keep them from moving the policy. - advantages = torch.nan_to_num(advantages, nan=0.0) - - is_std_zero = torch.isclose(std_rewards, torch.zeros_like(std_rewards)) # for logging - - # Slice to keep only the local part of the data - process_slice = slice( - self.accelerator.process_index * len(prompts), - (self.accelerator.process_index + 1) * len(prompts), - ) - all_process_advantages = advantages.clone() # keep the aggregated advantages for logging - advantages = advantages[process_slice] - - # Calculate and log the mean KL divergence between current and reference model - if self.beta != 0.0: - kl_stats = self.accelerator.reduce( - torch.stack([(per_token_kl * completion_mask).sum(), completion_mask.sum().float()]), reduction="sum" - ) - self._metrics[mode]["kl"].append((kl_stats[0] / kl_stats[1].clamp(min=1.0)).item()) - - # Calculate mean reward per function, but only for samples where the function was applied (non-NaN values) - for i, reward_func_name in enumerate(self.reward_func_names): - mean_rewards = torch.nanmean(rewards_per_func[:, i]).item() - self._metrics[mode][f"rewards/{reward_func_name}/mean"].append(mean_rewards) - std_func_rewards = nanstd(rewards_per_func[:, i]).item() - self._metrics[mode][f"rewards/{reward_func_name}/std"].append(std_func_rewards) - rewards = (rewards_per_func * self.reward_weights.to(rewards_per_func.device).unsqueeze(0)).nansum(dim=1) - rewards[unscorable_mask] = torch.nan # exclude unscorable rows from the logged reward stats - self._metrics[mode]["reward"].append(torch.nanmean(rewards).item()) - self._metrics[mode]["reward_std"].append(nanstd(rewards).item()) - self._metrics[mode]["frac_reward_zero_std"].append(is_std_zero.float().mean().item()) - - # Log prompts and completions - self._logs["prompt"].extend(gather_object(prompts)) - self._logs["completion"].extend(gather_object(completions)) - for i, name in enumerate(self.reward_func_names): - self._logs["rewards"][name].extend(rewards_per_func[:, i].tolist()) - self._logs["advantages"].extend(all_process_advantages.tolist()) - - # Flush user-logged extra columns (from log_extra), gathering across processes. - # Keys must be sorted so that all ranks call gather_object in the same order, otherwise values - # get mis-attributed across columns (dict insertion order may differ between processes). - for column in sorted(set(gather_object(list(self._pending_extra_logs)))): - values = self._pending_extra_logs.get(column, [None] * len(prompts)) - self._logs["extra"][column].extend(gather_object(values)) - self._pending_extra_logs.clear() - - # Flush user-logged metrics (from log_metric), averaging across processes. - # Keys must be sorted so that all ranks call accelerator.gather in the same order, otherwise values - # get mis-attributed across metrics (dict insertion order may differ between processes). - for name in sorted(set(gather_object(list(self._pending_metrics)))): - values = self._pending_metrics.get(name, []) - local_stats = torch.tensor([sum(values), len(values)], dtype=torch.float32, device=device) - total, count = self.accelerator.gather(local_stats).view(-1, 2).sum(dim=0).tolist() - self._metrics[mode][name].append(total / count) - self._pending_metrics.clear() - - if images is not None and self.log_multimodal: - self._logs["images"].extend(gather_object(images)) - - output = { - "prompt_ids": prompt_ids, - "prompt_mask": prompt_mask, - "completion_ids": completion_ids, - "completion_mask": completion_mask, - "old_logps": old_logps, - "advantages": advantages, - } - if "pixel_values" in forward_kwargs: - output["pixel_values"] = forward_kwargs["pixel_values"] - if "image_grid_thw" in forward_kwargs: - output["image_grid_thw"] = forward_kwargs["image_grid_thw"] - if "pixel_attention_mask" in forward_kwargs: - output["pixel_attention_mask"] = forward_kwargs["pixel_attention_mask"] - if "spatial_shapes" in forward_kwargs: - output["spatial_shapes"] = forward_kwargs["spatial_shapes"] - if "image_sizes" in forward_kwargs: - output["image_sizes"] = forward_kwargs["image_sizes"] - if "token_type_ids" in forward_kwargs: - output["token_type_ids"] = forward_kwargs["token_type_ids"] - if "mm_token_type_ids" in forward_kwargs: - output["mm_token_type_ids"] = forward_kwargs["mm_token_type_ids"] - if "image_position_ids" in forward_kwargs: - output["image_position_ids"] = forward_kwargs["image_position_ids"] - if images is not None: - output["num_images"] = num_images - if num_tiles is not None: - output["num_tiles"] = num_tiles - return output - - @profiling_decorator - def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): - if return_outputs: - raise ValueError("The RLOOTrainer does not support returning outputs") - return self._compute_loss(model, inputs) - - def _compute_loss(self, model, inputs): - # 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"] - input_ids = torch.cat([prompt_ids, completion_ids], dim=1) - attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) - logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens - - # Compute the per_token_logps and the entropy at each position in the completion - per_token_logps, entropies, aux_loss = self._get_per_token_logps_and_entropies( - model, - input_ids, - attention_mask, - logits_to_keep, - compute_entropy=True, - compute_aux_loss=self.aux_loss_enabled, - pixel_values=inputs.get("pixel_values"), - image_grid_thw=inputs.get("image_grid_thw"), - num_images=inputs.get("num_images"), - pixel_attention_mask=inputs.get("pixel_attention_mask"), - spatial_shapes=inputs.get("spatial_shapes"), - num_tiles=inputs.get("num_tiles"), - image_sizes=inputs.get("image_sizes"), - token_type_ids=inputs.get("token_type_ids"), - mm_token_type_ids=inputs.get("mm_token_type_ids"), - image_position_ids=inputs.get("image_position_ids"), - ) - - logps = (per_token_logps * completion_mask).sum(1) # mask out padding and tokens after EOS - old_logps = inputs["old_logps"] - log_ratio = logps - old_logps - - # Compute the loss - advantages = inputs["advantages"] - coef_1 = torch.exp(log_ratio) - coef_2 = torch.clamp(coef_1, 1 - self.epsilon_low, 1 + self.epsilon_high) - per_sequence_loss1 = coef_1 * advantages - per_sequence_loss2 = coef_2 * advantages - per_sequence_loss = -torch.min(per_sequence_loss1, per_sequence_loss2) - loss = per_sequence_loss.mean() - - # Log the metrics - mode = "train" if self.model.training else "eval" - - # RLOO returns an unscaled loss (the HF Trainer divides by gradient accumulation), so add the aux term unscaled - if self.aux_loss_enabled: - loss = loss + self.router_aux_loss_coef * aux_loss - self._metrics[mode]["aux_loss"].append(self.accelerator.gather_for_metrics(aux_loss).mean().item()) - - # Entropy - entropy_stats = self.accelerator.reduce( - torch.stack([(entropies * completion_mask).sum(), completion_mask.sum().float()]), reduction="sum" - ) - self._metrics[mode]["entropy"].append((entropy_stats[0] / entropy_stats[1].clamp(min=1.0)).item()) - - # 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 - gathered_low_clip = self.accelerator.gather(is_low_clipped.float()) - self._metrics[mode]["clip_ratio/low_mean"].append(gathered_low_clip.nanmean().item()) - self._metrics[mode]["clip_ratio/low_min"].append(nanmin(gathered_low_clip).item()) - gathered_high_clip = self.accelerator.gather(is_high_clipped.float()) - self._metrics[mode]["clip_ratio/high_mean"].append(gathered_high_clip.nanmean().item()) - self._metrics[mode]["clip_ratio/high_max"].append(nanmax(gathered_high_clip).item()) - gathered_clip_ratio = self.accelerator.gather(is_region_clipped.float()) - self._metrics[mode]["clip_ratio/region_mean"].append(gathered_clip_ratio.nanmean().item()) - return loss - - # During eval, Trainer calls prediction_step. If no labels are present in the inputs, it only runs forward and - # returns logits. We override prediction_step to force compute_loss, because this trainer doesn't involve labels. - def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys: list[str] | None = None): - inputs = self._prepare_inputs(inputs) - with torch.no_grad(): - with self.compute_loss_context_manager(): - loss = self.compute_loss(model, inputs) - loss = loss.mean().detach() - return loss, None, None - - def log(self, logs: dict[str, float], start_time: float | None = None) -> None: - mode = "train" if self.model.training else "eval" - # Average the metrics - metrics = {} - for key, val in self._metrics[mode].items(): - # Filter out NaN values before averaging. A reward function that returns None for all samples - # in a batch produces NaN for that batch's metric. With logging_steps > 1, a naive sum()/len() - # would let a single NaN contaminate valid data from other batches. Only return None when no - # valid values remain (e.g. JSON loggers crash on float NaN). - valid = [v for v in val if not math.isnan(v)] - metrics[key] = sum(valid) / len(valid) if valid else None - - # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` - # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. - if mode == "eval": - metrics = {f"eval_{key}": val for key, val in metrics.items()} - - logs.update(metrics) - super().log(logs, start_time) - self._metrics[mode].clear() - - if self.accelerator.is_main_process and self.log_completions: - if is_rich_available(): - print_prompt_completions_sample( - self._logs["prompt"], - self._logs["completion"], - self._logs["rewards"], - self._logs["advantages"], - self.state.global_step, - self.num_completions_to_print, - extra=dict(self._logs["extra"]), - ) - - logging_backends = [] - if self.args.report_to and "wandb" in self.args.report_to and wandb.run is not None: - logging_backends.append(wandb) - if self.args.report_to and "trackio" in self.args.report_to: - logging_backends.append(trackio) - - table = { - "step": [self.state.global_step] * len(self._logs["prompt"]), - "prompt": [strip_images_from_messages(messages) for messages in self._logs["prompt"]], - "completion": [strip_images_from_messages(messages) for messages in self._logs["completion"]], - **self._logs["rewards"], - **self._logs["extra"], - "advantage": self._logs["advantages"], - } - - df_base = pd.DataFrame(table) - images_raw = self._logs["images"] or [] - - for logging_backend in logging_backends: - if images_raw: - images = [] - for image_list in self._logs["images"]: - images.append([logging_backend.Image(image) for image in image_list]) - df = pd.concat( - [df_base, pd.Series(images, name="image")], - axis=1, - copy=False, - ) - else: - df = df_base - - if self.log_unique_prompts: - # Prompts are conversations (lists of dicts), so compare them by their text - df = df[~df["prompt"].astype(str).duplicated()] - - logging_backend.log({"completions": logging_backend.Table(dataframe=df)}) - - # Ensure the model card is saved along with the checkpoint - def _save_checkpoint(self, model, trial): - if self.args.hub_model_id is None: - model_name = Path(self.args.output_dir).name - else: - model_name = self.args.hub_model_id.split("/")[-1] - self.create_model_card(model_name=model_name) - super()._save_checkpoint(model, trial) -class UnslothRLOOTrainer(_UnslothRLOOTrainer): - """ - - Trainer for the Reinforce Leave One Out (RLOO) method. This algorithm was initially proposed in the paper [Back to - Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in - LLMs](https://huggingface.co/papers/2402.14740). - - Example: - - ```python - >>> from trl import RLOOTrainer - >>> from trl.rewards import accuracy_reward - >>> from datasets import load_dataset - - >>> dataset = load_dataset("trl-lib/DeepMath-103K", split="train") - - >>> trainer = RLOOTrainer( - ... model="Qwen/Qwen2.5-0.5B-Instruct", - ... reward_funcs=accuracy_reward, - ... train_dataset=dataset, - ... ) - >>> trainer.train() - ``` - - Args: - model (`str` or [`~transformers.PreTrainedModel`] or [`~peft.PeftModel`]): - Model to be trained. Can be either: - - - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a - path to a *directory* containing model weights saved using - [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded - using `.from_pretrained` (where `` is derived from the model - config) with the keyword arguments in `args.model_init_kwargs`. If `dtype` is not specified in - `args.model_init_kwargs`, it defaults to `float32`. This differs from - [`~transformers.PreTrainedModel.from_pretrained`], where (since Transformers v5) the dtype is inferred - from the model config. - - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. - - A [`~peft.PeftModel`] object. Only causal language models are supported. - reward_funcs (`RewardFunc | list[RewardFunc]`): - Reward functions to be used for computing the rewards. To compute the rewards, we call all the reward - functions with the prompts and completions and sum the rewards. Can be either: - - - A single reward function, such as: - - A string: The *model ID* of a pretrained model hosted inside a model repo on huggingface.co, or a - path to a *directory* containing model weights saved using - [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded - using [`~transformers.AutoModelForSequenceClassification.from_pretrained`] with `num_labels=1` and the - keyword arguments in `args.model_init_kwargs`. - - A [`~transformers.PreTrainedModel`] object: Only sequence classification models are supported. - - A custom reward function: The function is provided with the prompts and the generated completions, - plus any additional columns in the dataset. It should return a list of rewards. Custom reward - functions can be either synchronous or asynchronous and can also return `None` when the reward is - not applicable to those samples. This is useful for multi-task training where different reward - functions apply to different types of samples. When a reward function returns `None` for a sample, - that reward function is excluded from the reward calculation for that sample. For more details, see - [Using a custom reward - function](#using-a-custom-reward-function). - - The trainer's state is also passed to the reward function. The trainer's state is an instance of - [`~transformers.TrainerState`] and can be accessed by accessing the `trainer_state` argument to the - reward function's signature. - - A list of reward functions, where each item can independently be any of the above types. Mixing different - types within the list (e.g., a string model ID and a custom reward function) is allowed. - args ([`RLOOConfig`], *optional*): - Configuration for this trainer. If `None`, a default configuration is used. - train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): - Dataset to use for training. It must include a column `"prompt"`. Any additional columns in the dataset is - ignored. The format of the samples can be either: - - - [Standard](dataset_formats#standard): Each sample contains plain text. - - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role - and content). - - When `train_dataset` is an [`~datasets.IterableDataset`] (e.g. a streaming dataset), `max_steps` must be - set in the training arguments, since its length cannot be inferred and the total number of training steps - is required to bound the training loop and configure the learning rate scheduler. - eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`], [`~datasets.DatasetDict`], [`~datasets.IterableDatasetDict`] or `dict[str, Dataset | IterableDataset]`): - Dataset to use for evaluation. It must meet the same requirements as `train_dataset`. - processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.ProcessorMixin`], *optional*): - Processing class used to process the data. The padding side must be set to "left". If `None`, the - processing class is loaded from the model's name with [`~transformers.AutoProcessor.from_pretrained`]. A - padding token, `tokenizer.pad_token`, must be set. If the processing class has not set a padding token, - `tokenizer.eos_token` will be used as the default. - reward_processing_classes ([`~transformers.PreTrainedTokenizerBase`] or `list[PreTrainedTokenizerBase]`, *optional*): - Processing classes corresponding to the reward functions specified in `reward_funcs`. Can be either: - - - A single processing class: Used when `reward_funcs` contains only one reward function. - - A list of processing classes: Must match the order and length of the reward functions in `reward_funcs`. - If set to `None`, or if an element of the list corresponding to a [`~transformers.PreTrainedModel`] is - `None`, the tokenizer for the model is automatically loaded using - [`~transformers.AutoTokenizer.from_pretrained`]. For elements in `reward_funcs` that are custom reward - functions (not [`~transformers.PreTrainedModel`]), the corresponding entries in `reward_processing_classes` - are ignored. - callbacks (list of [`~transformers.TrainerCallback`], *optional*): - List of callbacks to customize the training loop. Will add those to the list of default callbacks detailed - in [here](https://huggingface.co/docs/transformers/main_classes/callback). - - If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`] - method. - optimizers (`tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None]`, *optional*, defaults to `(None, None)`): - A tuple containing the optimizer and the scheduler to use. Will default to an instance of `AdamW` on your - model and a scheduler given by [`~transformers.get_linear_schedule_with_warmup`] controlled by `args`. - quantization_config ([`~transformers.BitsAndBytesConfig`], *optional*): - Quantization configuration used when loading the model from a model identifier. Combine with `peft_config` - for QLoRA training. Ignored if the model is already instantiated. - peft_config ([`~peft.PeftConfig`], *optional*): - PEFT configuration used to wrap the model. If `None`, the model is not wrapped. - - """ - def __init__( - self, - model, - reward_funcs, - args = None, - train_dataset = None, - eval_dataset = None, - processing_class = None, - reward_processing_classes = None, - callbacks = None, - quantization_config = None, - peft_config = None, - **kwargs - ): - if args is None: args = UnslothRLOOConfig() - _unsloth_precision_requested = tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) - use_bf16 = getattr(args, 'bf16', False) - if type(use_bf16) is not bool: use_bf16 = False - use_fp16 = getattr(args, 'fp16', False) - if type(use_fp16) is not bool: use_fp16 = False - force_float32 = False - try: - from unsloth_zoo.device_type import device_is_bf16_supported as _bf16_supported - except Exception: - _bf16_supported = torch.cuda.is_bf16_supported - try: - from unsloth_zoo.device_type import DEVICE_TYPE as _unsloth_device_type - except Exception: - _unsloth_device_type = None - full_finetuning = getattr(model, '_unsloth_full_finetuning', None) - if full_finetuning is None: full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1' - model_forced_float32 = getattr(model, '_unsloth_forced_float32', None) - if model_forced_float32 is None: model_forced_float32 = os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1' - if model_forced_float32 and not (full_finetuning and _bf16_supported()): - print('Unsloth: Switching to float32 training since model cannot work with float16') - force_float32 = True - mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') - dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None) - if dtype is None: dtype = model.get_input_embeddings().weight.dtype - from unsloth_zoo.utils import _get_dtype - dtype = _get_dtype(dtype) - float16 = dtype == torch.float16 - bfloat16 = dtype == torch.bfloat16 - float32 = dtype == torch.float32 - user_float32 = bool(getattr(model, '_unsloth_user_float32', False)) - if full_finetuning: - if bfloat16 and use_fp16: use_fp16 = False - if float16 and use_bf16: use_bf16 = False - if not force_float32 and bfloat16 and use_fp16 and _unsloth_device_type == 'xpu' and not _bf16_supported(): - print('Unsloth: The model was loaded in bfloat16 on purpose, so training uses bfloat16 mixed precision instead of float16.') - use_fp16 = False - use_bf16 = True - args.fp16 = False - args.bf16 = True - if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`') - if not force_float32 and (bfloat16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`') - if force_float32: - # Forced float32 training - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32' and float32 and user_float32 and not _bf16_supported(): - print('Unsloth: Model is in float32 and this GPU has no bfloat16 support, so training stays in float32. Pass fp16 = True to force float16 mixed precision instead.') - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32': - # Mixed precision training. bf16 only if the GPU supports it; V100/T4 use fp16. - use_bf16_amp = (not float16) and (_bf16_supported() or (bfloat16 and _unsloth_device_type == 'xpu')) - args.fp16 = not use_bf16_amp - args.bf16 = use_bf16_amp - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16_amp else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16_amp else 'fp16' - # args.mixed_precision is a new argument which needs to be set now - elif mixed_precision_dtype == 'bfloat16': - # Both False since bfloat16 full finetuning doesn't do any autocasting. - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif use_bf16 or use_fp16: - # transformers <5 exported this itself from fp16/bf16; 5.x dropped the write, so an - # explicit flag left it unset and GRPO readers defaulted to 'fp16', wrapping a - # bfloat16 model in a float16 autocast. See unslothai/unsloth#4891. - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16 else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16 else 'fp16' - - if getattr(args, 'eval_dataset', None) is not None and getattr(args, 'eval_strategy', 'no') == 'no': - args.eval_strategy = 'steps' - if getattr(args, 'eval_steps', None) is None: args.eval_steps = 0.1 - ga_steps = getattr(args, 'gradient_accumulation_steps', None) - if ga_steps is not None and ga_steps > 1: - from transformers import __version__ as transformers_version - if Version(transformers_version) <= Version('4.45.2'): - print('**** Unsloth: Please use our fixed gradient_accumulation_steps by updating transformers, TRL and Unsloth!\n' - '`pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth transformers trl unsloth_zoo`') - if getattr(args, 'eval_strategy', 'no') != 'no': - eval_bsz = getattr(args, 'per_device_eval_batch_size', 8) - eval_generations = getattr(args, 'num_generations_eval', None) or getattr(args, 'num_generations', None) or 1 - if eval_bsz == 8 and args.per_device_train_batch_size < eval_bsz and (args.per_device_train_batch_size * args.world_size) % eval_generations == 0: args.per_device_eval_batch_size = args.per_device_train_batch_size - if getattr(args, 'eval_accumulation_steps', None) is None and ga_steps is not None: args.eval_accumulation_steps = ga_steps - fp16_full_eval = getattr(args, 'fp16_full_eval', False) - if type(fp16_full_eval) is not bool: fp16_full_eval = False - bf16_full_eval = getattr(args, 'bf16_full_eval', False) - if type(bf16_full_eval) is not bool: bf16_full_eval = False - if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True - if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False - if force_float32: - args.bf16_full_eval = False - args.fp16_full_eval = False - elif os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16': - args.bf16_full_eval = True - args.fp16_full_eval = False - elif not bf16_full_eval and not fp16_full_eval: - args.bf16_full_eval = args.bf16 - args.fp16_full_eval = args.fp16 - if tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) != _unsloth_precision_requested: - from unsloth.models.loader_utils import sync_deepspeed_mixed_precision - sync_deepspeed_mixed_precision(args) - if model is not None: - from unsloth.models.loader_utils import apply_deepspeed_zero3_model_setup - apply_deepspeed_zero3_model_setup(model) - if model is not None: - _warnings_issued = getattr(model, 'warnings_issued', None) - if _warnings_issued is None: - model.warnings_issued = {} - elif not isinstance(_warnings_issued, dict): - try: - model.warnings_issued = dict(_warnings_issued) - except Exception: - model.warnings_issued = {} - if 'max_seq_length' not in locals() and not hasattr(args, 'max_seq_length'): - pass - else: - model_max_seq_length = getattr(model, 'max_seq_length', None) - args_max_seq_length = getattr(args, 'max_seq_length', None) - if args_max_seq_length is None and model_max_seq_length is not None: - max_seq_length = model.max_seq_length - if hasattr(args, 'max_seq_length'): args.max_seq_length = max_seq_length - elif args_max_seq_length is not None and model_max_seq_length is not None: - if args_max_seq_length > model_max_seq_length: - print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but ' - 'the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.') - args.max_seq_length = model_max_seq_length - if model is not None and hasattr(model, 'for_training'): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right' - if 'processing_class' in locals(): - if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right' - if hasattr(processing_class, 'tokenizer') and hasattr(processing_class.tokenizer, 'padding_side'): processing_class.tokenizer.padding_side = 'right' - other_metrics = [] - - from unsloth_zoo.logging_utils import PatchRLStatistics - PatchRLStatistics('rloo_trainer', other_metrics) - - # [TODO] Fix up DataParallel multiplying batch sizes - # [TODO] DDP works, but DP seems to not work? [TODO] - if getattr(args, "parallel_mode", None) == ParallelMode.NOT_DISTRIBUTED and args.n_gpu > 1: - if getattr(args, "_n_gpu", 1) != 1: - args._n_gpu = 1 - if "model" in locals(): - try: - from unsloth.models._gc_config_default import honour_gradient_checkpointing_off - honour_gradient_checkpointing_off(args, model) - except Exception: - pass - if "model" in locals() and hasattr(model, "for_training"): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if "model" in locals(): - try: - from unsloth.models._ddp_defaults import unsloth_default_ddp_find_unused_parameters - unsloth_default_ddp_find_unused_parameters(args, model) - except Exception: - pass - super().__init__( - model = model, - reward_funcs = reward_funcs, - args = args, - train_dataset = train_dataset, - eval_dataset = eval_dataset, - processing_class = processing_class, - reward_processing_classes = reward_processing_classes, - callbacks = callbacks, - quantization_config = quantization_config, - peft_config = peft_config,**kwargs) - # A tensor parallel model (tp_native.tp_trainer_setup) stays on its training path: for_inference would - # route the decoder layers onto decode paths outside the tensor parallel regions. - if "model" in locals() and hasattr(model, "for_inference") and not getattr(self, "_unsloth_tensor_parallel", False): - model.for_inference() - if hasattr(self, 'neftune_hook_handle'): - self.neftune_hook_handle.remove() - if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle - if getattr(args, 'neftune_noise_alpha', None) is not None: - model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha - pass - if hasattr(self, 'accelerator'): - scaler = self.accelerator.scaler - current_model = model - while hasattr(current_model, 'model'): - current_model.accelerator_scaler = scaler - current_model = current_model.model - current_model.accelerator_scaler = scaler - pass - if hasattr(self, 'train'): - self.train = MethodType(prepare_for_training_mode(self.__class__.train), self) - pass - if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'): - _vllm_tok = self.llm.get_tokenizer() - _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None) - if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None: - _vllm_tok.chat_template = _pc.chat_template - pass - if getattr(self, 'aux_loss_enabled', False) and hasattr(getattr(self, 'model', None), 'config'): - _text_config = self.model.config - if hasattr(_text_config, 'get_text_config'): _text_config = _text_config.get_text_config() - _n_experts = [getattr(_text_config, _k) for _k in ('num_local_experts', 'num_experts', 'n_routed_experts', 'moe_num_experts') if isinstance(getattr(_text_config, _k, None), int)] - if _n_experts and all(_n <= 1 for _n in _n_experts): - self.aux_loss_enabled = False - _text_config.output_router_logits = False - pass - -pass - - -if hasattr(logger, "addFilter"): - import logging - class HideLoggingMessage(logging.Filter): - def __init__(self, text): self.text = text - def filter(self, x): return not (self.text in x.getMessage()) - pass - logger.addFilter(HideLoggingMessage("`use_cache=True`")) - diff --git a/unsloth_compiled_cache/UnslothRewardTrainer.py b/unsloth_compiled_cache/UnslothRewardTrainer.py deleted file mode 100644 index 6a6d52a229119941f1268cf83e6aaec4c23b6801..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/UnslothRewardTrainer.py +++ /dev/null @@ -1,1797 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from unsloth_zoo.temporary_patches.common import torch_compile -from unsloth_zoo.temporary_patches.common import _maybe_compile -import functools -import transformers -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from trl.trainer.reward_trainer import (Any, AutoModelForSequenceClassification, AutoTokenizer, BitsAndBytesConfig, Callable, DataCollator, DataCollatorForPreference, Dataset, DatasetDict, EvalPrediction, IterableDataset, IterableDatasetDict, Path, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, RewardConfig, RewardTrainer, TrainerCallback, Version, _BaseTrainer, _tokenize, accelerate, clone_chat_template, contextlib, create_model_from_path, dataclass, defaultdict, disable_dropout_in_model, get_act_offloading_ctx_manager, get_config_model_id, get_dataset_column_names, get_peft_model, global_then_local_main_first, is_conversational, is_peft_available, is_peft_model, json, logger, logging, nn, os, pad, re, set_seed, suppress_seqcls_warning, torch, transformers, warnings, Any, AutoModelForSequenceClassification, AutoTokenizer, BitsAndBytesConfig, Callable, DataCollator, DataCollatorForPreference, Dataset, DatasetDict, EvalPrediction, IterableDataset, IterableDatasetDict, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, RewardConfig, TrainerCallback, Version, accelerate, clone_chat_template, contextlib, create_model_from_path, defaultdict, disable_dropout_in_model, get_act_offloading_ctx_manager, get_config_model_id, get_peft_model, is_peft_available, is_peft_model, logger, os, pad, re, set_seed, suppress_seqcls_warning, torch, transformers, PeftModel, PreTrainedModel, is_peft_available, logger, os, re, torch) - - -import os -import math -import logging -from typing import * -from dataclasses import dataclass, field -from packaging.version import Version -import torch -import numpy as np -from contextlib import nullcontext -from torch.nn import functional as F -import inspect -from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling -from transformers.training_args import ParallelMode -from unsloth_zoo.device_type import DEVICE_TYPE, DEVICE_TYPE_TORCH, device_synchronize - -# Wrap trainer with padding to right and enable training mode -import functools -from types import MethodType -try: - from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers -except: - def reset_unsloth_gradient_checkpointing_buffers(): pass -# Canonical reset lives in unsloth.models._utils so the SFT auto-packing wrapper and the plain -# Trainer loop can import the same helper; fall back to a no-op only if it can't be imported. -try: - from unsloth.models._utils import _unsloth_reset_stray_compile_cache -except Exception: - def _unsloth_reset_stray_compile_cache(self): pass -try: - from unsloth.models._utils import _unsloth_dataset_column_names -except Exception: - def _unsloth_dataset_column_names(dataset): return dataset.column_names -# Drops/renames config arguments the installed TRL no longer accepts, so a -# script pinned to an older TRL keeps working after an upgrade. Falls back to -# the historical raw passthrough so this can never break trainer construction. -try: - from unsloth.models.rl_config_compat import filter_config_init_kwargs as _unsloth_filter_config_init_kwargs - # A cache file generated here can be imported by an older Unsloth whose filter - # predates `mirrored_from`, so drop the argument rather than raise TypeError. - if "mirrored_from" not in inspect.signature(_unsloth_filter_config_init_kwargs).parameters: - _unsloth_filter_config_init_kwargs_old = _unsloth_filter_config_init_kwargs - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): - return _unsloth_filter_config_init_kwargs_old(config_class, kwargs) -except Exception: - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): return kwargs -def prepare_for_training_mode(f): - @functools.wraps(f) - def wrapper(self, *args, **kwargs): - # Drop any torch.compile graph cache poisoned by a stray pre-train forward. - try: - _unsloth_reset_stray_compile_cache(self) - except Exception: - pass - # Finish the previous W&B run if this is a subsequent train() call. - # We do this at the START of train() (not the end) so that - # evaluate() / log() still work after train() completes. - # HF's WandbCallback.setup() will call wandb.init() for the new run. - # See: https://github.com/unslothai/unsloth/issues/3954 - if getattr(self, '_unsloth_training_completed', False): - try: - import wandb - if wandb.run is not None: - wandb.finish() - # Reset HF's WandbCallback so it calls wandb.init() for the new run - for cb in self.callback_handler.callbacks: - if type(cb).__name__ == 'WandbCallback': - cb._initialized = False - break - except: - pass - # Enable training mode - _was_training = None - # Restore the GC mode the model was configured with at setup; fall back to - # the training args only when it wasn't recorded (issue #4735). Use hasattr, - # not a None sentinel, so a deliberately-recorded None is restored verbatim. - _model = getattr(self, 'model', None) - if hasattr(_model, '_unsloth_gradient_checkpointing'): - use_gc = _model._unsloth_gradient_checkpointing - else: - use_gc = getattr(self.args, 'gradient_checkpointing', True) - if hasattr(self, 'model') and hasattr(self.model, "training"): - _was_training = self.model.training - if hasattr(self, 'model') and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - output = f(self, *args, **kwargs) - # Restore previous mode when possible - if hasattr(self, 'model') and hasattr(self.model, "for_inference"): - if _was_training is False: - self.model.for_inference() - elif _was_training is True and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - # Reset gradient checkpointing buffers to free memory while staying ready for next run - try: - reset_unsloth_gradient_checkpointing_buffers() - except: - pass - # Mark that training completed so the next train() call can - # finish this W&B run before starting a new one - self._unsloth_training_completed = True - return output - return wrapper -pass - -torch_compile_options = { - "epilogue_fusion" : True, - "max_autotune" : False, - "shape_padding" : True, - "trace.enabled" : False, - "triton.cudagraphs" : False, -} - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_hidden_states_selective_log_softmax( - hidden_states: torch.Tensor, - lm_head: torch.Tensor, - index: torch.Tensor, - chunks: int = 4, - logit_scale_multiply: float = 0.0, - logit_scale_divide: float = 0.0, - logit_softcapping: float = 0.0, - temperature: float = 1.0, - # Rows per chunk cap. Read HERE, in the default, not in the body: the body - # is traced with fullgraph = True, and `os.environ` is an unsupported op - # there on torch 2.4 -- Dynamo raises - # torch._dynamo.exc.Unsupported: const method call bytes.decode - # from os._Environ.__getitem__, which the eager fallback does not catch - # (it only catches recompile-limit and disabled-hook breaks), so the very - # first call would die even with the variable unset. A default is evaluated - # once when the def runs, which is import time, outside any traced region. - # It is also a plain int argument, so Dynamo guards on it instead of - # constant-folding an unguarded read (2.7+ never notice a later change). - # A non-numeric value is ignored rather than raised on, so a typo cannot - # break the import. 0 keeps the previous chunk boundaries exactly. - max_rows_per_chunk: int = ( - int(os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip()) - if os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip().isdigit() - else 0 - ), - lm_head_bias: Optional[torch.Tensor] = None, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, - # Read in the default for the same Dynamo reason as max_rows_per_chunk. Pairs with vLLM's - # head_dtype=float32 that load_vllm requests under the same flag. - fp32_lm_head: bool = ( - os.environ.get("UNSLOTH_VLLM_FP32_LM_HEAD", "0").strip().lower() in ("1", "true", "on", "yes") - ), - # Older Unsloth scores old / ref log-probs here and the policy with the fused kernel: one scorer keeps ratio at 1. - fused_linear: bool = ( - os.environ.get("UNSLOTH_FUSED_LINEAR", "1") != "0" - and os.environ.get("UNSLOTH_GRPO_CHUNKED_LOGPROBS", "1") != "0" - ), -) -> torch.Tensor: - # All Unsloth Zoo code licensed under AGPL3 - from unsloth_zoo.fused_losses import chunked_logprob as _chunked_logprob - if (fused_linear and _chunked_logprob._FUSED_LINEAR and not fp32_lm_head - and support_top_k == 0 and support_top_p == 1.0 and support_min_p == 0.0 - and support_excluded_ids is None): - return _chunked_logprob.unsloth_chunked_logprobs( - hidden_states, lm_head, index, - temperature = temperature, - logit_scale_multiply = logit_scale_multiply, - logit_scale_divide = logit_scale_divide, - logit_softcapping = logit_softcapping, - lm_head_bias = lm_head_bias, - max_rows_per_chunk = max_rows_per_chunk if max_rows_per_chunk > 0 else None, - ).reshape((hidden_states.shape[0], hidden_states.shape[1])) - # Reshape on this tensor's own last dim: a no-op, so a wrong-width caller - # cannot have its row count silently rewritten and instead fails at the - # matmul below, which prints both operands. Do not swap in a bare - # torch._check: it reports only "Expected cond to be True", naming neither - # operand, and Dynamo rejects a message-carrying one. Callers dispatch on - # the width first -- see `compute_logprobs_chunk`, the packed path and - # `_pg_grad_forward`. - flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) - flat_index = index.reshape(-1) - - # Each chunk materialises rows x vocab logits and then a float32 copy of - # them, all on the device holding the output head. With a large vocabulary - # and a fixed chunk count that grows with the batch, so the peak scales with - # the batch rather than staying bounded. max_rows_per_chunk caps the rows - # per chunk instead, which is pure loop splitting: more, smaller chunks, - # same concatenated result. 0 (the default) keeps the previous chunk - # boundaries exactly. - if max_rows_per_chunk > 0: - n_rows = flat_hidden_states.shape[0] - chunks = max(chunks, -(-n_rows // max_rows_per_chunk)) - chunks = min(chunks, max(n_rows, 1)) - - # A one-row chunk's head gradient is a K=1 matmul Inductor mis-lowers under dynamic shapes (illegal - # memory access or NaN head gradients, torch 2.13), so keep chunks at 2+ rows when the head trains. - # Exactly `chunks` row ranges, so the unrolled loop length never depends on T: torch.chunk - # returns ceil(T / ceil(T / chunks)) pieces and recompiled per new count; tensor_split guards - # on T % chunks. At most rows // 2 chunks when the head trains keeps every piece at 2+ rows. - lm_head_grad = torch.is_grad_enabled() and lm_head.requires_grad - if lm_head_grad and flat_hidden_states.shape[0] < 2 * chunks: - chunks = max(flat_hidden_states.shape[0] // 2, 1) - - n_rows = flat_hidden_states.shape[0] - chunks = max(1, min(chunks, n_rows)) - bounds = [(i * n_rows) // chunks for i in range(chunks + 1)] - chunked_hidden_states = [flat_hidden_states[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - - # Cast once, outside the loop, so every chunk shares one fp32 head copy (also for autograd). - if fp32_lm_head: - lm_head = lm_head.to(torch.float32) - if lm_head_bias is not None: - lm_head_bias = lm_head_bias.to(torch.float32) - - all_per_token_logps = [] - - for chunk_hidden_states, chunk_index in zip(chunked_hidden_states, chunked_index): - # When the model is dispatched over several devices, the output head can - # sit on a different one from the hidden states, because accelerate - # places the tail of the model on the last device it fills. Co-locate on - # the head's device before the matmul, otherwise this raises - # Unhandled FakeTensor Device Propagation for aten.mm.default, - # found two different devices cuda:0, cuda:1 - # On a single device every .to() here is a no-op and the result is - # bit-identical to before. - chunk_hidden_states = chunk_hidden_states.to(device = lm_head.device, dtype = lm_head.dtype) - chunk_index = chunk_index.to(lm_head.device) - if fp32_lm_head: - # The GRPO loss runs this under bf16 / fp16 autocast, which would downcast the matmul. - with torch.autocast(device_type = chunk_hidden_states.device.type, enabled = False): - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - elif lm_head_bias is None: - chunk_logits = chunk_hidden_states @ lm_head.t() - else: - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - - if logit_scale_multiply != 0.0: - chunk_logits = chunk_logits * logit_scale_multiply - if logit_scale_divide != 0.0: - chunk_logits = chunk_logits / logit_scale_divide - if logit_softcapping != 0.0: - chunk_logits = logit_softcapping * torch.tanh(chunk_logits / logit_softcapping) - - chunk_logits = chunk_logits.to(torch.float32) - - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - per_token_logps = _truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - ) - else: - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - # Return to the caller's device so the concatenation below and every - # downstream consumer see the device they started on. - all_per_token_logps.append(per_token_logps.to(hidden_states.device)) - - all_per_token_logps = torch.concat(all_per_token_logps) - - all_per_token_logps = all_per_token_logps.reshape((hidden_states.shape[0], hidden_states.shape[1])) - return all_per_token_logps - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_selective_log_softmax( - logits, - index, - temperature: float = 1.0, - chunks: int = 4, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, -): - # Fixed `chunks` row ranges; see chunked_hidden_states_selective_log_softmax. - flat_logits = logits.reshape(-1, logits.shape[-1]) - flat_index = index.reshape(-1) - # Fewer rows than chunks: one row per range, as torch.chunk did, never empty ranges. - chunks = max(1, min(chunks, flat_logits.shape[0])) - bounds = [(i * flat_logits.shape[0]) // chunks for i in range(chunks + 1)] - chunked_logits = [flat_logits[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - all_per_token_logps = [] - # Per-chunk selective_log_softmax. - for chunk_logits, chunk_index in zip(chunked_logits, chunked_index): - chunk_logits = chunk_logits.to(torch.float32) - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - all_per_token_logps.append(_truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - )) - continue - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - all_per_token_logps.append(per_token_logps) - pass - all_per_token_logps = torch.concat(all_per_token_logps) - all_per_token_logps = all_per_token_logps.reshape((logits.shape[0], logits.shape[1])) - return all_per_token_logps - -def calculate_pad_tokens_in_prompt( - input_ids: torch.Tensor, - logits_to_keep: int, - pad_token_id: int -) -> torch.Tensor: - """Count left-padded tokens per sequence, e.g. [pad, pad, pad, cat] -> 3.""" - if logits_to_keep >= input_ids.shape[1]: - raise ValueError("logits_to_keep must be smaller than the sequence length.") - - prompt_section = input_ids[:, :-logits_to_keep] - - padding_mask = (prompt_section == pad_token_id) - - pad_token_counts = padding_mask.sum(dim=1) - - return pad_token_counts - -def create_completion_attention_mask( - completion_input_ids: torch.Tensor, - left_pad_tokens_per_prompt: torch.Tensor, - max_left_pad: int, - pad_token_id: int -) -> torch.Tensor: - """Build a completion mask that zeros leading prompt and trailing pad tokens. - - For [p,p,p,c,c,c,pad,pad,pad] (p=sliced prompt, c=completion, pad=padding) - this returns [0,0,0,1,1,1,0,0,0]. - """ - batch_size, completion_len = completion_input_ids.shape - device = completion_input_ids.device - - num_tokens_to_mask = max_left_pad - left_pad_tokens_per_prompt - - indices = torch.arange(completion_len, device=device).unsqueeze(0) - shift_mask = indices >= num_tokens_to_mask.unsqueeze(1) - - non_padding_mask = (completion_input_ids != pad_token_id) - - final_mask = shift_mask & non_padding_mask - - return final_mask - -def left_pack_padding(tensor: torch.Tensor, pad_id: int) -> torch.Tensor: - """Move all padding tokens in each sequence to the right.""" - mask = (tensor != pad_id) - # stable=True since the binary mask is unordered. - sorted_indices = torch.argsort(mask, dim=1, descending=True, stable=True) - packed_tensor = torch.gather(tensor, 1, sorted_indices) - return packed_tensor - -def align_logprobs_with_mask( - logprob_tensor: torch.Tensor, - attention_mask: torch.Tensor, - pad_value: float = 0.0 -) -> torch.Tensor: - """Align a log probability tensor with a given attention mask.""" - - device = logprob_tensor.device - batch_size, logprob_seq_len = logprob_tensor.shape - mask_seq_len = attention_mask.shape[1] - - padded_logprobs = torch.full( - attention_mask.shape, - fill_value=pad_value, - dtype=logprob_tensor.dtype, - device=device - ) - - left_pad_counts = torch.argmax(attention_mask, dim=1) - - cols = torch.arange(logprob_seq_len, device=device) - dest_indices = left_pad_counts.unsqueeze(1) + cols - - # Destination row indices, shape [batch_size, logprob_seq_len]. - row_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand_as(dest_indices) - - # Keep only in-bounds destinations, then scatter via advanced indexing. - valid_mask = dest_indices < mask_seq_len - valid_rows = row_indices[valid_mask] - valid_cols = dest_indices[valid_mask] - valid_vals = logprob_tensor[valid_mask] - padded_logprobs[valid_rows, valid_cols] = valid_vals - - return padded_logprobs - -def align_completion_tool_mask( - tool_mask: torch.Tensor, - completion_mask: torch.Tensor, -) -> torch.Tensor: - """Align a raw completion-length tool/env mask with Unsloth's repacked loss mask.""" - if tool_mask is None: - return completion_mask - if tool_mask.shape[0] != completion_mask.shape[0]: - raise ValueError("tool_mask batch size must match completion_mask batch size.") - - tool_mask = tool_mask.to(device=completion_mask.device) - if tool_mask.shape == completion_mask.shape: - aligned_tool_mask = tool_mask - else: - aligned_tool_mask = align_logprobs_with_mask( - tool_mask, - completion_mask, - pad_value=0, - ) - return completion_mask * aligned_tool_mask.to(dtype=completion_mask.dtype) - -def autotune_batch_and_chunks( - total_input_rows, - seq_len, - hidden_size, - vocab_size, - dtype_bytes=16, - multiplier=None -): - if multiplier is None: - final_m = max(4, seq_len // 4096) - else: - final_m = multiplier - - if torch.cuda.is_available(): - free_bytes, _ = torch.cuda.mem_get_info() - limit_gb = (free_bytes / (1024**3))*.80 - elif hasattr(torch, "xpu") and torch.xpu.is_available(): - # XPU: estimate free memory as total - reserved. - total_mem = torch.xpu.get_device_properties(0).total_memory - reserved_mem = torch.xpu.memory_reserved() - free_bytes = total_mem - reserved_mem - limit_gb = (free_bytes / (1024**3)) * 0.80 - else: - # Fallback: assume 8GB available. - limit_gb = 8.0 - - bytes_to_gb = 1024**3 - - b_vals = torch.arange(total_input_rows, 0, -1, device='cpu', dtype=torch.float32) - - hidden_gb = (b_vals * seq_len * hidden_size * dtype_bytes) / bytes_to_gb - - base_logits = ((b_vals/total_input_rows) * b_vals * seq_len * vocab_size * dtype_bytes) / bytes_to_gb - logits_gb = base_logits / final_m - - total_mem_gb = hidden_gb + logits_gb - - valid_mask = total_mem_gb <= limit_gb - valid_indices = torch.nonzero(valid_mask, as_tuple=False) - - if valid_indices.shape[0] == 0: - #This means your GPU will OOM - # Capped at the row count: unsloth's no-grad pass divides rows by this without max(1, ...). - return max(1, min(4, total_input_rows)), final_m - - best_idx = valid_indices[0].item() - final_b = int(b_vals[best_idx].item()) - - return final_b, final_m - -def sanitize_logprob(logprob): - """Local port of trl.scripts.vllm_serve.sanitize_logprob. - Filters NaN logprobs from vLLM outputs.""" - value = logprob.logprob - if math.isnan(value): - logging.getLogger(__name__).warning( - f"Generated NaN logprob, token logprob '{logprob}' will be ignored" - ) - return None - return value -@dataclass -class UnslothRewardConfig(RewardConfig): - """ - - Configuration class for the [`RewardTrainer`]. - - This class includes only the parameters that are specific to Reward training. For a full list of training - arguments, please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this - class may differ from those in [`~transformers.TrainingArguments`]. - - Using [`~transformers.HfArgumentParser`] we can turn this class into - [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the - command line. - - Parameters: - > Parameters that control the model - - model_init_kwargs (`dict[str, Any]`, *optional*): - Keyword arguments for [`~transformers.AutoModelForSequenceClassification.from_pretrained`], used when the - `model` argument of the [`RewardTrainer`] is provided as a string. The `revision` value is also used when - loading the tokenizer. - trust_remote_code (`bool`, *optional*, defaults to `False`): - Whether to allow loading models and tokenizers that ship custom Python code from the Hub. Forwarded to - [`~transformers.AutoModelForSequenceClassification.from_pretrained`] and - [`~transformers.AutoTokenizer.from_pretrained`]. - chat_template_path (`str`, *optional*): - If specified, sets the model's chat template. This can either be the path to a tokenizer (local directory - or Hugging Face Hub model) or a direct path to a Jinja template file. When using a Jinja file, you must - ensure that any special tokens referenced in the template are added to the tokenizer and that the model's - embedding layer is resized accordingly. - disable_dropout (`bool`, *optional*, defaults to `True`): - Whether to disable dropout in the model. - - > Parameters that control the data preprocessing - - dataset_num_proc (`int`, *optional*): - Number of processes to use for processing the dataset. - eos_token (`str`, *optional*): - Token used to indicate the end of a turn or sequence. If `None`, it defaults to - `processing_class.eos_token`. - max_length (`int` or `None`, *optional*, defaults to `1024`): - Maximum length of the tokenized sequence. Samples are filtered out if either chosen or rejected sequence - exceeds this value. If `None`, no filtering is applied. - pad_to_multiple_of (`int`, *optional*): - If set, the sequences will be padded to a multiple of this value. - - > Parameters that control the training - - center_rewards_coefficient (`float`, *optional*): - Coefficient to incentivize the reward model to output mean-zero rewards (proposed by [this - paper](https://huggingface.co/papers/2312.09244), Eq. 2). Recommended value: `0.01`. - activation_offloading (`bool`, *optional*, defaults to `False`): - Whether to offload the activations to the CPU. - - > Deprecated parameters - - pad_token: - - - - Parameter `pad_token` is deprecated and will be removed in version v2.0.0. Set `tokenizer.pad_token` - directly and pass it as `processing_class` to the trainer instead. - - - - > [!NOTE] - > These parameters have default values different from [`~transformers.TrainingArguments`]: - > - `logging_steps`: Defaults to `10` instead of `500`. - > - `gradient_checkpointing`: Defaults to `True` instead of `False`. - > - `bf16`: Defaults to `True` if `fp16` is not set, instead of `False`. - > - `learning_rate`: Defaults to `1e-4` instead of `5e-5`. - - """ - _unsloth_accepts_adamw_flash = True - vllm_sampling_params: Optional[Any] = field( - default = None, - metadata = {'help': 'vLLM SamplingParams'}, - ) - unsloth_num_chunks : Optional[int] = field( - default = -1, - metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'}, - ) - unsloth_logit_chunk_multiplier : Optional[int] = field( - default = None, - metadata = {'help': 'Multiplier for chunked logit computations.'}, - ) - unsloth_grpo_mini_batch : Optional[int] = field( - default = None, - metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'}, - ) - max_seq_length : Optional[int] = field( - default = None, - metadata = {'help': 'Maximum sequence length to truncate to.'}, - ) - def __init__( - self, - output_dir = None, - per_device_train_batch_size = 4, - num_train_epochs = 3.0, - max_steps = -1, - learning_rate = 5e-05, - lr_scheduler_type = 'linear', - lr_scheduler_kwargs = None, - warmup_steps = 0.1, - optim = 'adamw_8bit', - optim_args = None, - weight_decay = 0.001, - adam_beta1 = 0.9, - adam_beta2 = 0.999, - adam_epsilon = 1e-08, - optim_target_modules = None, - gradient_accumulation_steps = 2, - average_tokens_across_devices = True, - max_grad_norm = 1.0, - label_smoothing_factor = 0.0, - bf16 = False, - fp16 = False, - bf16_full_eval = False, - fp16_full_eval = False, - tf32 = None, - gradient_checkpointing = True, - gradient_checkpointing_kwargs = None, - torch_compile = False, - torch_compile_backend = None, - torch_compile_mode = None, - use_liger_kernel = False, - liger_kernel_config = None, - use_cache = False, - neftune_noise_alpha = None, - torch_empty_cache_steps = 250, - auto_find_batch_size = False, - logging_strategy = 'steps', - logging_steps = 1, - logging_first_step = False, - log_on_each_node = True, - logging_nan_inf_filter = False, - include_num_input_tokens_seen = False, - log_level = 'passive', - log_level_replica = 'warning', - disable_tqdm = None, - report_to = 'none', - run_name = None, - project = 'huggingface', - trackio_space_id = 'trackio', - eval_strategy = 'no', - eval_steps = None, - eval_delay = 0, - per_device_eval_batch_size = 4, - prediction_loss_only = False, - eval_on_start = False, - eval_do_concat_batches = True, - eval_use_gather_object = False, - eval_accumulation_steps = 2, - batch_eval_metrics = False, - save_only_model = False, - save_strategy = 'steps', - save_steps = 500, - save_on_each_node = False, - save_total_limit = None, - enable_jit_checkpoint = False, - push_to_hub = False, - hub_token = None, - hub_private_repo = None, - hub_model_id = None, - hub_strategy = 'every_save', - hub_always_push = False, - hub_revision = None, - load_best_model_at_end = False, - metric_for_best_model = None, - greater_is_better = None, - ignore_data_skip = False, - restore_callback_states_from_checkpoint = False, - full_determinism = False, - seed = 3407, - data_seed = None, - use_cpu = False, - accelerator_config = None, - parallelism_config = None, - dataloader_drop_last = False, - dataloader_num_workers = 0, - dataloader_pin_memory = True, - dataloader_persistent_workers = False, - dataloader_prefetch_factor = None, - remove_unused_columns = True, - label_names = None, - train_sampling_strategy = 'random', - length_column_name = 'length', - ddp_find_unused_parameters = None, - ddp_bucket_cap_mb = None, - ddp_broadcast_buffers = None, - ddp_backend = None, - ddp_timeout = 1800, - fsdp = None, - fsdp_config = None, - deepspeed = None, - debug = '', - skip_memory_metrics = True, - do_train = False, - do_eval = False, - do_predict = False, - resume_from_checkpoint = None, - warmup_ratio = None, - logging_dir = None, - local_rank = -1, - model_init_kwargs = None, - trust_remote_code = False, - chat_template_path = None, - disable_dropout = True, - dataset_num_proc = None, - eos_token = None, - max_length = 1024, - pad_to_multiple_of = None, - center_rewards_coefficient = None, - activation_offloading = False, - pad_token = None, - vllm_sampling_params = None, - unsloth_num_chunks = -1, - unsloth_logit_chunk_multiplier = None, - unsloth_grpo_mini_batch = None, - max_seq_length = None, - **kwargs, - ): - if learning_rate < 1e-7: print(f'Unsloth: Your learning rate of `{learning_rate}` is too small and less than 1e-7! Consider increasing it, otherwise gradient updates will be close to 0!') - if learning_rate > 1: print(f'Unsloth: Your learning rate of `{learning_rate}` is way too larger > 1! Consider decreasing it to 1e-1, otherwise gradient updates will explode!') - _unsloth_fsdp2 = os.environ.get('FSDP_VERSION', '').strip() == '2' - _unsloth_fsdp_cfg = fsdp_config - if fsdp and isinstance(_unsloth_fsdp_cfg, str) and os.path.isfile(_unsloth_fsdp_cfg): - try: - import json as _unsloth_json - with open(_unsloth_fsdp_cfg, encoding = 'utf-8') as _unsloth_f: - _unsloth_fsdp_cfg = _unsloth_json.load(_unsloth_f) - except Exception: - _unsloth_fsdp_cfg = None - if fsdp and _unsloth_fsdp_cfg is None: - _unsloth_fsdp_cfg = {} - if fsdp and isinstance(_unsloth_fsdp_cfg, dict): - _unsloth_fsdp2 = _unsloth_fsdp2 or str(_unsloth_fsdp_cfg.get('fsdp_version', _unsloth_fsdp_cfg.get('version', 1))) == '2' - _unsloth_flash_adamw = str(getattr(optim, 'value', optim)).strip().lower() == 'adamw_flash' - if _unsloth_flash_adamw: optim = 'adamw_torch' - if _unsloth_fsdp2 and str(getattr(optim, 'value', optim)) == 'adamw_8bit': - print("Unsloth: FSDP2 shards parameters as DTensors, which bitsandbytes' 8-bit optimizers cannot step. Switching optim = 'adamw_8bit' to 'adamw_torch_fused'.") - optim = 'adamw_torch_fused' - if num_train_epochs is None: - num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override - if output_dir is None and save_strategy == 'steps' and save_steps == 500: - output_dir = 'unsloth_training_checkpoints' - save_strategy = 'no' - try: - from unsloth_zoo.dataset_num_proc import get_dataset_num_proc as _unsloth_get_dataset_num_proc - except Exception: - try: - from unsloth.dataset_num_proc import get_dataset_num_proc as _unsloth_get_dataset_num_proc - except Exception: - _unsloth_get_dataset_num_proc = None - if _unsloth_get_dataset_num_proc is not None: - dataset_num_proc = _unsloth_get_dataset_num_proc(dataset_num_proc, serial_as_none = True) - if os.environ.get('UNSLOTH_ENABLE_FLEX_ATTENTION', '0') == '1': - from unsloth_zoo.flex_attention import HAS_FLEX_ATTENTION - if HAS_FLEX_ATTENTION and pad_to_multiple_of is None: - from unsloth_zoo.flex_attention import FLEX_ATTENTION_BLOCK_SIZE - pad_to_multiple_of = FLEX_ATTENTION_BLOCK_SIZE - - if eval_steps is not None and eval_strategy != 'steps': - print(f'Unsloth: `eval_steps = {eval_steps}` is ignored because `eval_strategy` is {getattr(eval_strategy, "value", eval_strategy)!r}. Set `eval_strategy = "steps"` to evaluate every `eval_steps` steps.') - - - # One dict so the filter sees the mirrored parameters AND `**kwargs`: - # filtering kwargs alone would double-bind any argument TRL renamed, - # since the new name is itself a mirrored parameter. - _unsloth_config_arguments = dict( - output_dir = output_dir, - per_device_train_batch_size = per_device_train_batch_size, - num_train_epochs = num_train_epochs, - max_steps = max_steps, - learning_rate = learning_rate, - lr_scheduler_type = lr_scheduler_type, - lr_scheduler_kwargs = lr_scheduler_kwargs, - warmup_steps = warmup_steps, - optim = optim, - optim_args = optim_args, - weight_decay = weight_decay, - adam_beta1 = adam_beta1, - adam_beta2 = adam_beta2, - adam_epsilon = adam_epsilon, - optim_target_modules = optim_target_modules, - gradient_accumulation_steps = gradient_accumulation_steps, - average_tokens_across_devices = average_tokens_across_devices, - max_grad_norm = max_grad_norm, - label_smoothing_factor = label_smoothing_factor, - bf16 = bf16, - fp16 = fp16, - bf16_full_eval = bf16_full_eval, - fp16_full_eval = fp16_full_eval, - tf32 = tf32, - gradient_checkpointing = gradient_checkpointing, - gradient_checkpointing_kwargs = gradient_checkpointing_kwargs, - torch_compile = torch_compile, - torch_compile_backend = torch_compile_backend, - torch_compile_mode = torch_compile_mode, - use_liger_kernel = use_liger_kernel, - liger_kernel_config = liger_kernel_config, - use_cache = use_cache, - neftune_noise_alpha = neftune_noise_alpha, - torch_empty_cache_steps = torch_empty_cache_steps, - auto_find_batch_size = auto_find_batch_size, - logging_strategy = logging_strategy, - logging_steps = logging_steps, - logging_first_step = logging_first_step, - log_on_each_node = log_on_each_node, - logging_nan_inf_filter = logging_nan_inf_filter, - include_num_input_tokens_seen = include_num_input_tokens_seen, - log_level = log_level, - log_level_replica = log_level_replica, - disable_tqdm = disable_tqdm, - report_to = report_to, - run_name = run_name, - project = project, - trackio_space_id = trackio_space_id, - eval_strategy = eval_strategy, - eval_steps = eval_steps, - eval_delay = eval_delay, - per_device_eval_batch_size = per_device_eval_batch_size, - prediction_loss_only = prediction_loss_only, - eval_on_start = eval_on_start, - eval_do_concat_batches = eval_do_concat_batches, - eval_use_gather_object = eval_use_gather_object, - eval_accumulation_steps = eval_accumulation_steps, - batch_eval_metrics = batch_eval_metrics, - save_only_model = save_only_model, - save_strategy = save_strategy, - save_steps = save_steps, - save_on_each_node = save_on_each_node, - save_total_limit = save_total_limit, - enable_jit_checkpoint = enable_jit_checkpoint, - push_to_hub = push_to_hub, - hub_token = hub_token, - hub_private_repo = hub_private_repo, - hub_model_id = hub_model_id, - hub_strategy = hub_strategy, - hub_always_push = hub_always_push, - hub_revision = hub_revision, - load_best_model_at_end = load_best_model_at_end, - metric_for_best_model = metric_for_best_model, - greater_is_better = greater_is_better, - ignore_data_skip = ignore_data_skip, - restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint, - full_determinism = full_determinism, - seed = seed, - data_seed = data_seed, - use_cpu = use_cpu, - accelerator_config = accelerator_config, - parallelism_config = parallelism_config, - dataloader_drop_last = dataloader_drop_last, - dataloader_num_workers = dataloader_num_workers, - dataloader_pin_memory = dataloader_pin_memory, - dataloader_persistent_workers = dataloader_persistent_workers, - dataloader_prefetch_factor = dataloader_prefetch_factor, - remove_unused_columns = remove_unused_columns, - label_names = label_names, - train_sampling_strategy = train_sampling_strategy, - length_column_name = length_column_name, - ddp_find_unused_parameters = ddp_find_unused_parameters, - ddp_bucket_cap_mb = ddp_bucket_cap_mb, - ddp_broadcast_buffers = ddp_broadcast_buffers, - ddp_backend = ddp_backend, - ddp_timeout = ddp_timeout, - fsdp = fsdp, - fsdp_config = fsdp_config, - deepspeed = deepspeed, - debug = debug, - skip_memory_metrics = skip_memory_metrics, - do_train = do_train, - do_eval = do_eval, - do_predict = do_predict, - resume_from_checkpoint = resume_from_checkpoint, - warmup_ratio = warmup_ratio, - logging_dir = logging_dir, - local_rank = local_rank, - model_init_kwargs = model_init_kwargs, - trust_remote_code = trust_remote_code, - chat_template_path = chat_template_path, - disable_dropout = disable_dropout, - dataset_num_proc = dataset_num_proc, - eos_token = eos_token, - max_length = max_length, - pad_to_multiple_of = pad_to_multiple_of, - center_rewards_coefficient = center_rewards_coefficient, - activation_offloading = activation_offloading, - pad_token = pad_token,**kwargs) - super().__init__(**_unsloth_filter_config_init_kwargs(RewardConfig, _unsloth_config_arguments, mirrored_from = __class__)) - self.vllm_sampling_params = vllm_sampling_params - self.unsloth_num_chunks = unsloth_num_chunks - if unsloth_grpo_mini_batch is not None: - if self.generation_batch_size >= unsloth_grpo_mini_batch: - self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch - else: - raise ValueError( - f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, " - f"which is self.per_device_train_batch_size * gradient_accumulation_steps." - ) - self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier - self.max_seq_length = max_seq_length - # Unsloth: keep the reentrant checkpoint path - if getattr(self, 'gradient_checkpointing', False): - _gc_kwargs = getattr(self, 'gradient_checkpointing_kwargs', None) or {} - if _gc_kwargs.get('context_fn') is None and not _gc_kwargs.get('debug', False): - _gc_kwargs['use_reentrant'] = True - self.gradient_checkpointing_kwargs = _gc_kwargs - from unsloth.optimizers.flash_adamw import is_flash_adamw_marked, mark_flash_adamw, unmark_flash_adamw - if _unsloth_flash_adamw or is_flash_adamw_marked(self): - mark_flash_adamw(self) - else: - unmark_flash_adamw(self) - -pass - -class _UnslothRewardTrainer(_BaseTrainer): - """""" - - loss_is_scaled_for_ga = False - - _tag_names = ["trl", "reward-trainer"] - _name = "Reward" - _template_file = "rm_model_card.md" - - def __init__( - self, - model: "str | PreTrainedModel | PeftModel", - args: RewardConfig | None = None, - data_collator: DataCollator | None = None, - train_dataset: Dataset | IterableDataset | None = None, - eval_dataset: Dataset - | IterableDataset - | DatasetDict - | IterableDatasetDict - | dict[str, Dataset | IterableDataset] - | None = None, - processing_class: PreTrainedTokenizerBase | None = None, - compute_metrics: Callable[[EvalPrediction], dict] | None = None, - callbacks: list[TrainerCallback] | None = None, - optimizers: tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None] = (None, None), - optimizer_cls_and_kwargs: tuple[type[torch.optim.Optimizer], dict[str, Any]] | None = None, - preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, - quantization_config: "BitsAndBytesConfig | None" = None, - peft_config: "PeftConfig | None" = None, - ): - # Args - if args is None: - model_name = model if isinstance(model, str) else get_config_model_id(model.config) - model_name = model_name.split("/")[-1] - args = RewardConfig(f"{model_name}-Reward") - - if train_dataset is None: - raise ValueError("`train_dataset` is required") - elif isinstance(train_dataset, IterableDataset): - # IterableDataset requires dispatch_batches=False because Accelerate's dispatch mode may try to concatenate - # batches from multiple processes, leading to mismatch errors. - if args.accelerator_config.dispatch_batches is True: - logger.warning( - "You are using an `IterableDataset` for training with `dispatch_batches=True`. `dispatch_batches` " - "is forced to `False` when using an `IterableDataset`. To remove this warning, unset " - "`dispatch_batches` in `RewardConfig` or set it to `False`." - ) - args.accelerator_config.dispatch_batches = False - elif not isinstance(train_dataset, Dataset): - raise TypeError( - f"`train_dataset` must be a `Dataset` or `IterableDataset`, got `{type(train_dataset).__name__}`." - ) - - # Model - # As AutoModelForSequenceClassification.from_pretrained[] will add a random head for the model, set_seed must - # be done before loading the model to ensure reproducibility. - set_seed(args.seed) - if isinstance(model, str): - model_init_kwargs = dict(args.model_init_kwargs or {}) # copy to avoid mutating model_init_kwargs - if quantization_config is not None: - if "quantization_config" in model_init_kwargs: - raise ValueError( - "You set `quantization_config` both as a trainer argument and in `args.model_init_kwargs`. " - "Please set it in only one place, preferably as a trainer argument." - ) - model_init_kwargs["quantization_config"] = quantization_config - # Distributed training requires device_map=None ["auto" fails] - if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: - model_init_kwargs["device_map"] = None - model_init_kwargs["num_labels"] = 1 # the only output of the model is the reward score - model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) - model_revision = model_init_kwargs.get("revision") - with suppress_seqcls_warning(): - model = create_model_from_path(model, AutoModelForSequenceClassification, **model_init_kwargs) - else: - model_revision = None - if args.model_init_kwargs is not None: - logger.warning( - "You passed `model_init_kwargs` to the `RewardConfig`, but your model is already instantiated. " - "The `model_init_kwargs` will be ignored." - ) - if quantization_config is not None: - logger.warning( - "You passed `quantization_config` to the trainer, but your model is already instantiated. The " - "`quantization_config` will be ignored." - ) - # Validate that the model has num_labels = 1 [required for reward models] - if getattr(model.config, "num_labels", None) != 1: - raise ValueError( - f"The model has `num_labels={model.config.num_labels}`, but reward models require `num_labels=1` " - "to output a single scalar reward per sequence. Please instantiate your model with `num_labels=1` " - "or pass a model name as a string to have it configured automatically." - ) - # Non-quantized models do not have the `is_loaded_in_{8,4}bit` attributes, whereas quantized models do - _is_quantized_model = getattr(model, "is_loaded_in_4bit", False) or getattr(model, "is_loaded_in_8bit", False) - - # Processing class - if processing_class is None: - processing_class = AutoTokenizer.from_pretrained( - get_config_model_id(model.config), revision=model_revision, trust_remote_code=args.trust_remote_code - ) - - # Handle pad token for processors or tokenizers - if args.eos_token is not None: - if args.eos_token not in processing_class.get_vocab(): - raise ValueError( - f"The specified `eos_token` ('{args.eos_token}') is not found in the vocabulary of the given " - f"`processing_class` ({processing_class.__class__.__name__}). Ensure that the `eos_token` exists " - "in the vocabulary before using it as an EOS token." - ) - processing_class.eos_token = args.eos_token - # The model must agree with the tokenizer on the eos token from construction, so mirror it onto the model - # config [a sequence classification model has no generation config]. - model.config.get_text_config().eos_token_id = processing_class.eos_token_id - - if args.chat_template_path is not None: - if os.path.isfile(args.chat_template_path) and args.chat_template_path.endswith((".jinja", ".j2")): - with open(args.chat_template_path, encoding="utf-8") as chat_template_file: - processing_class.chat_template = chat_template_file.read() - added_tokens = [] - else: - model, processing_class, added_tokens = clone_chat_template( - model, processing_class, args.chat_template_path - ) - else: - added_tokens = [] - - # PEFT - if False: - if not is_peft_available(): - raise ImportError( - "You passed `peft_config` but the `peft` library is not installed. " - "Install it with `pip install trl[peft]`." - ) - if not isinstance(peft_config, PeftConfig): - raise TypeError( - f"`peft_config` must be a `peft.PeftConfig` instance (e.g. `peft.LoraConfig`), " - f"got {type(peft_config).__name__}." - ) - if is_peft_model(model): - raise ValueError( - "You passed a `PeftModel` instance together with a `peft_config` to the trainer. Please first merge " - "and unload the existing adapter, save the resulting base model, and then pass that base model along " - "with the new `peft_config` to the trainer." - ) - if added_tokens: - # Ensure that the added tokens are trainable - if peft_config.trainable_token_indices is None: - peft_config.trainable_token_indices = {"embed_tokens": added_tokens} - elif "embed_tokens" not in peft_config.trainable_token_indices: - peft_config.trainable_token_indices["embed_tokens"] = added_tokens - else: - peft_config.trainable_token_indices["embed_tokens"].extend(added_tokens) - # Ensure that the lm_head is trainable - if peft_config.modules_to_save is None or "lm_head" not in peft_config.modules_to_save: - logger.warning( - "Cloning chat template added new tokens to the tokenizer, but 'lm_head' is not in PEFT's " - "`modules_to_save`. As a result, the model may not learn to generate outputs with these new " - "tokens, leading to degraded generation quality. To fix this, add " - "`modules_to_save=['lm_head']` to your PEFT configuration." - ) - - if peft_config.modules_to_save is None: - peft_config.modules_to_save = ["lm_head"] - else: - peft_config.modules_to_save.append("lm_head") - # Create PEFT model - # ZeRO-3 + PEFT for non-quantized models: - # - PEFT's default autocast_adapter_dtype=True upcasts LoRA adapter params to fp32 even when the base model is bf16. - # - ZeRO-3's _allgather_params_coalesced allocates output buffers using the dtype of the first persistent parameter, - # so mixed-dtype persistent_parameters [bf16 base + fp32 LoRA] cause a TypeError on the first optimizer step. - # - Passing autocast_adapter_dtype=False keeps adapter params in the base model dtype [bf16], fixing the mismatch. - # - This is safe: the fp32 upcast is a QLoRA-specific concern [low-bit quantized base models], not needed for - # non-quantized bf16 training. - # - See: - # - TRL issue: https://github.com/huggingface/trl/issues/6089 - # - Upstream issue: https://github.com/deepspeedai/DeepSpeed/issues/8072 - get_peft_model_kwargs = {} - if args.deepspeed_plugin is not None and args.deepspeed_plugin.zero_stage == 3 and not _is_quantized_model: - get_peft_model_kwargs["autocast_adapter_dtype"] = False - model = get_peft_model(model, peft_config, **get_peft_model_kwargs) - - # PEFT + DeepSpeed ZeRO-3 requires reentrant checkpointing. For more details, see - # https://github.com/huggingface/trl/issues/2514#issuecomment-2692152703. - # Can be removed once https://github.com/deepspeedai/DeepSpeed/pull/8130 is merged and released. - if ( - is_peft_model(model) - and args.deepspeed_plugin is not None - and args.deepspeed_plugin.zero_stage == 3 - and args.gradient_checkpointing - ): - args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} - use_reentrant = args.gradient_checkpointing_kwargs.get("use_reentrant") - if use_reentrant is False: - logger.warning( - "You are using PEFT with DeepSpeed ZeRO-3 and gradient checkpointing with `use_reentrant=False`. " - "`use_reentrant` is forced to `True` in this configuration to ensure correct training. To remove " - "this warning, unset `use_reentrant` in `gradient_checkpointing_kwargs` or set it to `True`." - ) - args.gradient_checkpointing_kwargs["use_reentrant"] = True - - # When using gradient checkpointing with PEFT, we need to enable input gradients. transformers.Trainer normally - # handles this, but a bug currently prevents it; see https://github.com/huggingface/transformers/issues/42489 - if is_peft_model(model) and args.gradient_checkpointing: - model.enable_input_require_grads() - - # When using QLoRA, the PEFT adapter weights are converted to bf16 to follow the recommendations from the - # original paper [see https://huggingface.co/papers/2305.14314, paragraph 3]. Normally, this can be done by - # passing `autocast_adapter_dtype=False` to `get_peft_model`, but this option is not yet supported for - # quantized models. See: https://github.com/huggingface/peft/issues/2889 - # The DoRA magnitude vector is excluded: unlike LoRA A/B, its optimizer updates can be smaller than bf16 can - # represent, silently freezing it, see #7268. - if False: - for name, param in model.named_parameters(): - if param.requires_grad and "lora_magnitude_vector" not in name: - param.data = param.data.to(torch.bfloat16) - - # Disable dropout in the model - if args.disable_dropout: - disable_dropout_in_model(model) - - # Pad token [needed for SequenceClassification models] - # If not provided, use the one from the processing class or the eos token if the processing class does not have - # a pad token. - pad_token = args.pad_token or processing_class.pad_token or processing_class.eos_token - if pad_token not in processing_class.get_vocab(): - raise ValueError( - f"The specified `pad_token` ('{pad_token}') is not found in the vocabulary of the given " - f"`processing_class` ({processing_class.__class__.__name__}). Ensure that the `pad_token` exists " - "in the vocabulary before using it as a padding token." - ) - processing_class.pad_token = pad_token - # SequenceClassification models need `config.pad_token_id` to locate the last non-pad token. - model.config.get_text_config().pad_token_id = processing_class.pad_token_id - - # Data collator - if data_collator is None: - data_collator = DataCollatorForPreference( - pad_token_id=processing_class.pad_token_id, - pad_to_multiple_of=args.pad_to_multiple_of, - ) - - # Dataset - train_dataset = self._prepare_dataset(train_dataset, processing_class, args, "train") - if eval_dataset is not None: - if isinstance(eval_dataset, dict): - eval_dataset = { - key: self._prepare_dataset(dataset, processing_class, args, key) - for key, dataset in eval_dataset.items() - } - else: - eval_dataset = self._prepare_dataset(eval_dataset, processing_class, args, "eval") - - # Transformers explicitly set use_reentrant=True in the past to silence a PyTorch warning, but the default was - # never updated once PyTorch switched to recommending use_reentrant=False. Until that change lands upstream - # [see https://github.com/huggingface/transformers/pull/43203] and is released [most likely in 5.0.0], we - # default to the recommended non-reentrant behavior here, while preserving any user-provided value. - if args.gradient_checkpointing and Version(transformers.__version__) < Version("5.0.0"): - args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} - args.gradient_checkpointing_kwargs.setdefault("use_reentrant", False) - - super().__init__( - model=model, - args=args, - data_collator=data_collator, - train_dataset=train_dataset, - eval_dataset=eval_dataset, - processing_class=processing_class, - compute_metrics=compute_metrics, - callbacks=callbacks, - optimizers=optimizers, - optimizer_cls_and_kwargs=optimizer_cls_and_kwargs, - preprocess_logits_for_metrics=preprocess_logits_for_metrics, - ) - - # During evaluation, Trainer calls compute_loss[] only if can_return_loss is True and label_names is empty. - self.can_return_loss = True - self.label_names = [] - - # Initialize activation offloading context - if self.args.activation_offloading: - self.maybe_activation_offload_context = get_act_offloading_ctx_manager(model=self.model) - else: - self.maybe_activation_offload_context = contextlib.nullcontext() - - # Initialize the metrics - self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} - self._total_train_tokens = 0 - - # Tensor parallel ranks are given the same batch, so a token count gathered across all processes repeats - # every token once per rank in the group. Context and sequence parallelism shard the batch before the loss is - # computed, so they need no such correction. `parallelism_config` requires accelerate 1.12.0. - if Version(accelerate.__version__) >= Version("1.12.0") and self.accelerator.parallelism_config is not None: - self._tp_size = self.accelerator.parallelism_config.tp_size - else: - self._tp_size = 1 - - # Add tags to the model - self.model.add_model_tags(self._tag_names) - - def _prepare_dataset( - self, - dataset: Dataset | IterableDataset, - processing_class: PreTrainedTokenizerBase, - args: RewardConfig, - dataset_name: str, - ) -> Dataset | IterableDataset: - # If the dataset is already preprocessed (tokenized), skip the processing steps. - column_names = get_dataset_column_names(dataset) - is_processed = "chosen_ids" in column_names and "rejected_ids" in column_names - has_legacy_processed_columns = "chosen_input_ids" in column_names and "rejected_input_ids" in column_names - if has_legacy_processed_columns and not is_processed: - warnings.warn( - "Detected legacy dataset columns `chosen_input_ids`/`rejected_input_ids`; they are deprecated and " - "will not be supported in v1. Please migrate to `chosen_ids`/`rejected_ids`.", - FutureWarning, - stacklevel=2, - ) - dataset = dataset.rename_column("chosen_input_ids", "chosen_ids") - dataset = dataset.rename_column("rejected_input_ids", "rejected_ids") - is_processed = True - - # Build the kwargs for the `map` function - map_kwargs = {} - if isinstance(dataset, Dataset): # IterableDataset does not support num_proc - map_kwargs["num_proc"] = args.dataset_num_proc - - with global_then_local_main_first(): - if not is_processed: - # Add EOS token if needed: non-conversational only - first_example = next(iter(dataset)) - if not is_conversational(first_example): - if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` - map_kwargs["desc"] = f"Adding EOS to {dataset_name} dataset" - - def add_eos(example, eos_token): - if not example["chosen"].endswith(eos_token): - example["chosen"] = example["chosen"] + eos_token - if "rejected" in example and not example["rejected"].endswith(eos_token): - example["rejected"] = example["rejected"] + eos_token - return example - - dataset = dataset.map( - add_eos, - fn_kwargs={"eos_token": processing_class.eos_token}, - **map_kwargs, - ) - - # Tokenize the dataset - if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` - map_kwargs["desc"] = f"Tokenizing {dataset_name} dataset" - - def tokenize_fn(example, processing_class): - tools = example.get("tools") - tools = json.loads(tools) if isinstance(tools, str) else tools - apply_chat_template_kwargs = {"tools": tools, **example.get("chat_template_kwargs", {})} - if "prompt" in example: # explicit prompt case - example["chosen"] = example["prompt"] + example["chosen"] - example["rejected"] = example["prompt"] + example["rejected"] - - if is_conversational(example): - chosen_ids = _tokenize( - processing_class, - example["chosen"], - **apply_chat_template_kwargs, - )["input_ids"] - rejected_ids = _tokenize( - processing_class, - example["rejected"], - **apply_chat_template_kwargs, - )["input_ids"] - output = {"chosen_ids": chosen_ids, "rejected_ids": rejected_ids} - else: - output = { - "chosen_ids": _tokenize(processing_class, example["chosen"])["input_ids"], - "rejected_ids": _tokenize(processing_class, example["rejected"])["input_ids"], - } - return output - - dataset = dataset.map(tokenize_fn, fn_kwargs={"processing_class": processing_class}, **map_kwargs) - - # Filter samples that are longer than `max_length` - if args.max_length is not None: - if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` - map_kwargs["desc"] = f"Filtering {dataset_name} >{args.max_length} tokens" - dataset = dataset.filter( - lambda example: ( - len(example["chosen_ids"]) <= args.max_length - and len(example["rejected_ids"]) <= args.max_length - ), - **map_kwargs, - ) - - return dataset - - def _set_signature_columns_if_needed(self): - # If `self.args.remove_unused_columns` is True, non-signature columns are removed. - # By default, this method sets `self._signature_columns` to the model's expected inputs (usually, "input_ids" - # and "attention_mask"). - if self._signature_columns is None: - self._signature_columns = ["chosen_ids", "rejected_ids", "margin"] - - def evaluate( - self, - eval_dataset: Dataset - | IterableDataset - | DatasetDict - | IterableDatasetDict - | dict[str, Dataset | IterableDataset] - | None = None, - ignore_keys: list[str] | None = None, - metric_key_prefix: str = "eval", - ) -> dict[str, float]: - # When a dataset is passed directly to `evaluate` (e.g. a held-out test set), preprocess it the same way - # `__init__` does, so that `evaluate` accepts the same dataset types as the trainer. `_prepare_dataset` is - # idempotent: it skips datasets that are already tokenized. A `str` selects a dataset that was already prepared - # at init time, so it's left untouched. - if eval_dataset is not None and not isinstance(eval_dataset, str): - if isinstance(eval_dataset, dict): - eval_dataset = { - key: self._prepare_dataset(dataset, self.processing_class, self.args, key) - for key, dataset in eval_dataset.items() - } - else: - eval_dataset = self._prepare_dataset(eval_dataset, self.processing_class, self.args, "eval") - return super().evaluate( - eval_dataset=eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix - ) - - def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): - mode = "train" if self.model.training else "eval" - - # If not set, defaults from model config and may warn since cache isn't compatible with gradient checkpointing - inputs["use_cache"] = False - outputs = model(**inputs) - - # Split the rewards into chosen and rejected - rewards_chosen, rewards_rejected = torch.chunk(outputs.logits.squeeze(-1), chunks=2) - - # Calculate loss, optionally modulate with margin - if "margin" in inputs: - loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected - inputs["margin"]).mean() - else: - loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean() - - if self.args.center_rewards_coefficient is not None: - loss += self.args.center_rewards_coefficient * torch.mean((rewards_chosen + rewards_rejected) ** 2) - - if mode == "train": - num_tokens_in_batch = self.accelerator.gather_for_metrics(inputs["attention_mask"].sum()).sum().item() - self._total_train_tokens += num_tokens_in_batch // self._tp_size - self._metrics[mode]["num_tokens"] = [self._total_train_tokens] - - # Compute min, mean, max, accuracy and margin - with torch.no_grad(): - all_rewards = self.accelerator.gather(outputs.logits) - self._metrics[mode]["min_reward"].append(all_rewards.min().item()) - self._metrics[mode]["mean_reward"].append(all_rewards.mean().item()) - self._metrics[mode]["max_reward"].append(all_rewards.max().item()) - - mean_accuracy = (rewards_chosen > rewards_rejected).float().mean() - mean_accuracy = self.accelerator.gather_for_metrics(mean_accuracy).mean().item() - self._metrics[mode]["accuracy"].append(mean_accuracy) - - mean_margin = (rewards_chosen - rewards_rejected).mean() - mean_margin = self.accelerator.gather_for_metrics(mean_margin).mean() - self._metrics[mode]["margin"].append(mean_margin.item()) - - return (loss, outputs) if return_outputs else loss - - # Override training step to add activation offloading context. - def training_step(self, *args, **kwargs): - with self.maybe_activation_offload_context: - return super().training_step(*args, **kwargs) - - def log(self, logs: dict[str, float], start_time: float | None = None) -> None: - mode = "train" if self.model.training else "eval" - metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()} # average the metrics - - # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` - # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. - if mode == "eval": - metrics = {f"eval_{key}": val for key, val in metrics.items()} - - logs.update(metrics) - super().log(logs, start_time) - self._metrics[mode].clear() - - # Ensure the model card is saved along with the checkpoint - def _save_checkpoint(self, model, trial): - if self.args.hub_model_id is None: - model_name = Path(self.args.output_dir).name - else: - model_name = self.args.hub_model_id.split("/")[-1] - self.create_model_card(model_name=model_name) - super()._save_checkpoint(model, trial) -class UnslothRewardTrainer(_UnslothRewardTrainer): - """ - - Trainer for Outcome-supervised Reward Models (ORM). - - This class is a wrapper around the [`~transformers.Trainer`] class and inherits all of its attributes and methods. - - Example: - - ```python - >>> from trl import RewardTrainer - >>> from datasets import load_dataset - - >>> dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train") - - >>> trainer = RewardTrainer( - ... model="Qwen/Qwen2.5-0.5B-Instruct", - ... train_dataset=dataset, - ... ) - >>> trainer.train() - ``` - - Args: - model (`str` or [`~transformers.PreTrainedModel`] or [`~peft.PeftModel`]): - Model to be trained. Can be either: - - - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a - path to a *directory* containing model weights saved using - [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded - using `AutoModelForSequenceClassification.from_pretrained` with the keyword arguments in - `args.model_init_kwargs`. If `dtype` is not specified in `args.model_init_kwargs`, it defaults to - `float32`. This differs from [`~transformers.PreTrainedModel.from_pretrained`], where (since Transformers - v5) the dtype is inferred from the model config. - - A sequence classification [`~transformers.PreTrainedModel`] object. - - A sequence classification [`~peft.PeftModel`] object. - args ([`RewardConfig`], *optional*): - Configuration for this trainer. If `None`, a default configuration is used. - data_collator ([`~transformers.DataCollator`], *optional*): - Function to use to form a batch from a list of elements of the processed `train_dataset` or `eval_dataset`. - Will default to [`~trainer.reward_trainer.DataCollatorForPreference`]. - train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): - Dataset to use for training. This trainer supports [preference](#preference) type (both implicit and - explicit prompt). The format of the samples can be either: - - - [Standard](dataset_formats#standard): Each sample contains plain text. - - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role - and content). - - The trainer also supports processed datasets (tokenized) as long as they contain `chosen_ids` and - `rejected_ids` fields. - - When `train_dataset` is an [`~datasets.IterableDataset`] (e.g. a streaming dataset), `max_steps` must be - set in the training arguments, since its length cannot be inferred and the total number of training steps - is required to bound the training loop and configure the learning rate scheduler. - eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`], [`~datasets.DatasetDict`], [`~datasets.IterableDatasetDict`] or `dict[str, Dataset | IterableDataset]`): - Dataset to use for evaluation. It must meet the same requirements as `train_dataset`. - processing_class ([`~transformers.PreTrainedTokenizerBase`], *optional*): - Tokenizer used to process the data. If `None`, the tokenizer is loaded from the model's name with - [`~transformers.AutoTokenizer.from_pretrained`]. A padding token, `processing_class.pad_token`, must be - set. If the processing class has not set a padding token, `processing_class.eos_token` will be used as the - default. - compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): - The function that will be used to compute metrics at evaluation. Must take a - [`~transformers.EvalPrediction`] and return a dictionary string to metric values. When passing - [`RewardConfig`] with `batch_eval_metrics` set to `True`, your `compute_metrics` function must take a - boolean `compute_result` argument. This will be triggered after the last eval batch to signal that the - function needs to calculate and return the global summary statistics rather than accumulating the - batch-level statistics. - callbacks (list of [`~transformers.TrainerCallback`], *optional*): - List of callbacks to customize the training loop. Will add those to the list of default callbacks detailed - in [here](https://huggingface.co/docs/transformers/main_classes/callback). - - If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`] - method. - optimizers (`tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None]`, *optional*, defaults to `(None, None)`): - A tuple containing the optimizer and the scheduler to use. Will default to an instance of `AdamW` on your - model and a scheduler given by [`~transformers.get_linear_schedule_with_warmup`] controlled by `args`. - optimizer_cls_and_kwargs (`tuple[Type[torch.optim.Optimizer], Dict[str, Any]]`, *optional*): - A tuple containing the optimizer class and keyword arguments to use. Overrides `optim` and `optim_args` in - `args`. Incompatible with the `optimizers` argument. - - Unlike `optimizers`, this argument avoids the need to place model parameters on the correct devices before - initializing the Trainer. - preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`, *optional*): - A function that preprocess the logits right before caching them at each evaluation step. Must take two - tensors, the logits and the labels, and return the logits once processed as desired. The modifications made - by this function will be reflected in the predictions received by `compute_metrics`. - - Note that the labels (second parameter) will be `None` if the dataset does not have them. - quantization_config ([`~transformers.BitsAndBytesConfig`], *optional*): - Quantization configuration used when loading the model from a model identifier. Combine with `peft_config` - for QLoRA training. Ignored if the model is already instantiated, or if `quantization_config` is also set - in `args.model_init_kwargs`. - peft_config ([`~peft.PeftConfig`], *optional*): - PEFT configuration used to wrap the model. If `None`, the model is not wrapped. Note that if the loaded - model is a causal LM, it's highly recommended to set `modules_to_save=["score"]` in the PEFT configuration - to ensure that the reward head is properly trained. - - """ - def __init__( - self, - model, - args = None, - data_collator = None, - train_dataset = None, - eval_dataset = None, - processing_class = None, - compute_metrics = None, - callbacks = None, - optimizer_cls_and_kwargs = None, - preprocess_logits_for_metrics = None, - quantization_config = None, - peft_config = None, - **kwargs - ): - if args is None: args = UnslothRewardConfig() - _unsloth_precision_requested = tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) - use_bf16 = getattr(args, 'bf16', False) - if type(use_bf16) is not bool: use_bf16 = False - use_fp16 = getattr(args, 'fp16', False) - if type(use_fp16) is not bool: use_fp16 = False - force_float32 = False - try: - from unsloth_zoo.device_type import device_is_bf16_supported as _bf16_supported - except Exception: - _bf16_supported = torch.cuda.is_bf16_supported - try: - from unsloth_zoo.device_type import DEVICE_TYPE as _unsloth_device_type - except Exception: - _unsloth_device_type = None - full_finetuning = getattr(model, '_unsloth_full_finetuning', None) - if full_finetuning is None: full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1' - model_forced_float32 = getattr(model, '_unsloth_forced_float32', None) - if model_forced_float32 is None: model_forced_float32 = os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1' - if model_forced_float32 and not (full_finetuning and _bf16_supported()): - print('Unsloth: Switching to float32 training since model cannot work with float16') - force_float32 = True - mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') - dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None) - if dtype is None: dtype = model.get_input_embeddings().weight.dtype - from unsloth_zoo.utils import _get_dtype - dtype = _get_dtype(dtype) - float16 = dtype == torch.float16 - bfloat16 = dtype == torch.bfloat16 - float32 = dtype == torch.float32 - user_float32 = bool(getattr(model, '_unsloth_user_float32', False)) - if full_finetuning: - if bfloat16 and use_fp16: use_fp16 = False - if float16 and use_bf16: use_bf16 = False - if not force_float32 and bfloat16 and use_fp16 and _unsloth_device_type == 'xpu' and not _bf16_supported(): - print('Unsloth: The model was loaded in bfloat16 on purpose, so training uses bfloat16 mixed precision instead of float16.') - use_fp16 = False - use_bf16 = True - args.fp16 = False - args.bf16 = True - if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`') - if not force_float32 and (bfloat16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`') - if force_float32: - # Forced float32 training - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32' and float32 and user_float32 and not _bf16_supported(): - print('Unsloth: Model is in float32 and this GPU has no bfloat16 support, so training stays in float32. Pass fp16 = True to force float16 mixed precision instead.') - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32': - # Mixed precision training. bf16 only if the GPU supports it; V100/T4 use fp16. - use_bf16_amp = (not float16) and (_bf16_supported() or (bfloat16 and _unsloth_device_type == 'xpu')) - args.fp16 = not use_bf16_amp - args.bf16 = use_bf16_amp - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16_amp else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16_amp else 'fp16' - # args.mixed_precision is a new argument which needs to be set now - elif mixed_precision_dtype == 'bfloat16': - # Both False since bfloat16 full finetuning doesn't do any autocasting. - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif use_bf16 or use_fp16: - # transformers <5 exported this itself from fp16/bf16; 5.x dropped the write, so an - # explicit flag left it unset and GRPO readers defaulted to 'fp16', wrapping a - # bfloat16 model in a float16 autocast. See unslothai/unsloth#4891. - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16 else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16 else 'fp16' - - if getattr(args, 'eval_dataset', None) is not None and getattr(args, 'eval_strategy', 'no') == 'no': - args.eval_strategy = 'steps' - if getattr(args, 'eval_steps', None) is None: args.eval_steps = 0.1 - ga_steps = getattr(args, 'gradient_accumulation_steps', None) - if ga_steps is not None and ga_steps > 1: - from transformers import __version__ as transformers_version - if Version(transformers_version) <= Version('4.45.2'): - print('**** Unsloth: Please use our fixed gradient_accumulation_steps by updating transformers, TRL and Unsloth!\n' - '`pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth transformers trl unsloth_zoo`') - if getattr(args, 'eval_strategy', 'no') != 'no': - eval_bsz = getattr(args, 'per_device_eval_batch_size', 8) - eval_generations = getattr(args, 'num_generations_eval', None) or getattr(args, 'num_generations', None) or 1 - if eval_bsz == 8 and args.per_device_train_batch_size < eval_bsz and (args.per_device_train_batch_size * args.world_size) % eval_generations == 0: args.per_device_eval_batch_size = args.per_device_train_batch_size - if getattr(args, 'eval_accumulation_steps', None) is None and ga_steps is not None: args.eval_accumulation_steps = ga_steps - fp16_full_eval = getattr(args, 'fp16_full_eval', False) - if type(fp16_full_eval) is not bool: fp16_full_eval = False - bf16_full_eval = getattr(args, 'bf16_full_eval', False) - if type(bf16_full_eval) is not bool: bf16_full_eval = False - if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True - if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False - if force_float32: - args.bf16_full_eval = False - args.fp16_full_eval = False - elif os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16': - args.bf16_full_eval = True - args.fp16_full_eval = False - elif not bf16_full_eval and not fp16_full_eval: - args.bf16_full_eval = args.bf16 - args.fp16_full_eval = args.fp16 - if tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) != _unsloth_precision_requested: - from unsloth.models.loader_utils import sync_deepspeed_mixed_precision - sync_deepspeed_mixed_precision(args) - if model is not None: - from unsloth.models.loader_utils import apply_deepspeed_zero3_model_setup - apply_deepspeed_zero3_model_setup(model) - if model is not None: - _warnings_issued = getattr(model, 'warnings_issued', None) - if _warnings_issued is None: - model.warnings_issued = {} - elif not isinstance(_warnings_issued, dict): - try: - model.warnings_issued = dict(_warnings_issued) - except Exception: - model.warnings_issued = {} - if 'max_seq_length' not in locals() and not hasattr(args, 'max_seq_length'): - pass - else: - model_max_seq_length = getattr(model, 'max_seq_length', None) - args_max_seq_length = getattr(args, 'max_seq_length', None) - if args_max_seq_length is None and model_max_seq_length is not None: - max_seq_length = model.max_seq_length - if hasattr(args, 'max_seq_length'): args.max_seq_length = max_seq_length - elif args_max_seq_length is not None and model_max_seq_length is not None: - if args_max_seq_length > model_max_seq_length: - print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but ' - 'the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.') - args.max_seq_length = model_max_seq_length - if model is not None and hasattr(model, 'for_training'): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right' - if 'processing_class' in locals(): - if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right' - if hasattr(processing_class, 'tokenizer') and hasattr(processing_class.tokenizer, 'padding_side'): processing_class.tokenizer.padding_side = 'right' - __tokenizer = processing_class if 'processing_class' in locals() else tokenizer - from unsloth_zoo.vision_utils import UnslothVisionDataCollator - if not isinstance(data_collator, UnslothVisionDataCollator): - if isinstance(data_collator, DataCollatorForSeq2Seq) and 'labels' not in _unsloth_dataset_column_names(train_dataset): - data_collator = TransformersDataCollatorForLanguageModeling( - __tokenizer, - mlm = False, - mlm_probability = 0.0, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling) and 'labels' in _unsloth_dataset_column_names(train_dataset): - data_collator = DataCollatorForSeq2Seq( - __tokenizer, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - else: - if hasattr(args, 'remove_unused_columns'): args.remove_unused_columns = False - if hasattr(args, 'dataset_text_field'): args.dataset_text_field = '' - if hasattr(args, 'dataset_kwargs'): args.dataset_kwargs = {'skip_prepare_dataset': True} - if not isinstance(data_collator, UnslothVisionDataCollator): - if not hasattr(__tokenizer, 'pad') and hasattr(__tokenizer, 'tokenizer'): - if isinstance(data_collator, DataCollatorForSeq2Seq): - data_collator = DataCollatorForSeq2Seq( - __tokenizer.tokenizer, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling): - data_collator = TransformersDataCollatorForLanguageModeling( - __tokenizer.tokenizer, - mlm = False, - mlm_probability = 0.0, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - other_metrics = [] - - from unsloth_zoo.logging_utils import PatchRLStatistics - PatchRLStatistics('reward_trainer', other_metrics) - - # [TODO] Fix up DataParallel multiplying batch sizes - # [TODO] DDP works, but DP seems to not work? [TODO] - if getattr(args, "parallel_mode", None) == ParallelMode.NOT_DISTRIBUTED and args.n_gpu > 1: - if getattr(args, "_n_gpu", 1) != 1: - args._n_gpu = 1 - if "model" in locals(): - try: - from unsloth.models._gc_config_default import honour_gradient_checkpointing_off - honour_gradient_checkpointing_off(args, model) - except Exception: - pass - if "model" in locals() and hasattr(model, "for_training"): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if "model" in locals(): - try: - from unsloth.models._ddp_defaults import unsloth_default_ddp_find_unused_parameters - unsloth_default_ddp_find_unused_parameters(args, model) - except Exception: - pass - super().__init__( - model = model, - args = args, - data_collator = data_collator, - train_dataset = train_dataset, - eval_dataset = eval_dataset, - processing_class = processing_class, - compute_metrics = compute_metrics, - callbacks = callbacks, - optimizer_cls_and_kwargs = optimizer_cls_and_kwargs, - preprocess_logits_for_metrics = preprocess_logits_for_metrics, - quantization_config = quantization_config, - peft_config = peft_config,**kwargs) - # A tensor parallel model (tp_native.tp_trainer_setup) stays on its training path: for_inference would - # route the decoder layers onto decode paths outside the tensor parallel regions. - if "model" in locals() and hasattr(model, "for_inference") and not getattr(self, "_unsloth_tensor_parallel", False): - model.for_inference() - if hasattr(self, 'neftune_hook_handle'): - self.neftune_hook_handle.remove() - if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle - if getattr(args, 'neftune_noise_alpha', None) is not None: - model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha - pass - if hasattr(self, 'accelerator'): - scaler = self.accelerator.scaler - current_model = model - while hasattr(current_model, 'model'): - current_model.accelerator_scaler = scaler - current_model = current_model.model - current_model.accelerator_scaler = scaler - pass - if hasattr(self, 'train'): - self.train = MethodType(prepare_for_training_mode(self.__class__.train), self) - pass - if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'): - _vllm_tok = self.llm.get_tokenizer() - _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None) - if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None: - _vllm_tok.chat_template = _pc.chat_template - pass - if getattr(self, 'aux_loss_enabled', False) and hasattr(getattr(self, 'model', None), 'config'): - _text_config = self.model.config - if hasattr(_text_config, 'get_text_config'): _text_config = _text_config.get_text_config() - _n_experts = [getattr(_text_config, _k) for _k in ('num_local_experts', 'num_experts', 'n_routed_experts', 'moe_num_experts') if isinstance(getattr(_text_config, _k, None), int)] - if _n_experts and all(_n <= 1 for _n in _n_experts): - self.aux_loss_enabled = False - _text_config.output_router_logits = False - pass - -pass - - -if hasattr(logger, "addFilter"): - import logging - class HideLoggingMessage(logging.Filter): - def __init__(self, text): self.text = text - def filter(self, x): return not (self.text in x.getMessage()) - pass - logger.addFilter(HideLoggingMessage("`use_cache=True`")) - diff --git a/unsloth_compiled_cache/UnslothSDFTTrainer.py b/unsloth_compiled_cache/UnslothSDFTTrainer.py deleted file mode 100644 index e8e8de444358f8936b767bcf28f6b406190769fe..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/UnslothSDFTTrainer.py +++ /dev/null @@ -1,2286 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from unsloth_zoo.temporary_patches.common import torch_compile -from unsloth_zoo.temporary_patches.common import _maybe_compile -import functools -import transformers -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from trl.experimental.sdft.sdft_trainer import (Any, AutoProcessor, DataLoader, Dataset, DemonstrationTeacherContextBuilder, DistillationLogits, FSDP, GenerationConfig, IterableDataset, PEFTAdapterEMACallback, PeftConfig, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, RepeatSampler, SDFTConfig, SDFTTrainer, Sampler, SyncTeacherModelCallback, TrainerCallback, TrainingBatch, _BaseTrainer, add_tail_bucket, apply_importance_sampling_clipping, build_teacher_request_inputs, compute_divergence, compute_full_logit_self_distillation_loss, compute_sampled_token_self_distillation_loss, compute_topk_self_distillation_loss, create_model_from_path, dataclass, datasets, defaultdict, disable_dropout_in_model, get_config_model_id, identity, inspect, is_conversational, is_datasets_available, is_peft_available, is_peft_model, is_pure_lora_training, logger, nn, nullcontext, pad, partial, prepare_deepspeed, prepare_fsdp, seed_worker, selective_log_softmax, set_seed, split_tensor_dict, textwrap, torch, unwrap_model_for_generation, use_adapter, AutoProcessor, Dataset, DemonstrationTeacherContextBuilder, FSDP, GenerationConfig, IterableDataset, PeftConfig, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, SDFTConfig, SDFTTrainer, TrainerCallback, create_model_from_path, datasets, defaultdict, disable_dropout_in_model, get_config_model_id, identity, inspect, is_peft_available, is_peft_model, logger, nn, pad, set_seed, torch) - - -import os -import math -import logging -from typing import * -from dataclasses import dataclass, field -from packaging.version import Version -import torch -import numpy as np -from contextlib import nullcontext -from torch.nn import functional as F -import inspect -from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling -from transformers.training_args import ParallelMode -from unsloth_zoo.device_type import DEVICE_TYPE, DEVICE_TYPE_TORCH, device_synchronize - -# Wrap trainer with padding to right and enable training mode -import functools -from types import MethodType -try: - from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers -except: - def reset_unsloth_gradient_checkpointing_buffers(): pass -# Canonical reset lives in unsloth.models._utils so the SFT auto-packing wrapper and the plain -# Trainer loop can import the same helper; fall back to a no-op only if it can't be imported. -try: - from unsloth.models._utils import _unsloth_reset_stray_compile_cache -except Exception: - def _unsloth_reset_stray_compile_cache(self): pass -try: - from unsloth.models._utils import _unsloth_dataset_column_names -except Exception: - def _unsloth_dataset_column_names(dataset): return dataset.column_names -# Drops/renames config arguments the installed TRL no longer accepts, so a -# script pinned to an older TRL keeps working after an upgrade. Falls back to -# the historical raw passthrough so this can never break trainer construction. -try: - from unsloth.models.rl_config_compat import filter_config_init_kwargs as _unsloth_filter_config_init_kwargs - # A cache file generated here can be imported by an older Unsloth whose filter - # predates `mirrored_from`, so drop the argument rather than raise TypeError. - if "mirrored_from" not in inspect.signature(_unsloth_filter_config_init_kwargs).parameters: - _unsloth_filter_config_init_kwargs_old = _unsloth_filter_config_init_kwargs - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): - return _unsloth_filter_config_init_kwargs_old(config_class, kwargs) -except Exception: - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): return kwargs -def prepare_for_training_mode(f): - @functools.wraps(f) - def wrapper(self, *args, **kwargs): - # Drop any torch.compile graph cache poisoned by a stray pre-train forward. - try: - _unsloth_reset_stray_compile_cache(self) - except Exception: - pass - # Finish the previous W&B run if this is a subsequent train() call. - # We do this at the START of train() (not the end) so that - # evaluate() / log() still work after train() completes. - # HF's WandbCallback.setup() will call wandb.init() for the new run. - # See: https://github.com/unslothai/unsloth/issues/3954 - if getattr(self, '_unsloth_training_completed', False): - try: - import wandb - if wandb.run is not None: - wandb.finish() - # Reset HF's WandbCallback so it calls wandb.init() for the new run - for cb in self.callback_handler.callbacks: - if type(cb).__name__ == 'WandbCallback': - cb._initialized = False - break - except: - pass - # Enable training mode - _was_training = None - # Restore the GC mode the model was configured with at setup; fall back to - # the training args only when it wasn't recorded (issue #4735). Use hasattr, - # not a None sentinel, so a deliberately-recorded None is restored verbatim. - _model = getattr(self, 'model', None) - if hasattr(_model, '_unsloth_gradient_checkpointing'): - use_gc = _model._unsloth_gradient_checkpointing - else: - use_gc = getattr(self.args, 'gradient_checkpointing', True) - if hasattr(self, 'model') and hasattr(self.model, "training"): - _was_training = self.model.training - if hasattr(self, 'model') and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - output = f(self, *args, **kwargs) - # Restore previous mode when possible - if hasattr(self, 'model') and hasattr(self.model, "for_inference"): - if _was_training is False: - self.model.for_inference() - elif _was_training is True and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - # Reset gradient checkpointing buffers to free memory while staying ready for next run - try: - reset_unsloth_gradient_checkpointing_buffers() - except: - pass - # Mark that training completed so the next train() call can - # finish this W&B run before starting a new one - self._unsloth_training_completed = True - return output - return wrapper -pass - -torch_compile_options = { - "epilogue_fusion" : True, - "max_autotune" : False, - "shape_padding" : True, - "trace.enabled" : False, - "triton.cudagraphs" : False, -} - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_hidden_states_selective_log_softmax( - hidden_states: torch.Tensor, - lm_head: torch.Tensor, - index: torch.Tensor, - chunks: int = 4, - logit_scale_multiply: float = 0.0, - logit_scale_divide: float = 0.0, - logit_softcapping: float = 0.0, - temperature: float = 1.0, - # Rows per chunk cap. Read HERE, in the default, not in the body: the body - # is traced with fullgraph = True, and `os.environ` is an unsupported op - # there on torch 2.4 -- Dynamo raises - # torch._dynamo.exc.Unsupported: const method call bytes.decode - # from os._Environ.__getitem__, which the eager fallback does not catch - # (it only catches recompile-limit and disabled-hook breaks), so the very - # first call would die even with the variable unset. A default is evaluated - # once when the def runs, which is import time, outside any traced region. - # It is also a plain int argument, so Dynamo guards on it instead of - # constant-folding an unguarded read (2.7+ never notice a later change). - # A non-numeric value is ignored rather than raised on, so a typo cannot - # break the import. 0 keeps the previous chunk boundaries exactly. - max_rows_per_chunk: int = ( - int(os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip()) - if os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip().isdigit() - else 0 - ), - lm_head_bias: Optional[torch.Tensor] = None, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, - # Read in the default for the same Dynamo reason as max_rows_per_chunk. Pairs with vLLM's - # head_dtype=float32 that load_vllm requests under the same flag. - fp32_lm_head: bool = ( - os.environ.get("UNSLOTH_VLLM_FP32_LM_HEAD", "0").strip().lower() in ("1", "true", "on", "yes") - ), - # Older Unsloth scores old / ref log-probs here and the policy with the fused kernel: one scorer keeps ratio at 1. - fused_linear: bool = ( - os.environ.get("UNSLOTH_FUSED_LINEAR", "1") != "0" - and os.environ.get("UNSLOTH_GRPO_CHUNKED_LOGPROBS", "1") != "0" - ), -) -> torch.Tensor: - # All Unsloth Zoo code licensed under AGPL3 - from unsloth_zoo.fused_losses import chunked_logprob as _chunked_logprob - if (fused_linear and _chunked_logprob._FUSED_LINEAR and not fp32_lm_head - and support_top_k == 0 and support_top_p == 1.0 and support_min_p == 0.0 - and support_excluded_ids is None): - return _chunked_logprob.unsloth_chunked_logprobs( - hidden_states, lm_head, index, - temperature = temperature, - logit_scale_multiply = logit_scale_multiply, - logit_scale_divide = logit_scale_divide, - logit_softcapping = logit_softcapping, - lm_head_bias = lm_head_bias, - max_rows_per_chunk = max_rows_per_chunk if max_rows_per_chunk > 0 else None, - ).reshape((hidden_states.shape[0], hidden_states.shape[1])) - # Reshape on this tensor's own last dim: a no-op, so a wrong-width caller - # cannot have its row count silently rewritten and instead fails at the - # matmul below, which prints both operands. Do not swap in a bare - # torch._check: it reports only "Expected cond to be True", naming neither - # operand, and Dynamo rejects a message-carrying one. Callers dispatch on - # the width first -- see `compute_logprobs_chunk`, the packed path and - # `_pg_grad_forward`. - flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) - flat_index = index.reshape(-1) - - # Each chunk materialises rows x vocab logits and then a float32 copy of - # them, all on the device holding the output head. With a large vocabulary - # and a fixed chunk count that grows with the batch, so the peak scales with - # the batch rather than staying bounded. max_rows_per_chunk caps the rows - # per chunk instead, which is pure loop splitting: more, smaller chunks, - # same concatenated result. 0 (the default) keeps the previous chunk - # boundaries exactly. - if max_rows_per_chunk > 0: - n_rows = flat_hidden_states.shape[0] - chunks = max(chunks, -(-n_rows // max_rows_per_chunk)) - chunks = min(chunks, max(n_rows, 1)) - - # A one-row chunk's head gradient is a K=1 matmul Inductor mis-lowers under dynamic shapes (illegal - # memory access or NaN head gradients, torch 2.13), so keep chunks at 2+ rows when the head trains. - # Exactly `chunks` row ranges, so the unrolled loop length never depends on T: torch.chunk - # returns ceil(T / ceil(T / chunks)) pieces and recompiled per new count; tensor_split guards - # on T % chunks. At most rows // 2 chunks when the head trains keeps every piece at 2+ rows. - lm_head_grad = torch.is_grad_enabled() and lm_head.requires_grad - if lm_head_grad and flat_hidden_states.shape[0] < 2 * chunks: - chunks = max(flat_hidden_states.shape[0] // 2, 1) - - n_rows = flat_hidden_states.shape[0] - chunks = max(1, min(chunks, n_rows)) - bounds = [(i * n_rows) // chunks for i in range(chunks + 1)] - chunked_hidden_states = [flat_hidden_states[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - - # Cast once, outside the loop, so every chunk shares one fp32 head copy (also for autograd). - if fp32_lm_head: - lm_head = lm_head.to(torch.float32) - if lm_head_bias is not None: - lm_head_bias = lm_head_bias.to(torch.float32) - - all_per_token_logps = [] - - for chunk_hidden_states, chunk_index in zip(chunked_hidden_states, chunked_index): - # When the model is dispatched over several devices, the output head can - # sit on a different one from the hidden states, because accelerate - # places the tail of the model on the last device it fills. Co-locate on - # the head's device before the matmul, otherwise this raises - # Unhandled FakeTensor Device Propagation for aten.mm.default, - # found two different devices cuda:0, cuda:1 - # On a single device every .to() here is a no-op and the result is - # bit-identical to before. - chunk_hidden_states = chunk_hidden_states.to(device = lm_head.device, dtype = lm_head.dtype) - chunk_index = chunk_index.to(lm_head.device) - if fp32_lm_head: - # The GRPO loss runs this under bf16 / fp16 autocast, which would downcast the matmul. - with torch.autocast(device_type = chunk_hidden_states.device.type, enabled = False): - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - elif lm_head_bias is None: - chunk_logits = chunk_hidden_states @ lm_head.t() - else: - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - - if logit_scale_multiply != 0.0: - chunk_logits = chunk_logits * logit_scale_multiply - if logit_scale_divide != 0.0: - chunk_logits = chunk_logits / logit_scale_divide - if logit_softcapping != 0.0: - chunk_logits = logit_softcapping * torch.tanh(chunk_logits / logit_softcapping) - - chunk_logits = chunk_logits.to(torch.float32) - - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - per_token_logps = _truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - ) - else: - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - # Return to the caller's device so the concatenation below and every - # downstream consumer see the device they started on. - all_per_token_logps.append(per_token_logps.to(hidden_states.device)) - - all_per_token_logps = torch.concat(all_per_token_logps) - - all_per_token_logps = all_per_token_logps.reshape((hidden_states.shape[0], hidden_states.shape[1])) - return all_per_token_logps - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_selective_log_softmax( - logits, - index, - temperature: float = 1.0, - chunks: int = 4, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, -): - # Fixed `chunks` row ranges; see chunked_hidden_states_selective_log_softmax. - flat_logits = logits.reshape(-1, logits.shape[-1]) - flat_index = index.reshape(-1) - # Fewer rows than chunks: one row per range, as torch.chunk did, never empty ranges. - chunks = max(1, min(chunks, flat_logits.shape[0])) - bounds = [(i * flat_logits.shape[0]) // chunks for i in range(chunks + 1)] - chunked_logits = [flat_logits[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - all_per_token_logps = [] - # Per-chunk selective_log_softmax. - for chunk_logits, chunk_index in zip(chunked_logits, chunked_index): - chunk_logits = chunk_logits.to(torch.float32) - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - all_per_token_logps.append(_truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - )) - continue - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - all_per_token_logps.append(per_token_logps) - pass - all_per_token_logps = torch.concat(all_per_token_logps) - all_per_token_logps = all_per_token_logps.reshape((logits.shape[0], logits.shape[1])) - return all_per_token_logps - -def calculate_pad_tokens_in_prompt( - input_ids: torch.Tensor, - logits_to_keep: int, - pad_token_id: int -) -> torch.Tensor: - """Count left-padded tokens per sequence, e.g. [pad, pad, pad, cat] -> 3.""" - if logits_to_keep >= input_ids.shape[1]: - raise ValueError("logits_to_keep must be smaller than the sequence length.") - - prompt_section = input_ids[:, :-logits_to_keep] - - padding_mask = (prompt_section == pad_token_id) - - pad_token_counts = padding_mask.sum(dim=1) - - return pad_token_counts - -def create_completion_attention_mask( - completion_input_ids: torch.Tensor, - left_pad_tokens_per_prompt: torch.Tensor, - max_left_pad: int, - pad_token_id: int -) -> torch.Tensor: - """Build a completion mask that zeros leading prompt and trailing pad tokens. - - For [p,p,p,c,c,c,pad,pad,pad] (p=sliced prompt, c=completion, pad=padding) - this returns [0,0,0,1,1,1,0,0,0]. - """ - batch_size, completion_len = completion_input_ids.shape - device = completion_input_ids.device - - num_tokens_to_mask = max_left_pad - left_pad_tokens_per_prompt - - indices = torch.arange(completion_len, device=device).unsqueeze(0) - shift_mask = indices >= num_tokens_to_mask.unsqueeze(1) - - non_padding_mask = (completion_input_ids != pad_token_id) - - final_mask = shift_mask & non_padding_mask - - return final_mask - -def left_pack_padding(tensor: torch.Tensor, pad_id: int) -> torch.Tensor: - """Move all padding tokens in each sequence to the right.""" - mask = (tensor != pad_id) - # stable=True since the binary mask is unordered. - sorted_indices = torch.argsort(mask, dim=1, descending=True, stable=True) - packed_tensor = torch.gather(tensor, 1, sorted_indices) - return packed_tensor - -def align_logprobs_with_mask( - logprob_tensor: torch.Tensor, - attention_mask: torch.Tensor, - pad_value: float = 0.0 -) -> torch.Tensor: - """Align a log probability tensor with a given attention mask.""" - - device = logprob_tensor.device - batch_size, logprob_seq_len = logprob_tensor.shape - mask_seq_len = attention_mask.shape[1] - - padded_logprobs = torch.full( - attention_mask.shape, - fill_value=pad_value, - dtype=logprob_tensor.dtype, - device=device - ) - - left_pad_counts = torch.argmax(attention_mask, dim=1) - - cols = torch.arange(logprob_seq_len, device=device) - dest_indices = left_pad_counts.unsqueeze(1) + cols - - # Destination row indices, shape [batch_size, logprob_seq_len]. - row_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand_as(dest_indices) - - # Keep only in-bounds destinations, then scatter via advanced indexing. - valid_mask = dest_indices < mask_seq_len - valid_rows = row_indices[valid_mask] - valid_cols = dest_indices[valid_mask] - valid_vals = logprob_tensor[valid_mask] - padded_logprobs[valid_rows, valid_cols] = valid_vals - - return padded_logprobs - -def align_completion_tool_mask( - tool_mask: torch.Tensor, - completion_mask: torch.Tensor, -) -> torch.Tensor: - """Align a raw completion-length tool/env mask with Unsloth's repacked loss mask.""" - if tool_mask is None: - return completion_mask - if tool_mask.shape[0] != completion_mask.shape[0]: - raise ValueError("tool_mask batch size must match completion_mask batch size.") - - tool_mask = tool_mask.to(device=completion_mask.device) - if tool_mask.shape == completion_mask.shape: - aligned_tool_mask = tool_mask - else: - aligned_tool_mask = align_logprobs_with_mask( - tool_mask, - completion_mask, - pad_value=0, - ) - return completion_mask * aligned_tool_mask.to(dtype=completion_mask.dtype) - -def autotune_batch_and_chunks( - total_input_rows, - seq_len, - hidden_size, - vocab_size, - dtype_bytes=16, - multiplier=None -): - if multiplier is None: - final_m = max(4, seq_len // 4096) - else: - final_m = multiplier - - if torch.cuda.is_available(): - free_bytes, _ = torch.cuda.mem_get_info() - limit_gb = (free_bytes / (1024**3))*.80 - elif hasattr(torch, "xpu") and torch.xpu.is_available(): - # XPU: estimate free memory as total - reserved. - total_mem = torch.xpu.get_device_properties(0).total_memory - reserved_mem = torch.xpu.memory_reserved() - free_bytes = total_mem - reserved_mem - limit_gb = (free_bytes / (1024**3)) * 0.80 - else: - # Fallback: assume 8GB available. - limit_gb = 8.0 - - bytes_to_gb = 1024**3 - - b_vals = torch.arange(total_input_rows, 0, -1, device='cpu', dtype=torch.float32) - - hidden_gb = (b_vals * seq_len * hidden_size * dtype_bytes) / bytes_to_gb - - base_logits = ((b_vals/total_input_rows) * b_vals * seq_len * vocab_size * dtype_bytes) / bytes_to_gb - logits_gb = base_logits / final_m - - total_mem_gb = hidden_gb + logits_gb - - valid_mask = total_mem_gb <= limit_gb - valid_indices = torch.nonzero(valid_mask, as_tuple=False) - - if valid_indices.shape[0] == 0: - #This means your GPU will OOM - # Capped at the row count: unsloth's no-grad pass divides rows by this without max(1, ...). - return max(1, min(4, total_input_rows)), final_m - - best_idx = valid_indices[0].item() - final_b = int(b_vals[best_idx].item()) - - return final_b, final_m - -def sanitize_logprob(logprob): - """Local port of trl.scripts.vllm_serve.sanitize_logprob. - Filters NaN logprobs from vLLM outputs.""" - value = logprob.logprob - if math.isnan(value): - logging.getLogger(__name__).warning( - f"Generated NaN logprob, token logprob '{logprob}' will be ignored" - ) - return None - return value -@dataclass -class UnslothSDFTConfig(SDFTConfig): - """ - - Configuration class for the [`experimental.sdft.SDFTTrainer`]. - - Parameters: - > Parameters that control the SDFT loss - - distillation_alpha (`float`, *optional*, defaults to `0.5`): - Divergence interpolation coefficient for SDFT top-k logit distillation. - distillation_mode (`Literal["sampled_token", "full_logits", "topk_logits"]`, *optional*, defaults to `"topk_logits"`): - Distillation objective mode. SDFT defaults to top-k logit distillation. - distillation_topk (`int`, *optional*, defaults to `100`): - Number of top tokens used by the default SDFT top-k logit objective. - distillation_is_clip (`float`, *optional*, defaults to `2.0`): - Clipping coefficient for importance sampling in self-distillation. `None` disables clipping. - distillation_add_tail (`bool`, *optional*, defaults to `False`): - Whether to add a tail bucket for non-top-k probability mass. - num_loss_tokens_to_skip (`int`, *optional*, defaults to `0`): - Number of initial completion tokens to exclude from the distillation loss. - - > Parameters that control the teacher - - teacher_model_kind (`str`, *optional*, defaults to `"base"`): - Semantic teacher choice for SDFT. `base` uses the initial student, `live` uses the current student, and - `ema` uses an exponentially averaged teacher. - teacher_update_rate (`float`, *optional*, defaults to `0.05`): - EMA update rate used when `teacher_model_kind="ema"`. A value of `1.0` reduces the update to a hard - overwrite, periodically resyncing the teacher to the current student weights. - teacher_sync_steps (`int`, *optional*, defaults to `1`): - Number of optimizer steps between teacher updates. - - > Parameters that control teacher-conditioned generation - - generate_from_teacher (`bool`, *optional*, defaults to `False`): - Whether on-policy generation should use the teacher-conditioned prompt instead of the student prompt. - teacher_prompt_template (`str`, *optional*, defaults to `"{prompt}\n\n{privileged_context}"`): - Template used to combine the student prompt and privileged context into the teacher prompt. - - > Parameters that control the model - - model_init_kwargs (`dict[str, Any]`, *optional*): - Keyword arguments for `transformers.AutoModelForCausalLM.from_pretrained`, used when the `model` argument - of the `SDFTTrainer` is provided as a string. The `revision` value is also used when loading the processing - class. - trust_remote_code (`bool`, *optional*, defaults to `False`): - Whether to allow loading models and tokenizers that ship custom Python code from the Hub. Forwarded to - [`~transformers.AutoModelForCausalLM.from_pretrained`] and [`~transformers.AutoProcessor.from_pretrained`], - for both the student and teacher. - disable_dropout (`bool`, *optional*, defaults to `True`): - Whether to disable dropout in the student and teacher models. - - > Parameters that control data preprocessing - - remove_unused_columns (`bool`, *optional*, defaults to `False`): - Whether to only keep the columns required by the trainer in the dataset. Keep this to `False` if you - provide extra columns (such as `privileged_context`) that the trainer needs. - max_prompt_length (`int`, *optional*, defaults to `512`): - Maximum prompt length. Longer prompts are truncated from the left. - shuffle_dataset (`bool`, *optional*, defaults to `True`): - Whether to shuffle the training dataset. - - > Parameters that control generation - - num_generations (`int`, *optional*, defaults to `8`): - Number of generations to sample. The effective batch size (num_processes * per_device_batch_size * - gradient_accumulation_steps) must be evenly divisible by this value. - num_generations_eval (`int`, *optional*): - Number of generations to sample during evaluation. This allows using fewer generations during evaluation to - save computation. If `None`, uses the value of `num_generations`. - max_completion_length (`int`, *optional*, defaults to `256`): - Maximum length of the generated completion. - temperature (`float`, *optional*, defaults to `1.0`): - Temperature for sampling. The higher the temperature, the more random the completions. - top_p (`float`, *optional*, defaults to `1.0`): - Float that controls the cumulative probability of the top tokens to consider. Must be in (0, 1]. Set to 1.0 - to consider all tokens. - top_k (`int`, *optional*, defaults to `0`): - Number of highest probability vocabulary tokens to keep for top-k-filtering. If `0`, top-k-filtering is - disabled and all tokens are considered. - min_p (`float`, *optional*): - Minimum token probability, which will be scaled by the probability of the most likely token. It must be a - value between 0.0 and 1.0. Typical values are in the 0.01-0.2 range. - repetition_penalty (`float`, *optional*, defaults to `1.0`): - Float that penalizes new tokens based on whether they appear in the prompt and the generated text so far. - Values > 1.0 encourage the model to use new tokens, while values < 1.0 encourage the model to repeat - tokens. - cache_implementation (`str`, *optional*): - Implementation of the cache method for faster generation when use_vllm is set to False. - generation_kwargs (`dict[str, Any]`, *optional*): - Additional keyword arguments to pass to `GenerationConfig` (if using transformers) or `SamplingParams` (if - using vLLM) when sampling completions. This can be used to further customize the generation behavior, such - as setting `suppress_tokens`, `num_beams`, etc. If it contains keys that conflict with the other generation - parameters (like `min_p`, `top_p`, etc.), they will override them. - chat_template_kwargs (`dict[str, Any]`, *optional*): - Additional keyword arguments to pass to the `apply_chat_template` function when generating completions. - ds3_gather_for_generation (`bool`, *optional*, defaults to `True`): - This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for generation, - improving generation speed. However, disabling this option allows training models that exceed the VRAM - capacity of a single GPU, albeit at the cost of slower generation. Disabling this option is not compatible - with vLLM generation. - - > Parameters that control generation acceleration powered by vLLM - - use_vllm (`bool`, *optional*, defaults to `False`): - Whether to use vLLM for generating completions. If set to `True`, the trainer will use vLLM for generation - instead of the default model.generate(). Requires `vllm` to be installed. - use_teacher_server (`bool`, *optional*, defaults to `False`): - Compute teacher logprobs from the running vLLM generation server instead of a local teacher forward. Only - supported for `teacher_model_kind='live'` with `use_vllm=True` and `vllm_mode='server'`, and - `distillation_mode` in {'sampled_token', 'topk_logits'} (the server returns the teacher's top-k logprobs, - not the full vocabulary; `topk_logits` distills over the teacher's own top-k support). - vllm_mode (`str`, *optional*, defaults to `"colocate"`): - Mode to use for vLLM integration when `use_vllm` is set to `True`. Must be one of `'server'` or - `'colocate'`. `'server'`: The trainer will send generation requests to a separate vLLM server. Make sure a - vLLM server is running (start with `vllm serve`). `'colocate'`: vLLM will run in the same process and share - the training GPUs. This avoids the need for a separate server but may cause resource contention with - training. - vllm_model_impl (`str`, *optional*, defaults to `"vllm"`): - Model implementation to use for vLLM. Must be one of `transformers` or `vllm`. `transformers`: Use the - `transformers` backend for model implementation. `vllm`: Use the `vllm` library for model implementation. - vllm_enable_sleep_mode (`bool`, *optional*, defaults to `False`): - Enable vLLM sleep mode to offload weights/cache during the optimizer step. Keeps GPU memory usage low, but - waking the engine adds host–device transfer latency. - vllm_server_base_url (`str`, *optional*): - Base URL for the vLLM server (e.g., 'http://localhost:8000'). If provided, `vllm_server_host` and - `vllm_server_port` are ignored. - vllm_server_host (`str`, *optional*, defaults to `"0.0.0.0"`): - Host of the vLLM server to connect to. Ignored if vllm_server_base_url is provided. - vllm_server_port (`int`, *optional*, defaults to `8000`): - Port of the vLLM server to connect to. Ignored if vllm_server_base_url is provided. - vllm_group_port (`int`, *optional*, defaults to `51216`): - Port number for the weight update group. This is used to communicate with the vLLM server. Unless the port - is occupied, there is no need to change it. - vllm_server_timeout (`float`, *optional*, defaults to `240.0`): - Total timeout duration in seconds to wait for the vLLM server to be up. If the server is not up after the - timeout, a `ConnectionError` is raised. - vllm_tensor_parallel_size (`int`, *optional*, defaults to `1`): - Control the tensor parallel size for vLLM. This setting only applies when `vllm_mode` is set to - `'colocate'`. If you are using `vllm_mode='server'`, this parameter must be passed separately when - launching the vLLM server via the `--vllm_tensor_parallel_size` flag. - vllm_gpu_memory_utilization (`float`, *optional*, defaults to `0.3`): - Control the GPU memory utilization for vLLM. This setting only applies when `vllm_mode` is set to - `'colocate'`. If you are using `vllm_mode='server'`, this parameter must be passed separately when - launching the vLLM server via the `--vllm_gpu_memory_utilization` flag. - vllm_max_model_length (`int`, *optional*): - Context window for vLLM. Set it to at least the maximum prompt length in the dataset plus - `max_completion_length`; if omitted, it is inferred from the model config. - - > Parameters that control the training - - num_iterations (`int`, *optional*, defaults to `1`): - Number of iterations per batch (denoted as μ in the algorithm). - generation_batch_size (`int`, *optional*): - Batch size to use for generation. If `None`, it defaults to the effective training batch size: - `per_device_train_batch_size * num_processes * steps_per_generation`. - steps_per_generation (`int`, *optional*): - Number of steps per generation. If `None`, it defaults to `gradient_accumulation_steps`. - - """ - _unsloth_accepts_adamw_flash = True - vllm_sampling_params: Optional[Any] = field( - default = None, - metadata = {'help': 'vLLM SamplingParams'}, - ) - unsloth_num_chunks : Optional[int] = field( - default = -1, - metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'}, - ) - unsloth_logit_chunk_multiplier : Optional[int] = field( - default = None, - metadata = {'help': 'Multiplier for chunked logit computations.'}, - ) - unsloth_grpo_mini_batch : Optional[int] = field( - default = None, - metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'}, - ) - - def __init__( - self, - output_dir = None, - per_device_train_batch_size = 4, - num_train_epochs = 3.0, - max_steps = -1, - learning_rate = 5e-05, - lr_scheduler_type = 'linear', - lr_scheduler_kwargs = None, - warmup_steps = 0.1, - optim = 'adamw_8bit', - optim_args = None, - weight_decay = 0.001, - adam_beta1 = 0.9, - adam_beta2 = 0.999, - adam_epsilon = 1e-08, - optim_target_modules = None, - gradient_accumulation_steps = 2, - average_tokens_across_devices = True, - max_grad_norm = 1.0, - label_smoothing_factor = 0.0, - bf16 = False, - fp16 = False, - bf16_full_eval = False, - fp16_full_eval = False, - tf32 = None, - gradient_checkpointing = True, - gradient_checkpointing_kwargs = None, - torch_compile = False, - torch_compile_backend = None, - torch_compile_mode = None, - use_liger_kernel = False, - liger_kernel_config = None, - use_cache = False, - neftune_noise_alpha = None, - torch_empty_cache_steps = 250, - auto_find_batch_size = False, - logging_strategy = 'steps', - logging_steps = 1, - logging_first_step = False, - log_on_each_node = True, - logging_nan_inf_filter = False, - include_num_input_tokens_seen = False, - log_level = 'passive', - log_level_replica = 'warning', - disable_tqdm = None, - report_to = 'none', - run_name = None, - project = 'huggingface', - trackio_space_id = 'trackio', - eval_strategy = 'no', - eval_steps = None, - eval_delay = 0, - per_device_eval_batch_size = 8, - prediction_loss_only = False, - eval_on_start = False, - eval_do_concat_batches = True, - eval_use_gather_object = False, - eval_accumulation_steps = 2, - batch_eval_metrics = False, - save_only_model = False, - save_strategy = 'steps', - save_steps = 500, - save_on_each_node = False, - save_total_limit = None, - enable_jit_checkpoint = False, - push_to_hub = False, - hub_token = None, - hub_private_repo = None, - hub_model_id = None, - hub_strategy = 'every_save', - hub_always_push = False, - hub_revision = None, - load_best_model_at_end = False, - metric_for_best_model = None, - greater_is_better = None, - ignore_data_skip = False, - restore_callback_states_from_checkpoint = False, - full_determinism = False, - seed = 3407, - data_seed = None, - use_cpu = False, - accelerator_config = None, - parallelism_config = None, - dataloader_drop_last = False, - dataloader_num_workers = 0, - dataloader_pin_memory = True, - dataloader_persistent_workers = False, - dataloader_prefetch_factor = None, - remove_unused_columns = False, - label_names = None, - train_sampling_strategy = 'random', - length_column_name = 'length', - ddp_find_unused_parameters = None, - ddp_bucket_cap_mb = None, - ddp_broadcast_buffers = None, - ddp_backend = None, - ddp_timeout = 1800, - fsdp = None, - fsdp_config = None, - deepspeed = None, - debug = '', - skip_memory_metrics = True, - do_train = False, - do_eval = False, - do_predict = False, - resume_from_checkpoint = None, - warmup_ratio = None, - logging_dir = None, - local_rank = -1, - model_init_kwargs = None, - trust_remote_code = False, - disable_dropout = True, - max_prompt_length = 512, - num_generations = 8, - num_generations_eval = None, - max_completion_length = 256, - ds3_gather_for_generation = True, - shuffle_dataset = True, - generation_batch_size = None, - steps_per_generation = None, - temperature = 1.0, - top_p = 1.0, - top_k = 0, - min_p = None, - generation_kwargs = {}, - chat_template_kwargs = None, - repetition_penalty = 1.0, - cache_implementation = None, - use_vllm = False, - use_teacher_server = False, - vllm_mode = 'colocate', - vllm_model_impl = 'vllm', - vllm_enable_sleep_mode = False, - vllm_server_base_url = None, - vllm_server_host = '0.0.0.0', - vllm_server_port = 8000, - vllm_group_port = 51216, - vllm_server_timeout = 240.0, - vllm_tensor_parallel_size = 1, - vllm_gpu_memory_utilization = 0.3, - vllm_max_model_length = None, - num_iterations = 1, - teacher_model_kind = 'base', - teacher_update_rate = 0.05, - teacher_sync_steps = 1, - distillation_alpha = 0.5, - distillation_mode = 'topk_logits', - distillation_topk = 100, - distillation_is_clip = 2.0, - distillation_add_tail = False, - generate_from_teacher = False, - teacher_prompt_template = '{prompt}\n\n{privileged_context}', - num_loss_tokens_to_skip = 0, - vllm_sampling_params = None, - unsloth_num_chunks = -1, - unsloth_logit_chunk_multiplier = None, - unsloth_grpo_mini_batch = None, - - **kwargs, - ): - if learning_rate < 1e-7: print(f'Unsloth: Your learning rate of `{learning_rate}` is too small and less than 1e-7! Consider increasing it, otherwise gradient updates will be close to 0!') - if learning_rate > 1: print(f'Unsloth: Your learning rate of `{learning_rate}` is way too larger > 1! Consider decreasing it to 1e-1, otherwise gradient updates will explode!') - _unsloth_fsdp2 = os.environ.get('FSDP_VERSION', '').strip() == '2' - _unsloth_fsdp_cfg = fsdp_config - if fsdp and isinstance(_unsloth_fsdp_cfg, str) and os.path.isfile(_unsloth_fsdp_cfg): - try: - import json as _unsloth_json - with open(_unsloth_fsdp_cfg, encoding = 'utf-8') as _unsloth_f: - _unsloth_fsdp_cfg = _unsloth_json.load(_unsloth_f) - except Exception: - _unsloth_fsdp_cfg = None - if fsdp and _unsloth_fsdp_cfg is None: - _unsloth_fsdp_cfg = {} - if fsdp and isinstance(_unsloth_fsdp_cfg, dict): - _unsloth_fsdp2 = _unsloth_fsdp2 or str(_unsloth_fsdp_cfg.get('fsdp_version', _unsloth_fsdp_cfg.get('version', 1))) == '2' - _unsloth_flash_adamw = str(getattr(optim, 'value', optim)).strip().lower() == 'adamw_flash' - if _unsloth_flash_adamw: optim = 'adamw_torch' - if _unsloth_fsdp2 and str(getattr(optim, 'value', optim)) == 'adamw_8bit': - print("Unsloth: FSDP2 shards parameters as DTensors, which bitsandbytes' 8-bit optimizers cannot step. Switching optim = 'adamw_8bit' to 'adamw_torch_fused'.") - optim = 'adamw_torch_fused' - if num_train_epochs is None: - num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override - if output_dir is None and save_strategy == 'steps' and save_steps == 500: - output_dir = 'unsloth_training_checkpoints' - save_strategy = 'no' - if steps_per_generation is None and generation_batch_size is None: - ga = gradient_accumulation_steps - world_size = int(os.environ.get('WORLD_SIZE', '1')) - if (ga * world_size * per_device_train_batch_size) % num_generations != 0: - print('Unsloth: We now expect `per_device_train_batch_size` * `gradient_accumulation_steps` * `world_size` to be a multiple of `num_generations`.\nWe will change the batch size of ' + str(per_device_train_batch_size) + ' to the `num_generations` of ' + str(num_generations)) - per_device_train_batch_size = num_generations - - if eval_steps is not None and eval_strategy != 'steps': - print(f'Unsloth: `eval_steps = {eval_steps}` is ignored because `eval_strategy` is {getattr(eval_strategy, "value", eval_strategy)!r}. Set `eval_strategy = "steps"` to evaluate every `eval_steps` steps.') - - if temperature <= 0: - raise ValueError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.') - elif temperature >= 10: - raise ValueError('Unsloth: Please set a positive non-zero temperature less than 10, since sampling will be quite erratic.') - - - # One dict so the filter sees the mirrored parameters AND `**kwargs`: - # filtering kwargs alone would double-bind any argument TRL renamed, - # since the new name is itself a mirrored parameter. - _unsloth_config_arguments = dict( - output_dir = output_dir, - per_device_train_batch_size = per_device_train_batch_size, - num_train_epochs = num_train_epochs, - max_steps = max_steps, - learning_rate = learning_rate, - lr_scheduler_type = lr_scheduler_type, - lr_scheduler_kwargs = lr_scheduler_kwargs, - warmup_steps = warmup_steps, - optim = optim, - optim_args = optim_args, - weight_decay = weight_decay, - adam_beta1 = adam_beta1, - adam_beta2 = adam_beta2, - adam_epsilon = adam_epsilon, - optim_target_modules = optim_target_modules, - gradient_accumulation_steps = gradient_accumulation_steps, - average_tokens_across_devices = average_tokens_across_devices, - max_grad_norm = max_grad_norm, - label_smoothing_factor = label_smoothing_factor, - bf16 = bf16, - fp16 = fp16, - bf16_full_eval = bf16_full_eval, - fp16_full_eval = fp16_full_eval, - tf32 = tf32, - gradient_checkpointing = gradient_checkpointing, - gradient_checkpointing_kwargs = gradient_checkpointing_kwargs, - torch_compile = torch_compile, - torch_compile_backend = torch_compile_backend, - torch_compile_mode = torch_compile_mode, - use_liger_kernel = use_liger_kernel, - liger_kernel_config = liger_kernel_config, - use_cache = use_cache, - neftune_noise_alpha = neftune_noise_alpha, - torch_empty_cache_steps = torch_empty_cache_steps, - auto_find_batch_size = auto_find_batch_size, - logging_strategy = logging_strategy, - logging_steps = logging_steps, - logging_first_step = logging_first_step, - log_on_each_node = log_on_each_node, - logging_nan_inf_filter = logging_nan_inf_filter, - include_num_input_tokens_seen = include_num_input_tokens_seen, - log_level = log_level, - log_level_replica = log_level_replica, - disable_tqdm = disable_tqdm, - report_to = report_to, - run_name = run_name, - project = project, - trackio_space_id = trackio_space_id, - eval_strategy = eval_strategy, - eval_steps = eval_steps, - eval_delay = eval_delay, - per_device_eval_batch_size = per_device_eval_batch_size, - prediction_loss_only = prediction_loss_only, - eval_on_start = eval_on_start, - eval_do_concat_batches = eval_do_concat_batches, - eval_use_gather_object = eval_use_gather_object, - eval_accumulation_steps = eval_accumulation_steps, - batch_eval_metrics = batch_eval_metrics, - save_only_model = save_only_model, - save_strategy = save_strategy, - save_steps = save_steps, - save_on_each_node = save_on_each_node, - save_total_limit = save_total_limit, - enable_jit_checkpoint = enable_jit_checkpoint, - push_to_hub = push_to_hub, - hub_token = hub_token, - hub_private_repo = hub_private_repo, - hub_model_id = hub_model_id, - hub_strategy = hub_strategy, - hub_always_push = hub_always_push, - hub_revision = hub_revision, - load_best_model_at_end = load_best_model_at_end, - metric_for_best_model = metric_for_best_model, - greater_is_better = greater_is_better, - ignore_data_skip = ignore_data_skip, - restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint, - full_determinism = full_determinism, - seed = seed, - data_seed = data_seed, - use_cpu = use_cpu, - accelerator_config = accelerator_config, - parallelism_config = parallelism_config, - dataloader_drop_last = dataloader_drop_last, - dataloader_num_workers = dataloader_num_workers, - dataloader_pin_memory = dataloader_pin_memory, - dataloader_persistent_workers = dataloader_persistent_workers, - dataloader_prefetch_factor = dataloader_prefetch_factor, - remove_unused_columns = remove_unused_columns, - label_names = label_names, - train_sampling_strategy = train_sampling_strategy, - length_column_name = length_column_name, - ddp_find_unused_parameters = ddp_find_unused_parameters, - ddp_bucket_cap_mb = ddp_bucket_cap_mb, - ddp_broadcast_buffers = ddp_broadcast_buffers, - ddp_backend = ddp_backend, - ddp_timeout = ddp_timeout, - fsdp = fsdp, - fsdp_config = fsdp_config, - deepspeed = deepspeed, - debug = debug, - skip_memory_metrics = skip_memory_metrics, - do_train = do_train, - do_eval = do_eval, - do_predict = do_predict, - resume_from_checkpoint = resume_from_checkpoint, - warmup_ratio = warmup_ratio, - logging_dir = logging_dir, - local_rank = local_rank, - model_init_kwargs = model_init_kwargs, - trust_remote_code = trust_remote_code, - disable_dropout = disable_dropout, - max_prompt_length = max_prompt_length, - num_generations = num_generations, - num_generations_eval = num_generations_eval, - max_completion_length = max_completion_length, - ds3_gather_for_generation = ds3_gather_for_generation, - shuffle_dataset = shuffle_dataset, - generation_batch_size = generation_batch_size, - steps_per_generation = steps_per_generation, - temperature = temperature, - top_p = top_p, - top_k = top_k, - min_p = min_p, - generation_kwargs = generation_kwargs, - chat_template_kwargs = chat_template_kwargs, - repetition_penalty = repetition_penalty, - cache_implementation = cache_implementation, - use_vllm = use_vllm, - use_teacher_server = use_teacher_server, - vllm_mode = vllm_mode, - vllm_model_impl = vllm_model_impl, - vllm_enable_sleep_mode = vllm_enable_sleep_mode, - vllm_server_base_url = vllm_server_base_url, - vllm_server_host = vllm_server_host, - vllm_server_port = vllm_server_port, - vllm_group_port = vllm_group_port, - vllm_server_timeout = vllm_server_timeout, - vllm_tensor_parallel_size = vllm_tensor_parallel_size, - vllm_gpu_memory_utilization = vllm_gpu_memory_utilization, - vllm_max_model_length = vllm_max_model_length, - num_iterations = num_iterations, - teacher_model_kind = teacher_model_kind, - teacher_update_rate = teacher_update_rate, - teacher_sync_steps = teacher_sync_steps, - distillation_alpha = distillation_alpha, - distillation_mode = distillation_mode, - distillation_topk = distillation_topk, - distillation_is_clip = distillation_is_clip, - distillation_add_tail = distillation_add_tail, - generate_from_teacher = generate_from_teacher, - teacher_prompt_template = teacher_prompt_template, - num_loss_tokens_to_skip = num_loss_tokens_to_skip,**kwargs) - super().__init__(**_unsloth_filter_config_init_kwargs(SDFTConfig, _unsloth_config_arguments, mirrored_from = __class__)) - self.vllm_sampling_params = vllm_sampling_params - self.unsloth_num_chunks = unsloth_num_chunks - if unsloth_grpo_mini_batch is not None: - if self.generation_batch_size >= unsloth_grpo_mini_batch: - self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch - else: - raise ValueError( - f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, " - f"which is self.per_device_train_batch_size * gradient_accumulation_steps." - ) - self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier - - # Unsloth: keep the reentrant checkpoint path - if getattr(self, 'gradient_checkpointing', False): - _gc_kwargs = getattr(self, 'gradient_checkpointing_kwargs', None) or {} - if _gc_kwargs.get('context_fn') is None and not _gc_kwargs.get('debug', False): - _gc_kwargs['use_reentrant'] = True - self.gradient_checkpointing_kwargs = _gc_kwargs - from unsloth.optimizers.flash_adamw import is_flash_adamw_marked, mark_flash_adamw, unmark_flash_adamw - if _unsloth_flash_adamw or is_flash_adamw_marked(self): - mark_flash_adamw(self) - else: - unmark_flash_adamw(self) - -pass - -class _UnslothSDFTTrainer(_BaseTrainer): - """""" - - loss_is_scaled_for_ga = True - - _tag_names = ["trl", "sdft"] - _name = "SDFT" - config_cls = SDFTConfig - # docstyle-ignore - _paper = { - "title": "Self-Distillation Enables Continual Learning", - "id": "2601.19897", - "citation": textwrap.dedent("""\ - @article{shenfeld2026selfdistillation, - title = {{Self-Distillation Enables Continual Learning}}, - author = {Idan Shenfeld and Mehul Damani and Jonas H\\"ubotter and Pulkit Agrawal}, - year = 2026, - eprint = {arXiv:2601.19897} - }"""), - } - - def __init__( - self, - model: str | PreTrainedModel | nn.Module, - args: SDFTConfig | None = None, - train_dataset: Dataset | IterableDataset | None = None, - eval_dataset: Dataset | IterableDataset | dict[str, Dataset | IterableDataset] | None = None, - processing_class: PreTrainedTokenizerBase | ProcessorMixin | None = None, - callbacks: list[TrainerCallback] | None = None, - optimizers: tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None] = (None, None), - peft_config: "PeftConfig | None" = None, - ): - - if hasattr(model, 'vllm_engine') and hasattr(args, 'use_vllm'): - _unsloth_fsdp_cfg = getattr(args, 'fsdp_config', None) - if getattr(args, 'fsdp', None) and isinstance(_unsloth_fsdp_cfg, dict) and str(_unsloth_fsdp_cfg.get('version', _unsloth_fsdp_cfg.get('fsdp_version', 1))) == '2': - raise NotImplementedError('Unsloth: `fast_inference = True` (vLLM) is not supported with FSDP2.\nLoad the model with `fast_inference = False` to train with FSDP2, or launch without FSDP2 to use vLLM.') - if (getattr(args, 'use_vllm', False) == False): - args.use_vllm = True - if getattr(args, 'top_k', -1) is None or getattr(args, 'top_k', -1) == 0: - args.top_k = -1 - if isinstance(train_dataset, IterableDataset): - raise NotImplementedError("Iterable datasets are not yet supported in SDFTTrainer.") - if isinstance(eval_dataset, IterableDataset) or ( - isinstance(eval_dataset, dict) and any(isinstance(ds, IterableDataset) for ds in eval_dataset.values()) - ): - raise NotImplementedError("Iterable eval datasets are not yet supported in SDFTTrainer.") - - self.num_loss_tokens_to_skip = args.num_loss_tokens_to_skip - self.teacher_context_builder = DemonstrationTeacherContextBuilder(self) - - if train_dataset is None: - raise ValueError("`train_dataset` is required") - - # PEFT initializes the adapter weights randomly, so set_seed must be done before creating the model to ensure - # reproducibility. - set_seed(args.seed) - model_revision = None - if isinstance(model, str): - model_init_kwargs = args.model_init_kwargs or {} - if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: - model_init_kwargs["device_map"] = None - model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) - model_revision = model_init_kwargs.get("revision") - model = create_model_from_path(model, **model_init_kwargs) - elif args.model_init_kwargs is not None: - logger.warning( - "You passed `model_init_kwargs` to the self-distillation config, but `model` is already " - "instantiated. The `model_init_kwargs` will be ignored." - ) - - self.model_kwarg_keys = ( - inspect.signature(model.forward).parameters.keys() - if not hasattr(model, "get_base_model") - else inspect.signature(model.get_base_model().forward).parameters.keys() - ) - - if peft_config is None and getattr(model, "peft_config", None) is not None: - logger.warning( - "The provided self-distillation student model already contains a PEFT adapter. " - "This setup is accepted but not directly supported. In particular, `teacher_model_kind='base'` " - "may refer to the underlying base weights rather than the exact initially loaded student state " - "including its adapter. For unambiguous teacher behavior, start from a merged/non-adapter model " - "or manage separate adapters explicitly." - ) - if False: - if not is_peft_available(): - raise ImportError( - "You passed `peft_config` but the `peft` library is not installed. " - "Install it with `pip install trl[peft]`." - ) - if not isinstance(peft_config, PeftConfig): - raise TypeError( - f"`peft_config` must be a `peft.PeftConfig` instance (e.g. `peft.LoraConfig`), " - f"got {type(peft_config).__name__}." - ) - if is_peft_model(model): - raise ValueError( - "You passed a `PeftModel` instance together with a `peft_config`. Pass either a base " - "model with `peft_config`, or a pre-wrapped PEFT model." - ) - if peft_config is not None or (is_peft_available() and getattr(model, "peft_config", None) is not None): - pass - - if processing_class is None: - processing_class = AutoProcessor.from_pretrained( - get_config_model_id(model.config), - revision=model_revision, - truncation_side="left", - padding_side="left", - trust_remote_code=args.trust_remote_code, - ) - - if isinstance(processing_class, ProcessorMixin): - self._tokenizer = processing_class.tokenizer - elif isinstance(processing_class, PreTrainedTokenizerBase): - self._tokenizer = processing_class - else: - raise TypeError("The `processing_class` must be either a `PreTrainedTokenizerBase` or a `ProcessorMixin`") - - if self._tokenizer.pad_token is None: - self._tokenizer.pad_token = self._tokenizer.eos_token - # The model must agree with the tokenizer on the pad token from construction, so mirror it onto the model - # configs. - model.config.get_text_config().pad_token_id = self._tokenizer.pad_token_id - model.generation_config.pad_token_id = self._tokenizer.pad_token_id - - # Every eos id the model declares [e.g. an end-of-turn token] ends a completion, not only the tokenizer's - eos_token_ids = model.generation_config.eos_token_id - if eos_token_ids is None: - eos_token_ids = [] - elif isinstance(eos_token_ids, int): - eos_token_ids = [eos_token_ids] - if self._tokenizer.eos_token_id not in eos_token_ids: - eos_token_ids = [self._tokenizer.eos_token_id, *eos_token_ids] - self.eos_token_ids = eos_token_ids - - self.max_prompt_length = args.max_prompt_length - self.max_completion_length = args.max_completion_length - self.num_generations = args.num_generations - self.num_generations_eval = args.num_generations_eval or args.num_generations - self.num_iterations = args.num_iterations - self.shuffle_dataset = args.shuffle_dataset - self.temperature = args.temperature - self.generate_from_teacher = args.generate_from_teacher - self.use_vllm = args.use_vllm - self.chat_template_kwargs = args.chat_template_kwargs or {} - self._step = 0 - self._buffered_inputs = None - self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} - - self.generation_kwargs = { - "max_new_tokens": self.max_completion_length, - "do_sample": True, - "pad_token_id": self._tokenizer.pad_token_id, - "bos_token_id": self._tokenizer.bos_token_id, - "eos_token_id": self.eos_token_ids, - "temperature": args.temperature, - "top_p": args.top_p, - "top_k": args.top_k, - "min_p": args.min_p, - "repetition_penalty": args.repetition_penalty, - "cache_implementation": args.cache_implementation, - } - if args.generation_kwargs is not None: - self.generation_kwargs.update(args.generation_kwargs) - self.generation_config = GenerationConfig(**self.generation_kwargs, disable_compile=True) - - if hasattr(model, "warnings_issued"): - model.warnings_issued["estimate_tokens"] = True - - # Teacher logprobs from the running vLLM generation server [reuses the synced student weights] instead of a - # local teacher forward. MVP: `live` teacher only — the generation server holds the current student weights. - self.use_teacher_server = args.use_teacher_server - if self.use_teacher_server: - if not (args.use_vllm and args.vllm_mode == "server"): - raise ValueError( - "`use_teacher_server=True` requires `use_vllm=True` and `vllm_mode='server'`: teacher logprobs are " - "served by the same vLLM server used for generation." - ) - if args.teacher_model_kind != "live": - raise ValueError( - "`use_teacher_server=True` only supports `teacher_model_kind='live'` (the generation server holds " - f"the current student weights), got {args.teacher_model_kind!r}." - ) - if args.distillation_mode not in ("sampled_token", "topk_logits"): - raise ValueError( - "`use_teacher_server=True` only supports `distillation_mode` in {'sampled_token', 'topk_logits'}, " - f"got {args.distillation_mode!r}. The server returns the teacher's top-k logprobs, not the full " - "vocabulary, so `full_logits` is unavailable. Note `topk_logits` distills over the teacher's own " - "top-k support (the server cannot score the student's top-k indices)." - ) - super().__init__( - model=model, - args=args, - data_collator=identity, - train_dataset=train_dataset, - eval_dataset=eval_dataset, - processing_class=processing_class, - callbacks=callbacks, - optimizers=optimizers, - ) - - self._last_loaded_step = -1 if self.use_vllm else 0 - if self.use_vllm: - from trl.generation.vllm_generation import VLLMGeneration - - self.vllm_generation = VLLMGeneration( - model=self.model, - accelerator=self.accelerator, - processing_class=self.processing_class, - mode=args.vllm_mode, - server_base_url=args.vllm_server_base_url, - server_host=args.vllm_server_host, - server_port=args.vllm_server_port, - group_port=args.vllm_group_port, - server_timeout=args.vllm_server_timeout, - tensor_parallel_size=args.vllm_tensor_parallel_size, - gpu_memory_utilization=args.vllm_gpu_memory_utilization, - max_model_length=args.vllm_max_model_length, - max_num_seqs=args.per_device_train_batch_size - * args.vllm_tensor_parallel_size - * args.steps_per_generation, - enable_sleep_mode=args.vllm_enable_sleep_mode, - model_impl=args.vllm_model_impl, - trust_remote_code=args.trust_remote_code, - repetition_penalty=args.repetition_penalty, - temperature=self.temperature, - top_p=args.top_p, - top_k=args.top_k, - min_p=args.min_p, - max_completion_length=self.max_completion_length, - logprobs=None, - generation_kwargs=args.generation_kwargs, - ) - - # Per-rank read-only client to the same generation server for teacher scoring [weights are synced there by - # `VLLMGeneration`; scoring needs no weight-update communicator]. Mirrors the distillation trainer's - # `teacher_client`. - self.teacher_client = None - if self.use_teacher_server: - from trl.generation.vllm_client import VLLMClient - - base_url = args.vllm_server_base_url or f"http://{args.vllm_server_host}:{args.vllm_server_port}" - self.teacher_client = VLLMClient(base_url=base_url, connection_timeout=args.vllm_server_timeout) - - if args.disable_dropout: - disable_dropout_in_model(self.model) - - if hasattr(self.model, "add_model_tags"): - self.model.add_model_tags(self._tag_names) - - self._setup_teacher_model() - - def _set_signature_columns_if_needed(self): - if self._signature_columns is None: - self._signature_columns = ["prompt", "privileged_context"] - - def _dispatch_self_distillation_callback(self, event_name: str, **payload) -> None: - for callback in self.callback_handler.callbacks: - callback_fn = getattr(callback, event_name, None) - if callback_fn is not None: - callback_fn( - args=self.args, - state=self.state, - control=self.control, - model=self.model, - processing_class=self.processing_class, - **payload, - ) - - def _setup_teacher_model(self) -> None: - """Prepare teacher state according to the semantic teacher choice. - - Resolve `teacher_model_kind` × PEFT state into the effective teacher: - - - `"live"` (any model): - Teacher is the student. No divergence, no callback. - - `"base"` + PEFT model: - Teacher reuses `self.model`; the base weights are recovered downstream by disabling the adapter via - `use_adapter` during teacher forward. - - `"base"` + non-PEFT model: - Teacher is a frozen deepcopy of the initial student (falls through to the copy branch below). - - `"ema"` + pure-LoRA training: - Teacher reuses `self.model`; a dedicated `"teacher"` LoRA adapter is attached and updated by - `PEFTAdapterEMACallback`. Teacher forward switches to that adapter downstream. - - `"ema"` (otherwise): - Teacher is a frozen deepcopy synchronized each step by `SyncTeacherModelCallback`. - - Must be called after `super().__init__` so that `self.callback_handler` is available. - """ - - teacher_model_kind = self.args.teacher_model_kind - - if teacher_model_kind == "live": - self.teacher_model = self.model - return - - if teacher_model_kind == "base" and is_peft_model(self.model): - self.teacher_model = self.model - return - - if self._use_peft_ema_teacher_adapter(): - # Must run after super().__init__ so self.callback_handler exists. - self.add_callback( - PEFTAdapterEMACallback( - model=self.model, - teacher_adapter_name="teacher", - update_rate=self.args.teacher_update_rate, - sync_steps=self.args.teacher_sync_steps, - accelerator=self.accelerator, - ) - ) - self.teacher_model = self.model - return - - if is_peft_model(self.model): - raise ValueError( - "`teacher_model_kind='ema'` with a non-pure-LoRA PEFT model is not supported: the separate EMA " - "teacher cannot be parameter-matched to the PEFT student. Use pure-LoRA training, a non-PEFT model, " - "or `teacher_model_kind` in {'live', 'base'}." - ) - - # Build the teacher from the model path (like the GRPO/DPO reference model) rather than deep-copying the - # student: under ZeRO-3 the student params are already sharded, so a deep copy would clone empty shards. - model_init_kwargs = self.args.model_init_kwargs or {} - if self.args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: - model_init_kwargs["device_map"] = None - model_init_kwargs.setdefault("trust_remote_code", self.args.trust_remote_code) - self.teacher_model = create_model_from_path(get_config_model_id(self.model.config), **model_init_kwargs) - self.teacher_model.requires_grad_(False) - self.teacher_model.eval() - if self.is_deepspeed_enabled: - self.teacher_model = prepare_deepspeed(self.teacher_model, self.accelerator) - elif self.is_fsdp_enabled: - self.teacher_model = prepare_fsdp(self.teacher_model, self.accelerator) - else: - self.teacher_model = self.accelerator.prepare_model(self.teacher_model, evaluation_mode=True) - - if teacher_model_kind == "ema": - self.add_callback(SyncTeacherModelCallback(teacher_model=self.teacher_model, accelerator=self.accelerator)) - - def _use_peft_ema_teacher_adapter(self) -> bool: - return self.args.teacher_model_kind == "ema" and is_pure_lora_training(self.model, self.accelerator) - - def get_train_dataloader(self): - if self.train_dataset is None: - raise ValueError("Trainer: training requires a train_dataset.") - - train_dataset = self.train_dataset - data_collator = self.data_collator - if is_datasets_available() and isinstance(train_dataset, datasets.Dataset): - train_dataset = self._remove_unused_columns(train_dataset, description="training") - else: - data_collator = self._get_collator_with_removed_columns(data_collator, description="training") - - dataloader_params = { - "batch_size": self._train_batch_size * self.args.steps_per_generation, - "collate_fn": data_collator, - "num_workers": self.args.dataloader_num_workers, - "pin_memory": self.args.dataloader_pin_memory, - "persistent_workers": self.args.dataloader_persistent_workers, - } - if not isinstance(train_dataset, torch.utils.data.IterableDataset): - dataloader_params["sampler"] = self._get_train_sampler() - dataloader_params["drop_last"] = self.args.dataloader_drop_last - dataloader_params["worker_init_fn"] = partial( - seed_worker, num_workers=self.args.dataloader_num_workers, rank=self.args.process_index - ) - dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor - return self.accelerator.prepare(DataLoader(train_dataset, **dataloader_params)) - - def _get_train_sampler(self, dataset=None) -> Sampler: - if dataset is None: - dataset = self.train_dataset - return RepeatSampler( - data_source=dataset, - mini_repeat_count=self.num_generations, - batch_size=self.args.generation_batch_size // self.num_generations, - repeat_count=self.num_iterations * self.args.steps_per_generation, - shuffle=self.shuffle_dataset, - seed=self.args.seed, - ) - - def _get_eval_sampler(self, eval_dataset) -> Sampler: - return RepeatSampler( - data_source=eval_dataset, - mini_repeat_count=self.num_generations_eval, - seed=self.args.seed, - ) - - def training_step(self, model, inputs, num_items_in_batch): - output = super().training_step(model, inputs, num_items_in_batch) - self._step += 1 - return output - - def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None): - if not isinstance(inputs, dict): - inputs = self._prepare_inputs(inputs) - with torch.no_grad(): - with self.compute_loss_context_manager(): - loss = self.compute_loss(model, inputs) - return loss.detach(), None, None - - def _prepare_inputs(self, generation_batch): - """Return the per-step training batch, regenerating rollouts and buffering them for reuse in train mode. - - In train mode, rollouts are generated once every `steps_per_generation * num_iterations` steps and split into - per-step slices reused until the next regeneration. In eval mode, every batch is freshly prepared. - """ - mode = "train" if self.model.training else "eval" - if mode == "train": - generate_every = self.args.steps_per_generation * self.num_iterations - if self._step % generate_every == 0 or self._buffered_inputs is None: - buffered_batch = self._prepare_training_batch(generation_batch) - self._buffered_inputs = split_tensor_dict(buffered_batch, self.args.steps_per_generation) - self._dispatch_self_distillation_callback( - "on_generation_batch_built", - generate_every=generate_every, - steps_per_generation=self.args.steps_per_generation, - ) - return self._buffered_inputs[self._step % self.args.steps_per_generation] - return self._prepare_training_batch(generation_batch) - - def _prepare_training_batch(self, inputs: list[dict[str, Any]]) -> TrainingBatch: - """Sample student rollouts and construct teacher prompts""" - batch = self.sample_rollouts(inputs) - mode = "train" if self.model.training else "eval" - self._record_completion_metrics(mode, batch) - prompts = [example["prompt"] for example in inputs] - privileged_contexts = [example.get("privileged_context") for example in inputs] - teacher_batch = self.teacher_context_builder.build( - prompts, - privileged_contexts, - batch["completion_ids"], - batch["completion_mask"], - ) - batch.update( - { - "teacher_input_ids": teacher_batch["teacher_input_ids"], - "teacher_attention_mask": teacher_batch["teacher_attention_mask"], - } - ) - - self._dispatch_self_distillation_callback( - "on_self_distillation_batch_prepared", - old_per_token_logps=batch.get("old_per_token_logps"), - prompt_ids=batch["prompt_ids"], - completion_ids=batch["completion_ids"], - teacher_input_ids=batch["teacher_input_ids"], - teacher_attention_mask=batch["teacher_attention_mask"], - ) - return batch - - def sample_rollouts(self, inputs: list[dict[str, Any]]) -> TrainingBatch: - """Generate completions for a batch of prompts and assemble the training batch.""" - prompts = [example["prompt"] for example in inputs] - privileged_contexts = [example.get("privileged_context") for example in inputs] - - student_prompt_ids_list = self._tokenize_prompts(prompts) - if self.generate_from_teacher: - generation_prompts = self.teacher_context_builder.select_generation_prompts(prompts, privileged_contexts) - generation_prompt_ids_list = self._tokenize_prompts(generation_prompts) - else: - generation_prompts = prompts - generation_prompt_ids_list = student_prompt_ids_list - - self._dispatch_self_distillation_callback( - "on_generation_prompts_selected", - generation_prompts=generation_prompts, - generation_prompt_text=None, - ) - - _, completion_ids_list = self._generate(generation_prompt_ids_list) - device = self.accelerator.device - prompt_ids = [torch.tensor(ids) for ids in student_prompt_ids_list] - prompt_mask = [torch.ones_like(ids, dtype=torch.long) for ids in prompt_ids] - prompt_ids = pad(prompt_ids, padding_value=self._tokenizer.pad_token_id, padding_side="left").to(device=device) - prompt_mask = pad(prompt_mask, padding_value=0, padding_side="left").to(device=device) - - completion_ids = [torch.tensor(ids) for ids in completion_ids_list] - completion_mask = [torch.ones_like(ids, dtype=torch.long) for ids in completion_ids] - completion_ids = pad(completion_ids, padding_value=self._tokenizer.pad_token_id, padding_side="right").to( - device=device - ) - completion_mask = pad(completion_mask, padding_value=0, padding_side="right").to(device=device) - - old_per_token_logps = None - if not self.generate_from_teacher: - old_per_token_logps = self._compute_rollout_logps( - prompt_ids=prompt_ids, - prompt_mask=prompt_mask, - completion_ids=completion_ids, - completion_mask=completion_mask, - ) - batch: TrainingBatch = { - "prompt_ids": prompt_ids, - "prompt_mask": prompt_mask, - "completion_ids": completion_ids, - "completion_mask": completion_mask, - "raw_completion_lengths": torch.tensor( - [len(ids) for ids in completion_ids_list], device=device, dtype=torch.long - ), - } - if old_per_token_logps is not None: - batch["old_per_token_logps"] = old_per_token_logps - return batch - - def _get_completion_ids_list(self, batch: TrainingBatch) -> list[list[int]]: - raw_completion_lengths = batch["raw_completion_lengths"].detach().cpu().tolist() - return [ - ids[:length].tolist() - for ids, length in zip(batch["completion_ids"].detach().cpu(), raw_completion_lengths, strict=True) - ] - - def _tokenize_prompts(self, prompts: list[Any]) -> list[list[int]]: - if is_conversational({"prompt": prompts[0]}): - tokenized = self.processing_class.apply_chat_template( - conversation=prompts, - add_generation_prompt=True, - tokenize=True, - return_dict=True, - **self.chat_template_kwargs, - ) - prompt_ids = tokenized["input_ids"] - else: - prompt_ids = self.processing_class(text=prompts)["input_ids"] - if self.max_prompt_length is not None: - prompt_ids = [ids[-self.max_prompt_length :] for ids in prompt_ids] - return prompt_ids - - def _generate(self, prompt_ids: list[list[int]]) -> tuple[list[list[int]], list[list[int]]]: - if self.use_vllm: - return self._generate_vllm(prompt_ids) - return self._generate_transformers(prompt_ids) - - def _generate_vllm(self, prompt_ids: list[list[int]]) -> tuple[list[list[int]], list[list[int]]]: - if self.state.global_step != self._last_loaded_step: - self.vllm_generation.sync_weights() - self._last_loaded_step = self.state.global_step - - mode = "train" if self.model.training else "eval" - num_generations = self.num_generations if mode == "train" else self.num_generations_eval - prompt_ids_out, completion_ids_list, _, _ = self.vllm_generation.generate( - prompts=prompt_ids, - images=None, - num_generations=num_generations, - ) - return prompt_ids_out, completion_ids_list - - def _generate_transformers(self, prompt_ids: list[list[int]]) -> tuple[list[list[int]], list[list[int]]]: - device = self.accelerator.device - prompt_tensors = [torch.tensor(ids) for ids in prompt_ids] - padded_ids = pad(prompt_tensors, padding_value=self._tokenizer.pad_token_id, padding_side="left").to( - device=device - ) - attention_mask = pad([torch.ones_like(t) for t in prompt_tensors], padding_value=0, padding_side="left").to( - device=device - ) - generate_inputs: dict[str, torch.Tensor | Any] = {"input_ids": padded_ids, "attention_mask": attention_mask} - - with ( - unwrap_model_for_generation( - self.model_wrapped, - self.accelerator, - gather_deepspeed3_params=self.args.ds3_gather_for_generation, - generation_kwargs=self.generation_kwargs, - ) as unwrapped_model, - torch.no_grad(), - FSDP.summon_full_params(self.model_wrapped, recurse=False) if self.is_fsdp_enabled else nullcontext(), - ): - prompt_completion_ids = unwrapped_model.generate( - **generate_inputs, generation_config=self.generation_config - ) - - prompt_length = generate_inputs["input_ids"].size(1) - completion_ids = prompt_completion_ids[:, prompt_length:] - is_eos = torch.isin(completion_ids, torch.tensor(self.eos_token_ids, device=device)) - eos_idx = torch.full((is_eos.size(0),), is_eos.size(1), dtype=torch.long, device=device) - eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)] - seq_idx = torch.arange(is_eos.size(1), device=device).expand(is_eos.size(0), -1) - completion_mask = (seq_idx <= eos_idx.unsqueeze(1)).int() - completion_ids_list = [ - c[m].tolist() for c, m in zip(completion_ids.cpu(), completion_mask.bool().cpu(), strict=True) - ] - return prompt_ids, completion_ids_list - - def _compute_rollout_logps( - self, - prompt_ids: torch.Tensor, - prompt_mask: torch.Tensor, - completion_ids: torch.Tensor, - completion_mask: torch.Tensor, - ) -> torch.Tensor | None: - generate_every = self.args.steps_per_generation * self.num_iterations - old_per_token_logps = None - - if self.args.gradient_accumulation_steps % generate_every != 0: - prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1) - attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) - logits_to_keep = completion_ids.size(1) - with torch.no_grad(): - logits = self._forward_logits( - self.model, - prompt_completion_ids, - attention_mask, - logits_to_keep, - ) - old_per_token_logps = selective_log_softmax(logits, completion_ids) - - return old_per_token_logps - - def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): - if return_outputs: - raise ValueError("The SDFTTrainer does not support returning outputs") - - if self.use_teacher_server: - loss = self._compute_server_distillation_loss(model, inputs) - else: - distillation_logits = self._compute_teacher_student_logits(model, self.teacher_model, inputs) - loss = self._compute_self_distillation_loss(model, inputs, distillation_logits) - accumulation_scale = self.current_gradient_accumulation_steps if self.model.training else 1.0 - return loss / accumulation_scale - - def _compute_self_distillation_loss( - self, - model, - inputs: TrainingBatch, - distillation_logits: DistillationLogits, - ) -> torch.Tensor: - """Compute the per-token distillation loss and aggregate by normalizing over sequence length. - - Dispatches between three objectives based on `distillation_mode`: - - - `"topk_logits"`: top-k approximation of the divergence, optionally with a tail bucket for the remaining - probability mass (`distillation_add_tail`). - - `"full_logits"`: full-vocab divergence. - - `"sampled_token"`: token-level (reverse-KL) distillation on sampled `completion_ids`. - - When `distillation_is_clip` is set and `old_per_token_logps` are available, the loss is corrected by a clipped - importance-sampling ratio between the current student and the student at rollout time. - """ - if distillation_logits.loss_mask.sum() == 0: - mode = "train" if model.training else "eval" - self._log_self_distillation_metric(mode, 0.0) - # Keep the zero loss attached to the student graph so backward produces zero gradients instead of stopping. - return distillation_logits.student_logits.sum() * 0.0 - - if self.args.distillation_mode == "topk_logits": - if self.args.distillation_topk is None: - raise ValueError("`distillation_mode='topk_logits'` requires `distillation_topk` to be set.") - per_token_loss = compute_topk_self_distillation_loss( - distillation_logits.student_logits, - distillation_logits.teacher_logits, - distillation_topk=self.args.distillation_topk, - distillation_alpha=self.args.distillation_alpha, - distillation_add_tail=self.args.distillation_add_tail, - ) - elif self.args.distillation_mode == "full_logits": - per_token_loss = compute_full_logit_self_distillation_loss( - distillation_logits.student_logits, - distillation_logits.teacher_logits, - distillation_alpha=self.args.distillation_alpha, - ) - elif self.args.distillation_mode == "sampled_token": - per_token_loss = compute_sampled_token_self_distillation_loss( - distillation_logits.student_logits, - distillation_logits.teacher_logits, - distillation_logits.completion_ids, - distillation_alpha=self.args.distillation_alpha, - ) - else: - raise ValueError( - "distillation_mode must be one of: 'sampled_token', 'full_logits', 'topk_logits', " - f"got {self.args.distillation_mode!r}" - ) - - old_per_token_logps = inputs.get("old_per_token_logps") - if self.args.distillation_is_clip is not None and old_per_token_logps is not None: - student_per_token_logps = selective_log_softmax( - distillation_logits.student_logits, - distillation_logits.completion_ids, - ) - per_token_loss = apply_importance_sampling_clipping( - per_token_loss, - student_per_token_logps, - old_per_token_logps, - self.args.distillation_is_clip, - ) - - loss = (per_token_loss * distillation_logits.loss_mask).sum(-1) / distillation_logits.loss_mask.sum(-1).clamp( - min=1.0 - ) - loss = loss.mean() - - mode = "train" if model.training else "eval" - mean_distill_loss = ( - per_token_loss * distillation_logits.loss_mask - ).sum() / distillation_logits.loss_mask.sum().clamp(min=1.0) - self._log_self_distillation_metric( - mode, - self.accelerator.gather(mean_distill_loss).mean().item(), - ) - return loss - - def _compute_server_distillation_loss(self, model, inputs: TrainingBatch) -> torch.Tensor: - """Distillation loss with teacher logprobs served by the vLLM generation server (`teacher_model_kind='live'`). - - The student is forwarded locally (grad) for its logits; the teacher logprobs are fetched from the server (no - local teacher forward). `sampled_token` distills the realized token (reverse KL); `topk_logits` distills over - the teacher's own top-k support (the server cannot score the student's top-k indices). - """ - # Buffered batches are reused across optimizer steps (`num_iterations > 1`), so the server weights can lag the - # live student; re-sync before scoring. - if self.state.global_step != self._last_loaded_step: - self.vllm_generation.sync_weights() - # Only the main process pushes weights to the server, so hold every rank here until the push completes; - # otherwise non-main ranks could score against half-updated weights. - self.accelerator.wait_for_everyone() - self._last_loaded_step = self.state.global_step - - completion_ids = inputs["completion_ids"] - completion_mask = inputs["completion_mask"] - logits_to_keep = completion_ids.size(1) - - # SDFT skips the first few completion tokens in the distillation loss to suppress teacher-prompt artifacts. - loss_mask = completion_mask - if self.num_loss_tokens_to_skip > 0: - token_positions = torch.arange(completion_mask.size(1), device=completion_mask.device).unsqueeze(0) - loss_mask = completion_mask * (token_positions >= self.num_loss_tokens_to_skip).long() - - student_input_ids = torch.cat([inputs["prompt_ids"], completion_ids], dim=1) - student_attention_mask = torch.cat([inputs["prompt_mask"], completion_mask], dim=1) - student_logits = self._forward_logits( - model=model, - input_ids=student_input_ids, - attention_mask=student_attention_mask, - logits_to_keep=logits_to_keep, - ) - - mode = "train" if model.training else "eval" - teacher = self._get_teacher_token_logprobs_from_server(inputs, logits_to_keep) - # Padding positions come back as `-inf`; the masked-out positions are neutralized to finite values below so the - # divergence does not leak `+inf` through them (mirrors the distillation trainer). - required = loss_mask.bool() - student_per_token_logps = selective_log_softmax(student_logits, completion_ids) - - teacher_per_token_logps = teacher["actual_logprobs"] - if (required & ~torch.isfinite(teacher_per_token_logps)).any(): - raise ValueError("Teacher server returned no logprob for a required completion token.") - teacher_per_token_logps = torch.where( - required, teacher_per_token_logps, torch.zeros_like(teacher_per_token_logps) - ) - - # Diagnostic for disagreement between local student scores and server teacher scores on realized tokens. - # Sudden jumps can indicate stale server weights or numerical drift. - abs_diff = ((student_per_token_logps.detach() - teacher_per_token_logps).abs() * loss_mask).sum() / ( - loss_mask.sum().clamp(min=1.0) - ) - self._metrics[mode]["self_distillation/server_logprob_abs_diff"].append( - self.accelerator.gather(abs_diff).mean().item() - ) - - if self.args.distillation_mode == "sampled_token": - # Reverse-KL token-level objective, matching `compute_sampled_token_self_distillation_loss`. - per_token_loss = (student_per_token_logps - teacher_per_token_logps).detach() * student_per_token_logps - else: - teacher_topk_logps = teacher["topk_logprobs"] - teacher_topk_ids = teacher["topk_token_ids"] - # Project the student onto the teacher's top-k support, then renormalize (or add a tail bucket) before the - # divergence, matching `compute_topk_self_distillation_loss` on the teacher's support. - keep = required.unsqueeze(-1) - student_topk_logps = torch.gather( - torch.log_softmax(student_logits, dim=-1), dim=-1, index=teacher_topk_ids - ) - student_topk_logps = torch.where(keep, student_topk_logps, torch.zeros_like(student_topk_logps)) - teacher_topk_logps = torch.where(keep, teacher_topk_logps, torch.zeros_like(teacher_topk_logps)) - if self.args.distillation_add_tail: - student_topk_logps = add_tail_bucket(student_topk_logps) - teacher_topk_logps = add_tail_bucket(teacher_topk_logps) - else: - student_topk_logps = student_topk_logps - torch.logsumexp(student_topk_logps, dim=-1, keepdim=True) - teacher_topk_logps = teacher_topk_logps - torch.logsumexp(teacher_topk_logps, dim=-1, keepdim=True) - per_token_loss = compute_divergence(student_topk_logps, teacher_topk_logps, self.args.distillation_alpha) - - old_per_token_logps = inputs.get("old_per_token_logps") - if self.args.distillation_is_clip is not None and old_per_token_logps is not None: - per_token_loss = apply_importance_sampling_clipping( - per_token_loss, student_per_token_logps, old_per_token_logps, self.args.distillation_is_clip - ) - - loss = (per_token_loss * loss_mask).sum(-1) / loss_mask.sum(-1).clamp(min=1.0) - loss = loss.mean() - mean_distill_loss = (per_token_loss * loss_mask).sum() / loss_mask.sum().clamp(min=1.0) - self._log_self_distillation_metric(mode, self.accelerator.gather(mean_distill_loss).mean().item()) - return loss - - def _get_teacher_token_logprobs_from_server( - self, inputs: TrainingBatch, logits_to_keep: int - ) -> dict[str, torch.Tensor]: - """Per-token teacher logprobs from the vLLM server. - - Returns a dict with `actual_logprobs` (`(B, T)`, the realized-token logprob for reverse KL) and `topk_logprobs` - / `topk_token_ids` (`(B, T, K)`, the teacher's top-k support for top-k logit distillation). Completions are - right-padded, so server values fill the leading positions of each row and the trailing padding keeps the `-inf` - / `0` sentinels (neutralized by the caller). - """ - import numpy as np - - teacher_input_ids = inputs["teacher_input_ids"] - teacher_attention_mask = inputs["teacher_attention_mask"] - device = teacher_input_ids.device - batch_size = teacher_input_ids.size(0) - prompt_attention_mask = teacher_attention_mask[:, : teacher_attention_mask.size(1) - logits_to_keep] - if self.args.distillation_mode == "topk_logits" and self.args.distillation_topk is None: - raise ValueError("`distillation_mode='topk_logits'` requires `distillation_topk` to be set.") - top_logprobs = self.args.distillation_topk if self.args.distillation_mode == "topk_logits" else 1 - - sequences, prompt_lengths, _ = build_teacher_request_inputs( - teacher_input_ids, teacher_attention_mask, prompt_attention_mask - ) - result = self.teacher_client.get_sequence_logprobs( - sequences=sequences, - prompt_lengths=prompt_lengths, - top_logprobs=top_logprobs, - temperature=self.temperature, - ) - - actual = np.full((batch_size, logits_to_keep), float("-inf"), dtype=np.float32) - topk = np.full((batch_size, logits_to_keep, top_logprobs), float("-inf"), dtype=np.float32) - topk_ids = np.zeros((batch_size, logits_to_keep, top_logprobs), dtype=np.int64) - for i in range(batch_size): - seq_actual = result["actual_logprobs"][i] - n = min(len(seq_actual), logits_to_keep) - if n: - actual[i, :n] = np.array(seq_actual, dtype=np.float32)[:n, 0] - topk[i, :n] = np.array(result["logprobs"][i], dtype=np.float32)[:n] - topk_ids[i, :n] = np.array(result["logprob_token_ids"][i], dtype=np.int64)[:n] - return { - "actual_logprobs": torch.from_numpy(actual).to(device), - "topk_logprobs": torch.from_numpy(topk).to(device), - "topk_token_ids": torch.from_numpy(topk_ids).to(device), - } - - def _compute_teacher_student_logits( - self, - model, - teacher_model, - inputs: TrainingBatch, - ) -> DistillationLogits: - """Compute the per-token logits of the student and teacher over the completion tokens. - - The student is forwarded on its own input (original prompt plus the sampled completion) while the teacher is - forwarded on its input (prompt, privileged context, and the same completion). Both sets of logits are aligned - to the completion tokens so they can be compared position-by-position in the distillation loss. - - The teacher forward runs under `torch.no_grad()` and the context resolved by - `_get_teacher_context_for_self_distillation`, which routes it to the correct weights. - """ - prompt_ids = inputs["prompt_ids"] - prompt_mask = inputs["prompt_mask"] - completion_ids = inputs["completion_ids"] - completion_mask = inputs["completion_mask"] - logits_to_keep = completion_ids.size(1) - - # SDFT skips the first few completion tokens in the distillation loss to suppress teacher-prompt artifacts. - loss_mask = completion_mask - if self.num_loss_tokens_to_skip > 0: - token_positions = torch.arange(completion_mask.size(1), device=completion_mask.device).unsqueeze(0) - loss_mask = completion_mask * (token_positions >= self.num_loss_tokens_to_skip).long() - - student_input_ids = torch.cat([prompt_ids, completion_ids], dim=1) - student_attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) - student_logits = self._forward_logits( - model=model, - input_ids=student_input_ids, - attention_mask=student_attention_mask, - logits_to_keep=logits_to_keep, - ) - - with torch.no_grad(), self._get_teacher_context_for_self_distillation(): - teacher_logits = self._forward_logits( - model=teacher_model, - input_ids=inputs["teacher_input_ids"], - attention_mask=inputs["teacher_attention_mask"], - logits_to_keep=logits_to_keep, - ) - - return DistillationLogits( - completion_ids=completion_ids, - loss_mask=loss_mask, - student_logits=student_logits, - teacher_logits=teacher_logits, - ) - - def _forward_logits( - self, - model, - input_ids: torch.Tensor, - attention_mask: torch.Tensor, - logits_to_keep: int, - ) -> torch.Tensor: - """Forward the model and return temperature-scaled logits aligned to the completion tokens.""" - model_inputs = { - "input_ids": input_ids, - "attention_mask": attention_mask, - "use_cache": False, - } - if "logits_to_keep" in self.model_kwarg_keys: - model_inputs["logits_to_keep"] = logits_to_keep + 1 - - logits = model(**model_inputs).logits - logits = logits[:, :-1, :] - logits = logits[:, -logits_to_keep:, :] - return logits / self.temperature - - def _get_teacher_context_for_self_distillation(self): - """Return the context manager that routes the teacher forward to the correct weights. - - For non-PEFT models this is a no-op. For PEFT models: - - - `teacher_model_kind == "base"`: disable the student adapter so the teacher forward uses the base weights. - - `teacher_model_kind == "ema"` under pure-LoRA training: switch to the `"teacher"` LoRA adapter. - - otherwise: no-op; the teacher is a separate deepcopy. - """ - teacher_model_kind = self.args.teacher_model_kind - if not is_peft_model(self.model): - return nullcontext() - - target_model = self.accelerator.unwrap_model(self.teacher_model) - - if teacher_model_kind == "base": - return use_adapter(target_model, adapter_name=None) - if teacher_model_kind == "ema" and self._use_peft_ema_teacher_adapter(): - return use_adapter(target_model, adapter_name="teacher") - return nullcontext() - - def _record_completion_metrics(self, mode: str, batch: TrainingBatch) -> None: - device = self.accelerator.device - completion_ids_list = self._get_completion_ids_list(batch) - agg_completion_lengths = self.accelerator.gather(batch["raw_completion_lengths"]) - # Fail clearly if the generation backend returned no completions (avoids a cryptic min() error below). - if agg_completion_lengths.numel() == 0: - raise RuntimeError( - "No completions were generated. This usually means the generation backend failed to return any " - "results; see the generation logs above for the underlying error." - ) - self._metrics[mode]["completions/mean_length"].append(agg_completion_lengths.float().mean().item()) - self._metrics[mode]["completions/min_length"].append(agg_completion_lengths.float().min().item()) - self._metrics[mode]["completions/max_length"].append(agg_completion_lengths.float().max().item()) - - eos_and_pad = [*self.eos_token_ids, self._tokenizer.pad_token_id] - is_truncated = torch.tensor([ids[-1] not in eos_and_pad for ids in completion_ids_list], device=device) - agg_is_truncated = self.accelerator.gather(is_truncated) - self._metrics[mode]["completions/clipped_ratio"].append(agg_is_truncated.float().mean().item()) - term_completion_lengths = agg_completion_lengths[~agg_is_truncated] - if len(term_completion_lengths) == 0: - term_completion_lengths = torch.zeros(1, device=device) - self._metrics[mode]["completions/mean_terminated_length"].append(term_completion_lengths.float().mean().item()) - self._metrics[mode]["completions/min_terminated_length"].append(term_completion_lengths.float().min().item()) - self._metrics[mode]["completions/max_terminated_length"].append(term_completion_lengths.float().max().item()) - - def _log_self_distillation_metric(self, mode: str, value: float) -> None: - metric_prefix = self._name.lower().replace(" ", "_") - self._metrics[mode]["self_distillation/distillation_loss"].append(value) - self._metrics[mode][f"{metric_prefix}/distillation_loss"].append(value) - - def log(self, logs: dict[str, float], start_time: float | None = None) -> None: - mode = "train" if self.model.training else "eval" - metrics = {k: sum(v) / len(v) for k, v in self._metrics[mode].items() if v} - if mode == "eval": - metrics = {f"eval_{k}": v for k, v in metrics.items()} - logs = {**logs, **metrics} - super().log(logs, start_time) - self._metrics[mode].clear() -class UnslothSDFTTrainer(_UnslothSDFTTrainer): - """ - Trainer for SDFT-style on-policy self-distillation with explicit teacher prompts. - """ - def __init__( - self, - model, - args = None, - train_dataset = None, - eval_dataset = None, - processing_class = None, - callbacks = None, - peft_config = None, - **kwargs - ): - if args is None: args = UnslothSDFTConfig() - _unsloth_precision_requested = tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) - use_bf16 = getattr(args, 'bf16', False) - if type(use_bf16) is not bool: use_bf16 = False - use_fp16 = getattr(args, 'fp16', False) - if type(use_fp16) is not bool: use_fp16 = False - force_float32 = False - try: - from unsloth_zoo.device_type import device_is_bf16_supported as _bf16_supported - except Exception: - _bf16_supported = torch.cuda.is_bf16_supported - try: - from unsloth_zoo.device_type import DEVICE_TYPE as _unsloth_device_type - except Exception: - _unsloth_device_type = None - full_finetuning = getattr(model, '_unsloth_full_finetuning', None) - if full_finetuning is None: full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1' - model_forced_float32 = getattr(model, '_unsloth_forced_float32', None) - if model_forced_float32 is None: model_forced_float32 = os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1' - if model_forced_float32 and not (full_finetuning and _bf16_supported()): - print('Unsloth: Switching to float32 training since model cannot work with float16') - force_float32 = True - mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') - dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None) - if dtype is None: dtype = model.get_input_embeddings().weight.dtype - from unsloth_zoo.utils import _get_dtype - dtype = _get_dtype(dtype) - float16 = dtype == torch.float16 - bfloat16 = dtype == torch.bfloat16 - float32 = dtype == torch.float32 - user_float32 = bool(getattr(model, '_unsloth_user_float32', False)) - if full_finetuning: - if bfloat16 and use_fp16: use_fp16 = False - if float16 and use_bf16: use_bf16 = False - if not force_float32 and bfloat16 and use_fp16 and _unsloth_device_type == 'xpu' and not _bf16_supported(): - print('Unsloth: The model was loaded in bfloat16 on purpose, so training uses bfloat16 mixed precision instead of float16.') - use_fp16 = False - use_bf16 = True - args.fp16 = False - args.bf16 = True - if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`') - if not force_float32 and (bfloat16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`') - if force_float32: - # Forced float32 training - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32' and float32 and user_float32 and not _bf16_supported(): - print('Unsloth: Model is in float32 and this GPU has no bfloat16 support, so training stays in float32. Pass fp16 = True to force float16 mixed precision instead.') - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32': - # Mixed precision training. bf16 only if the GPU supports it; V100/T4 use fp16. - use_bf16_amp = (not float16) and (_bf16_supported() or (bfloat16 and _unsloth_device_type == 'xpu')) - args.fp16 = not use_bf16_amp - args.bf16 = use_bf16_amp - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16_amp else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16_amp else 'fp16' - # args.mixed_precision is a new argument which needs to be set now - elif mixed_precision_dtype == 'bfloat16': - # Both False since bfloat16 full finetuning doesn't do any autocasting. - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif use_bf16 or use_fp16: - # transformers <5 exported this itself from fp16/bf16; 5.x dropped the write, so an - # explicit flag left it unset and GRPO readers defaulted to 'fp16', wrapping a - # bfloat16 model in a float16 autocast. See unslothai/unsloth#4891. - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16 else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16 else 'fp16' - - if getattr(args, 'eval_dataset', None) is not None and getattr(args, 'eval_strategy', 'no') == 'no': - args.eval_strategy = 'steps' - if getattr(args, 'eval_steps', None) is None: args.eval_steps = 0.1 - ga_steps = getattr(args, 'gradient_accumulation_steps', None) - if ga_steps is not None and ga_steps > 1: - from transformers import __version__ as transformers_version - if Version(transformers_version) <= Version('4.45.2'): - print('**** Unsloth: Please use our fixed gradient_accumulation_steps by updating transformers, TRL and Unsloth!\n' - '`pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth transformers trl unsloth_zoo`') - if getattr(args, 'eval_strategy', 'no') != 'no': - eval_bsz = getattr(args, 'per_device_eval_batch_size', 8) - eval_generations = getattr(args, 'num_generations_eval', None) or getattr(args, 'num_generations', None) or 1 - if eval_bsz == 8 and args.per_device_train_batch_size < eval_bsz and (args.per_device_train_batch_size * args.world_size) % eval_generations == 0: args.per_device_eval_batch_size = args.per_device_train_batch_size - if getattr(args, 'eval_accumulation_steps', None) is None and ga_steps is not None: args.eval_accumulation_steps = ga_steps - fp16_full_eval = getattr(args, 'fp16_full_eval', False) - if type(fp16_full_eval) is not bool: fp16_full_eval = False - bf16_full_eval = getattr(args, 'bf16_full_eval', False) - if type(bf16_full_eval) is not bool: bf16_full_eval = False - if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True - if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False - if force_float32: - args.bf16_full_eval = False - args.fp16_full_eval = False - elif os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16': - args.bf16_full_eval = True - args.fp16_full_eval = False - elif not bf16_full_eval and not fp16_full_eval: - args.bf16_full_eval = args.bf16 - args.fp16_full_eval = args.fp16 - if tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) != _unsloth_precision_requested: - from unsloth.models.loader_utils import sync_deepspeed_mixed_precision - sync_deepspeed_mixed_precision(args) - if model is not None: - from unsloth.models.loader_utils import apply_deepspeed_zero3_model_setup - apply_deepspeed_zero3_model_setup(model) - if model is not None: - _warnings_issued = getattr(model, 'warnings_issued', None) - if _warnings_issued is None: - model.warnings_issued = {} - elif not isinstance(_warnings_issued, dict): - try: - model.warnings_issued = dict(_warnings_issued) - except Exception: - model.warnings_issued = {} - if 'max_seq_length' not in locals() and not hasattr(args, 'max_seq_length'): - pass - else: - model_max_seq_length = getattr(model, 'max_seq_length', None) - args_max_seq_length = getattr(args, 'max_seq_length', None) - if args_max_seq_length is None and model_max_seq_length is not None: - max_seq_length = model.max_seq_length - if hasattr(args, 'max_seq_length'): args.max_seq_length = max_seq_length - elif args_max_seq_length is not None and model_max_seq_length is not None: - if args_max_seq_length > model_max_seq_length: - print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but ' - 'the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.') - args.max_seq_length = model_max_seq_length - if model is not None and hasattr(model, 'for_training'): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right' - if 'processing_class' in locals(): - if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right' - if hasattr(processing_class, 'tokenizer') and hasattr(processing_class.tokenizer, 'padding_side'): processing_class.tokenizer.padding_side = 'right' - other_metrics = [] - - from unsloth_zoo.logging_utils import PatchRLStatistics - PatchRLStatistics('sdft_trainer', other_metrics) - - # [TODO] Fix up DataParallel multiplying batch sizes - # [TODO] DDP works, but DP seems to not work? [TODO] - if getattr(args, "parallel_mode", None) == ParallelMode.NOT_DISTRIBUTED and args.n_gpu > 1: - if getattr(args, "_n_gpu", 1) != 1: - args._n_gpu = 1 - if "model" in locals(): - try: - from unsloth.models._gc_config_default import honour_gradient_checkpointing_off - honour_gradient_checkpointing_off(args, model) - except Exception: - pass - if "model" in locals() and hasattr(model, "for_training"): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if "model" in locals(): - try: - from unsloth.models._ddp_defaults import unsloth_default_ddp_find_unused_parameters - unsloth_default_ddp_find_unused_parameters(args, model) - except Exception: - pass - super().__init__( - model = model, - args = args, - train_dataset = train_dataset, - eval_dataset = eval_dataset, - processing_class = processing_class, - callbacks = callbacks, - peft_config = peft_config,**kwargs) - # A tensor parallel model (tp_native.tp_trainer_setup) stays on its training path: for_inference would - # route the decoder layers onto decode paths outside the tensor parallel regions. - if "model" in locals() and hasattr(model, "for_inference") and not getattr(self, "_unsloth_tensor_parallel", False): - model.for_inference() - if hasattr(self, 'neftune_hook_handle'): - self.neftune_hook_handle.remove() - if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle - if getattr(args, 'neftune_noise_alpha', None) is not None: - model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha - pass - if hasattr(self, 'accelerator'): - scaler = self.accelerator.scaler - current_model = model - while hasattr(current_model, 'model'): - current_model.accelerator_scaler = scaler - current_model = current_model.model - current_model.accelerator_scaler = scaler - pass - if hasattr(self, 'train'): - self.train = MethodType(prepare_for_training_mode(self.__class__.train), self) - pass - if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'): - _vllm_tok = self.llm.get_tokenizer() - _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None) - if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None: - _vllm_tok.chat_template = _pc.chat_template - pass - if getattr(self, 'aux_loss_enabled', False) and hasattr(getattr(self, 'model', None), 'config'): - _text_config = self.model.config - if hasattr(_text_config, 'get_text_config'): _text_config = _text_config.get_text_config() - _n_experts = [getattr(_text_config, _k) for _k in ('num_local_experts', 'num_experts', 'n_routed_experts', 'moe_num_experts') if isinstance(getattr(_text_config, _k, None), int)] - if _n_experts and all(_n <= 1 for _n in _n_experts): - self.aux_loss_enabled = False - _text_config.output_router_logits = False - pass - -pass - - -if hasattr(logger, "addFilter"): - import logging - class HideLoggingMessage(logging.Filter): - def __init__(self, text): self.text = text - def filter(self, x): return not (self.text in x.getMessage()) - pass - logger.addFilter(HideLoggingMessage("`use_cache=True`")) - diff --git a/unsloth_compiled_cache/UnslothSFTTrainer.py b/unsloth_compiled_cache/UnslothSFTTrainer.py deleted file mode 100644 index c36ec12036e97c95777e9308f1f4c4f20937ece2..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/UnslothSFTTrainer.py +++ /dev/null @@ -1,2922 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from unsloth_zoo.temporary_patches.common import torch_compile -from unsloth_zoo.temporary_patches.common import _maybe_compile -import functools -import transformers -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from trl.trainer.sft_trainer import (Any, AutoProcessor, BaseTunerLayer, BitsAndBytesConfig, Callable, DataCollator, DataCollatorForLanguageModeling, DataCollatorForVisionLanguageModeling, Dataset, DatasetDict, EvalPrediction, FLASH_ATTENTION_VARIANTS, IterableDataset, IterableDatasetDict, Path, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, SFTConfig, SFTTrainer, TrainerCallback, TrainingArguments, Version, _BaseTrainer, _is_package_version_below, _tokenize, accelerate, add_fused_lm_head, apply_chat_template, clone_chat_template, contextlib, create_model_from_path, dataclass, defaultdict, enable_selective_activation_checkpointing, get_act_offloading_ctx_manager, get_config_model_id, get_peft_model, get_training_chat_template, has_generation_markers, is_chat_template_stop_token_trained, is_conversational, is_peft_available, is_peft_model, json, logger, os, pack_dataset, pad, prepare_multimodal_messages, set_seed, torch, transformers, warnings, Any, AutoProcessor, BaseTunerLayer, BitsAndBytesConfig, Callable, DataCollator, DataCollatorForLanguageModeling, DataCollatorForVisionLanguageModeling, Dataset, DatasetDict, EvalPrediction, FLASH_ATTENTION_VARIANTS, IterableDataset, IterableDatasetDict, PeftConfig, PeftModel, PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin, SFTConfig, SFTTrainer, TrainerCallback, TrainingArguments, Version, accelerate, add_fused_lm_head, clone_chat_template, contextlib, create_model_from_path, defaultdict, enable_selective_activation_checkpointing, get_act_offloading_ctx_manager, get_config_model_id, get_peft_model, get_training_chat_template, has_generation_markers, is_chat_template_stop_token_trained, is_conversational, is_peft_available, is_peft_model, logger, os, pad, set_seed, torch, transformers, warnings, Callable, DataCollator, DataCollatorForLanguageModeling, Dataset, IterableDataset, apply_chat_template, is_conversational, json, os, pack_dataset, pad, prepare_multimodal_messages, transformers, warnings, DataCollator, DataCollatorForLanguageModeling, Dataset, os, PeftModel, PreTrainedModel, is_peft_available, logger, os, torch, add_fused_lm_head, os, torch) - - -import os -import math -import logging -from typing import * -from dataclasses import dataclass, field -from packaging.version import Version -import torch -import numpy as np -from contextlib import nullcontext -from torch.nn import functional as F -import inspect -from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling -from transformers.training_args import ParallelMode -from unsloth_zoo.device_type import DEVICE_TYPE, DEVICE_TYPE_TORCH, device_synchronize - -# Wrap trainer with padding to right and enable training mode -import functools -from types import MethodType -try: - from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers -except: - def reset_unsloth_gradient_checkpointing_buffers(): pass -# Canonical reset lives in unsloth.models._utils so the SFT auto-packing wrapper and the plain -# Trainer loop can import the same helper; fall back to a no-op only if it can't be imported. -try: - from unsloth.models._utils import _unsloth_reset_stray_compile_cache -except Exception: - def _unsloth_reset_stray_compile_cache(self): pass -try: - from unsloth.models._utils import _unsloth_dataset_column_names -except Exception: - def _unsloth_dataset_column_names(dataset): return dataset.column_names -# Drops/renames config arguments the installed TRL no longer accepts, so a -# script pinned to an older TRL keeps working after an upgrade. Falls back to -# the historical raw passthrough so this can never break trainer construction. -try: - from unsloth.models.rl_config_compat import filter_config_init_kwargs as _unsloth_filter_config_init_kwargs - # A cache file generated here can be imported by an older Unsloth whose filter - # predates `mirrored_from`, so drop the argument rather than raise TypeError. - if "mirrored_from" not in inspect.signature(_unsloth_filter_config_init_kwargs).parameters: - _unsloth_filter_config_init_kwargs_old = _unsloth_filter_config_init_kwargs - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): - return _unsloth_filter_config_init_kwargs_old(config_class, kwargs) -except Exception: - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): return kwargs -def prepare_for_training_mode(f): - @functools.wraps(f) - def wrapper(self, *args, **kwargs): - # Drop any torch.compile graph cache poisoned by a stray pre-train forward. - try: - _unsloth_reset_stray_compile_cache(self) - except Exception: - pass - # Finish the previous W&B run if this is a subsequent train() call. - # We do this at the START of train() (not the end) so that - # evaluate() / log() still work after train() completes. - # HF's WandbCallback.setup() will call wandb.init() for the new run. - # See: https://github.com/unslothai/unsloth/issues/3954 - if getattr(self, '_unsloth_training_completed', False): - try: - import wandb - if wandb.run is not None: - wandb.finish() - # Reset HF's WandbCallback so it calls wandb.init() for the new run - for cb in self.callback_handler.callbacks: - if type(cb).__name__ == 'WandbCallback': - cb._initialized = False - break - except: - pass - # Enable training mode - _was_training = None - # Restore the GC mode the model was configured with at setup; fall back to - # the training args only when it wasn't recorded (issue #4735). Use hasattr, - # not a None sentinel, so a deliberately-recorded None is restored verbatim. - _model = getattr(self, 'model', None) - if hasattr(_model, '_unsloth_gradient_checkpointing'): - use_gc = _model._unsloth_gradient_checkpointing - else: - use_gc = getattr(self.args, 'gradient_checkpointing', True) - if hasattr(self, 'model') and hasattr(self.model, "training"): - _was_training = self.model.training - if hasattr(self, 'model') and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - output = f(self, *args, **kwargs) - # Restore previous mode when possible - if hasattr(self, 'model') and hasattr(self.model, "for_inference"): - if _was_training is False: - self.model.for_inference() - elif _was_training is True and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - # Reset gradient checkpointing buffers to free memory while staying ready for next run - try: - reset_unsloth_gradient_checkpointing_buffers() - except: - pass - # Mark that training completed so the next train() call can - # finish this W&B run before starting a new one - self._unsloth_training_completed = True - return output - return wrapper -pass - -torch_compile_options = { - "epilogue_fusion" : True, - "max_autotune" : False, - "shape_padding" : True, - "trace.enabled" : False, - "triton.cudagraphs" : False, -} - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_hidden_states_selective_log_softmax( - hidden_states: torch.Tensor, - lm_head: torch.Tensor, - index: torch.Tensor, - chunks: int = 4, - logit_scale_multiply: float = 0.0, - logit_scale_divide: float = 0.0, - logit_softcapping: float = 0.0, - temperature: float = 1.0, - # Rows per chunk cap. Read HERE, in the default, not in the body: the body - # is traced with fullgraph = True, and `os.environ` is an unsupported op - # there on torch 2.4 -- Dynamo raises - # torch._dynamo.exc.Unsupported: const method call bytes.decode - # from os._Environ.__getitem__, which the eager fallback does not catch - # (it only catches recompile-limit and disabled-hook breaks), so the very - # first call would die even with the variable unset. A default is evaluated - # once when the def runs, which is import time, outside any traced region. - # It is also a plain int argument, so Dynamo guards on it instead of - # constant-folding an unguarded read (2.7+ never notice a later change). - # A non-numeric value is ignored rather than raised on, so a typo cannot - # break the import. 0 keeps the previous chunk boundaries exactly. - max_rows_per_chunk: int = ( - int(os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip()) - if os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip().isdigit() - else 0 - ), - lm_head_bias: Optional[torch.Tensor] = None, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, - # Read in the default for the same Dynamo reason as max_rows_per_chunk. Pairs with vLLM's - # head_dtype=float32 that load_vllm requests under the same flag. - fp32_lm_head: bool = ( - os.environ.get("UNSLOTH_VLLM_FP32_LM_HEAD", "0").strip().lower() in ("1", "true", "on", "yes") - ), - # Older Unsloth scores old / ref log-probs here and the policy with the fused kernel: one scorer keeps ratio at 1. - fused_linear: bool = ( - os.environ.get("UNSLOTH_FUSED_LINEAR", "1") != "0" - and os.environ.get("UNSLOTH_GRPO_CHUNKED_LOGPROBS", "1") != "0" - ), -) -> torch.Tensor: - # All Unsloth Zoo code licensed under AGPL3 - from unsloth_zoo.fused_losses import chunked_logprob as _chunked_logprob - if (fused_linear and _chunked_logprob._FUSED_LINEAR and not fp32_lm_head - and support_top_k == 0 and support_top_p == 1.0 and support_min_p == 0.0 - and support_excluded_ids is None): - return _chunked_logprob.unsloth_chunked_logprobs( - hidden_states, lm_head, index, - temperature = temperature, - logit_scale_multiply = logit_scale_multiply, - logit_scale_divide = logit_scale_divide, - logit_softcapping = logit_softcapping, - lm_head_bias = lm_head_bias, - max_rows_per_chunk = max_rows_per_chunk if max_rows_per_chunk > 0 else None, - ).reshape((hidden_states.shape[0], hidden_states.shape[1])) - # Reshape on this tensor's own last dim: a no-op, so a wrong-width caller - # cannot have its row count silently rewritten and instead fails at the - # matmul below, which prints both operands. Do not swap in a bare - # torch._check: it reports only "Expected cond to be True", naming neither - # operand, and Dynamo rejects a message-carrying one. Callers dispatch on - # the width first -- see `compute_logprobs_chunk`, the packed path and - # `_pg_grad_forward`. - flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) - flat_index = index.reshape(-1) - - # Each chunk materialises rows x vocab logits and then a float32 copy of - # them, all on the device holding the output head. With a large vocabulary - # and a fixed chunk count that grows with the batch, so the peak scales with - # the batch rather than staying bounded. max_rows_per_chunk caps the rows - # per chunk instead, which is pure loop splitting: more, smaller chunks, - # same concatenated result. 0 (the default) keeps the previous chunk - # boundaries exactly. - if max_rows_per_chunk > 0: - n_rows = flat_hidden_states.shape[0] - chunks = max(chunks, -(-n_rows // max_rows_per_chunk)) - chunks = min(chunks, max(n_rows, 1)) - - # A one-row chunk's head gradient is a K=1 matmul Inductor mis-lowers under dynamic shapes (illegal - # memory access or NaN head gradients, torch 2.13), so keep chunks at 2+ rows when the head trains. - # Exactly `chunks` row ranges, so the unrolled loop length never depends on T: torch.chunk - # returns ceil(T / ceil(T / chunks)) pieces and recompiled per new count; tensor_split guards - # on T % chunks. At most rows // 2 chunks when the head trains keeps every piece at 2+ rows. - lm_head_grad = torch.is_grad_enabled() and lm_head.requires_grad - if lm_head_grad and flat_hidden_states.shape[0] < 2 * chunks: - chunks = max(flat_hidden_states.shape[0] // 2, 1) - - n_rows = flat_hidden_states.shape[0] - chunks = max(1, min(chunks, n_rows)) - bounds = [(i * n_rows) // chunks for i in range(chunks + 1)] - chunked_hidden_states = [flat_hidden_states[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - - # Cast once, outside the loop, so every chunk shares one fp32 head copy (also for autograd). - if fp32_lm_head: - lm_head = lm_head.to(torch.float32) - if lm_head_bias is not None: - lm_head_bias = lm_head_bias.to(torch.float32) - - all_per_token_logps = [] - - for chunk_hidden_states, chunk_index in zip(chunked_hidden_states, chunked_index): - # When the model is dispatched over several devices, the output head can - # sit on a different one from the hidden states, because accelerate - # places the tail of the model on the last device it fills. Co-locate on - # the head's device before the matmul, otherwise this raises - # Unhandled FakeTensor Device Propagation for aten.mm.default, - # found two different devices cuda:0, cuda:1 - # On a single device every .to() here is a no-op and the result is - # bit-identical to before. - chunk_hidden_states = chunk_hidden_states.to(device = lm_head.device, dtype = lm_head.dtype) - chunk_index = chunk_index.to(lm_head.device) - if fp32_lm_head: - # The GRPO loss runs this under bf16 / fp16 autocast, which would downcast the matmul. - with torch.autocast(device_type = chunk_hidden_states.device.type, enabled = False): - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - elif lm_head_bias is None: - chunk_logits = chunk_hidden_states @ lm_head.t() - else: - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - - if logit_scale_multiply != 0.0: - chunk_logits = chunk_logits * logit_scale_multiply - if logit_scale_divide != 0.0: - chunk_logits = chunk_logits / logit_scale_divide - if logit_softcapping != 0.0: - chunk_logits = logit_softcapping * torch.tanh(chunk_logits / logit_softcapping) - - chunk_logits = chunk_logits.to(torch.float32) - - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - per_token_logps = _truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - ) - else: - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - # Return to the caller's device so the concatenation below and every - # downstream consumer see the device they started on. - all_per_token_logps.append(per_token_logps.to(hidden_states.device)) - - all_per_token_logps = torch.concat(all_per_token_logps) - - all_per_token_logps = all_per_token_logps.reshape((hidden_states.shape[0], hidden_states.shape[1])) - return all_per_token_logps - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_selective_log_softmax( - logits, - index, - temperature: float = 1.0, - chunks: int = 4, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, -): - # Fixed `chunks` row ranges; see chunked_hidden_states_selective_log_softmax. - flat_logits = logits.reshape(-1, logits.shape[-1]) - flat_index = index.reshape(-1) - # Fewer rows than chunks: one row per range, as torch.chunk did, never empty ranges. - chunks = max(1, min(chunks, flat_logits.shape[0])) - bounds = [(i * flat_logits.shape[0]) // chunks for i in range(chunks + 1)] - chunked_logits = [flat_logits[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - all_per_token_logps = [] - # Per-chunk selective_log_softmax. - for chunk_logits, chunk_index in zip(chunked_logits, chunked_index): - chunk_logits = chunk_logits.to(torch.float32) - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - all_per_token_logps.append(_truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - )) - continue - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - all_per_token_logps.append(per_token_logps) - pass - all_per_token_logps = torch.concat(all_per_token_logps) - all_per_token_logps = all_per_token_logps.reshape((logits.shape[0], logits.shape[1])) - return all_per_token_logps - -def calculate_pad_tokens_in_prompt( - input_ids: torch.Tensor, - logits_to_keep: int, - pad_token_id: int -) -> torch.Tensor: - """Count left-padded tokens per sequence, e.g. [pad, pad, pad, cat] -> 3.""" - if logits_to_keep >= input_ids.shape[1]: - raise ValueError("logits_to_keep must be smaller than the sequence length.") - - prompt_section = input_ids[:, :-logits_to_keep] - - padding_mask = (prompt_section == pad_token_id) - - pad_token_counts = padding_mask.sum(dim=1) - - return pad_token_counts - -def create_completion_attention_mask( - completion_input_ids: torch.Tensor, - left_pad_tokens_per_prompt: torch.Tensor, - max_left_pad: int, - pad_token_id: int -) -> torch.Tensor: - """Build a completion mask that zeros leading prompt and trailing pad tokens. - - For [p,p,p,c,c,c,pad,pad,pad] (p=sliced prompt, c=completion, pad=padding) - this returns [0,0,0,1,1,1,0,0,0]. - """ - batch_size, completion_len = completion_input_ids.shape - device = completion_input_ids.device - - num_tokens_to_mask = max_left_pad - left_pad_tokens_per_prompt - - indices = torch.arange(completion_len, device=device).unsqueeze(0) - shift_mask = indices >= num_tokens_to_mask.unsqueeze(1) - - non_padding_mask = (completion_input_ids != pad_token_id) - - final_mask = shift_mask & non_padding_mask - - return final_mask - -def left_pack_padding(tensor: torch.Tensor, pad_id: int) -> torch.Tensor: - """Move all padding tokens in each sequence to the right.""" - mask = (tensor != pad_id) - # stable=True since the binary mask is unordered. - sorted_indices = torch.argsort(mask, dim=1, descending=True, stable=True) - packed_tensor = torch.gather(tensor, 1, sorted_indices) - return packed_tensor - -def align_logprobs_with_mask( - logprob_tensor: torch.Tensor, - attention_mask: torch.Tensor, - pad_value: float = 0.0 -) -> torch.Tensor: - """Align a log probability tensor with a given attention mask.""" - - device = logprob_tensor.device - batch_size, logprob_seq_len = logprob_tensor.shape - mask_seq_len = attention_mask.shape[1] - - padded_logprobs = torch.full( - attention_mask.shape, - fill_value=pad_value, - dtype=logprob_tensor.dtype, - device=device - ) - - left_pad_counts = torch.argmax(attention_mask, dim=1) - - cols = torch.arange(logprob_seq_len, device=device) - dest_indices = left_pad_counts.unsqueeze(1) + cols - - # Destination row indices, shape [batch_size, logprob_seq_len]. - row_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand_as(dest_indices) - - # Keep only in-bounds destinations, then scatter via advanced indexing. - valid_mask = dest_indices < mask_seq_len - valid_rows = row_indices[valid_mask] - valid_cols = dest_indices[valid_mask] - valid_vals = logprob_tensor[valid_mask] - padded_logprobs[valid_rows, valid_cols] = valid_vals - - return padded_logprobs - -def align_completion_tool_mask( - tool_mask: torch.Tensor, - completion_mask: torch.Tensor, -) -> torch.Tensor: - """Align a raw completion-length tool/env mask with Unsloth's repacked loss mask.""" - if tool_mask is None: - return completion_mask - if tool_mask.shape[0] != completion_mask.shape[0]: - raise ValueError("tool_mask batch size must match completion_mask batch size.") - - tool_mask = tool_mask.to(device=completion_mask.device) - if tool_mask.shape == completion_mask.shape: - aligned_tool_mask = tool_mask - else: - aligned_tool_mask = align_logprobs_with_mask( - tool_mask, - completion_mask, - pad_value=0, - ) - return completion_mask * aligned_tool_mask.to(dtype=completion_mask.dtype) - -def autotune_batch_and_chunks( - total_input_rows, - seq_len, - hidden_size, - vocab_size, - dtype_bytes=16, - multiplier=None -): - if multiplier is None: - final_m = max(4, seq_len // 4096) - else: - final_m = multiplier - - if torch.cuda.is_available(): - free_bytes, _ = torch.cuda.mem_get_info() - limit_gb = (free_bytes / (1024**3))*.80 - elif hasattr(torch, "xpu") and torch.xpu.is_available(): - # XPU: estimate free memory as total - reserved. - total_mem = torch.xpu.get_device_properties(0).total_memory - reserved_mem = torch.xpu.memory_reserved() - free_bytes = total_mem - reserved_mem - limit_gb = (free_bytes / (1024**3)) * 0.80 - else: - # Fallback: assume 8GB available. - limit_gb = 8.0 - - bytes_to_gb = 1024**3 - - b_vals = torch.arange(total_input_rows, 0, -1, device='cpu', dtype=torch.float32) - - hidden_gb = (b_vals * seq_len * hidden_size * dtype_bytes) / bytes_to_gb - - base_logits = ((b_vals/total_input_rows) * b_vals * seq_len * vocab_size * dtype_bytes) / bytes_to_gb - logits_gb = base_logits / final_m - - total_mem_gb = hidden_gb + logits_gb - - valid_mask = total_mem_gb <= limit_gb - valid_indices = torch.nonzero(valid_mask, as_tuple=False) - - if valid_indices.shape[0] == 0: - #This means your GPU will OOM - # Capped at the row count: unsloth's no-grad pass divides rows by this without max(1, ...). - return max(1, min(4, total_input_rows)), final_m - - best_idx = valid_indices[0].item() - final_b = int(b_vals[best_idx].item()) - - return final_b, final_m - -def sanitize_logprob(logprob): - """Local port of trl.scripts.vllm_serve.sanitize_logprob. - Filters NaN logprobs from vLLM outputs.""" - value = logprob.logprob - if math.isnan(value): - logging.getLogger(__name__).warning( - f"Generated NaN logprob, token logprob '{logprob}' will be ignored" - ) - return None - return value -@dataclass -class UnslothSFTConfig(SFTConfig): - """ - - Configuration class for the [`SFTTrainer`]. - - This class includes only the parameters that are specific to SFT training. For a full list of training arguments, - please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may - differ from those in [`~transformers.TrainingArguments`]. - - Using [`~transformers.HfArgumentParser`] we can turn this class into - [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the - command line. - - Parameters: - > Parameters that control the model - - model_init_kwargs (`dict[str, Any]`, *optional*): - Keyword arguments for [`~transformers.AutoModelForCausalLM.from_pretrained`], used when the `model` - argument of the [`SFTTrainer`] is provided as a string. The `revision` value is also used when loading the - processing class. - trust_remote_code (`bool`, *optional*, defaults to `False`): - Whether to allow loading models and tokenizers that ship custom Python code from the Hub. Forwarded to - [`~transformers.AutoModelForCausalLM.from_pretrained`] and [`~transformers.AutoProcessor.from_pretrained`]. - router_aux_loss_coef (`float`, *optional*): - Coefficient of the load-balancing auxiliary loss for Mixture-of-Experts (MoE) models, added to the training - loss with this weight. When not set, the value declared by the model config is used. Set to `0.0` to - disable it. Fails when used with a non-MoE model. - chat_template_path (`str`, *optional*): - If specified, sets the model's chat template. This can either be the path to a tokenizer (local directory - or Hugging Face Hub model) or a direct path to a Jinja template file. When using a Jinja file, you must - ensure that any special tokens referenced in the template are added to the tokenizer and that the model's - embedding layer is resized accordingly. - - > Parameters that control the data preprocessing - - dataset_text_field (`str`, *optional*, defaults to `"text"`): - Name of the column that contains text data in the dataset. - dataset_kwargs (`dict[str, Any]`, *optional*): - Dictionary of optional keyword arguments for the dataset preparation. The only supported key is - `skip_prepare_dataset`. When the dataset contains images, `skip_prepare_dataset` is automatically treated - as `True` regardless of the provided value, since preprocessing is done on the fly. - dataset_num_proc (`int`, *optional*): - Number of processes to use for processing the dataset. - eos_token (`str`, *optional*): - Token used to indicate the end of a turn or sequence. If `None`, it defaults to - `processing_class.eos_token`. - max_length (`int` or `None`, *optional*, defaults to `1024`): - Maximum length of the tokenized sequence. Sequences longer than `max_length` are truncated from the left or - right depending on `truncation_mode`. If `None`, no truncation is applied. When packing is enabled, this - value sets the sequence length. - truncation_mode (`str`, *optional*, defaults to `"keep_start"`): - Truncation mode to use when the sequence exceeds `max_length`. The only supported value is `"keep_start"`. - The `"keep_end"` value is deprecated and will be removed in v2.0.0. - shuffle_dataset (`bool`, *optional*, defaults to `False`): - Whether to shuffle the dataset. - packing (`bool`, *optional*, defaults to `False`): - Whether to group multiple sequences into fixed-length blocks to improve computational efficiency and reduce - padding. Uses `max_length` to define sequence length. - packing_strategy (`str`, *optional*, defaults to `"bfd"`): - Strategy for packing sequences. Can be `"bfd"` (best-fit decreasing, truncates overflow), `"bfd_split"` - (best-fit decreasing, splits overflow sequences), or `"wrapped"` (aggressive, cuts mid-sequence). - padding_free (`bool`, *optional*, defaults to `False`): - Whether to perform forward passes without padding by flattening all sequences in the batch into a single - continuous sequence. This reduces memory usage by eliminating padding overhead. Currently, this is only - supported with the FlashAttention 2 or 3, which can efficiently handle the flattened batch structure. When - packing is enabled with strategy `"bfd"`, padding-free is enabled, regardless of the value of this - parameter. - pad_to_multiple_of (`int`, *optional*): - If set, the sequences will be padded to a multiple of this value. - eval_packing (`bool`, *optional*): - Whether to pack the eval dataset. If `None`, uses the same value as `packing`. - - > Parameters that control the training - - completion_only_loss (`bool`, *optional*): - Whether to compute loss only on the completion part of the sequence. If set to `True`, loss is computed - only on the completion, which is supported only for [prompt-completion](#prompt-completion) datasets. If - `False`, loss is computed on the entire sequence. If `None` (default), the behavior depends on the dataset: - loss is computed on the completion for [prompt-completion](#prompt-completion) datasets, and on the full - sequence for [language modeling](#language-modeling) datasets. - assistant_only_loss (`bool`, *optional*, defaults to `False`): - Whether to compute loss only on the assistant part of the sequence. If set to `True`, loss is computed only - on the assistant responses, which is supported only for [conversational](#conversational) datasets. If - `False`, loss is computed on the entire sequence. - loss_type (`str`, *optional*, defaults to `"nll"`): - Type of loss to use. Possible values are: - - - `"nll"`: standard negative log-likelihood. - - `"dft"`: Dynamic Fine-Tuning, as described in [this paper](https://huggingface.co/papers/2508.05629). - - `"chunked_nll"`: deprecated alias of `"nll"`, which computes the loss the same way. It will be removed in - v2.0.0. - - activation_offloading (`bool`, *optional*, defaults to `False`): - Whether to offload the activations to the CPU. - - > Deprecated parameters - - pad_token: - - - - Parameter `pad_token` is deprecated and will be removed in version v2.0.0. Set `tokenizer.pad_token` - directly and pass it as `processing_class` to the trainer instead. - - - - > [!NOTE] - > These parameters have default values different from [`~transformers.TrainingArguments`]: - > - `logging_steps`: Defaults to `10` instead of `500`. - > - `gradient_checkpointing`: Defaults to `True` instead of `False`. - > - `bf16`: Defaults to `True` if `fp16` is not set, instead of `False`. - > - `learning_rate`: Defaults to `2e-5` instead of `5e-5`. - - """ - _unsloth_accepts_adamw_flash = True - vllm_sampling_params: Optional[Any] = field( - default = None, - metadata = {'help': 'vLLM SamplingParams'}, - ) - unsloth_num_chunks : Optional[int] = field( - default = -1, - metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'}, - ) - unsloth_logit_chunk_multiplier : Optional[int] = field( - default = None, - metadata = {'help': 'Multiplier for chunked logit computations.'}, - ) - unsloth_grpo_mini_batch : Optional[int] = field( - default = None, - metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'}, - ) - max_seq_length : Optional[int] = field( - default = None, - metadata = {'help': 'Maximum sequence length to truncate to.'}, - ) - context_parallel_size : Optional[int] = field( - default = 1, - metadata = {'help': 'Ranks per context parallel group. 1 disables it.'}, - ) - context_parallel_backend : Optional[str] = field( - default = 'ring', - metadata = {'help': 'Context parallel attention backend: ring or ulysses (keeps packing).'}, - ) - context_parallel_rotate_method : Optional[str] = field( - default = 'allgather', - metadata = {'help': 'Ring KV rotation: allgather or alltoall.'}, - ) - def __init__( - self, - output_dir = None, - per_device_train_batch_size = 4, - num_train_epochs = 3.0, - max_steps = -1, - learning_rate = 5e-05, - lr_scheduler_type = 'linear', - lr_scheduler_kwargs = None, - warmup_steps = 0.1, - optim = 'adamw_8bit', - optim_args = None, - weight_decay = 0.001, - adam_beta1 = 0.9, - adam_beta2 = 0.999, - adam_epsilon = 1e-08, - optim_target_modules = None, - gradient_accumulation_steps = 2, - average_tokens_across_devices = True, - max_grad_norm = 1.0, - label_smoothing_factor = 0.0, - bf16 = False, - fp16 = False, - bf16_full_eval = False, - fp16_full_eval = False, - tf32 = None, - gradient_checkpointing = True, - gradient_checkpointing_kwargs = None, - torch_compile = False, - torch_compile_backend = None, - torch_compile_mode = None, - use_liger_kernel = False, - liger_kernel_config = None, - use_cache = False, - neftune_noise_alpha = None, - torch_empty_cache_steps = 250, - auto_find_batch_size = False, - logging_strategy = 'steps', - logging_steps = 1, - logging_first_step = False, - log_on_each_node = True, - logging_nan_inf_filter = False, - include_num_input_tokens_seen = False, - log_level = 'passive', - log_level_replica = 'warning', - disable_tqdm = None, - report_to = 'none', - run_name = None, - project = 'huggingface', - trackio_space_id = 'trackio', - eval_strategy = 'no', - eval_steps = None, - eval_delay = 0, - per_device_eval_batch_size = 4, - prediction_loss_only = False, - eval_on_start = False, - eval_do_concat_batches = True, - eval_use_gather_object = False, - eval_accumulation_steps = 2, - batch_eval_metrics = False, - save_only_model = False, - save_strategy = 'steps', - save_steps = 500, - save_on_each_node = False, - save_total_limit = None, - enable_jit_checkpoint = False, - push_to_hub = False, - hub_token = None, - hub_private_repo = None, - hub_model_id = None, - hub_strategy = 'every_save', - hub_always_push = False, - hub_revision = None, - load_best_model_at_end = False, - metric_for_best_model = None, - greater_is_better = None, - ignore_data_skip = False, - restore_callback_states_from_checkpoint = False, - full_determinism = False, - seed = 3407, - data_seed = None, - use_cpu = False, - accelerator_config = None, - parallelism_config = None, - dataloader_drop_last = False, - dataloader_num_workers = 0, - dataloader_pin_memory = True, - dataloader_persistent_workers = False, - dataloader_prefetch_factor = None, - remove_unused_columns = True, - label_names = None, - train_sampling_strategy = 'random', - length_column_name = 'length', - ddp_find_unused_parameters = None, - ddp_bucket_cap_mb = None, - ddp_broadcast_buffers = None, - ddp_backend = None, - ddp_timeout = 1800, - fsdp = None, - fsdp_config = None, - deepspeed = None, - debug = '', - skip_memory_metrics = True, - do_train = False, - do_eval = False, - do_predict = False, - resume_from_checkpoint = None, - warmup_ratio = None, - logging_dir = None, - local_rank = -1, - model_init_kwargs = None, - trust_remote_code = False, - router_aux_loss_coef = None, - chat_template_path = None, - dataset_text_field = 'text', - dataset_kwargs = None, - dataset_num_proc = None, - eos_token = None, - max_length = 1024, - truncation_mode = 'keep_start', - shuffle_dataset = False, - packing = False, - packing_strategy = 'bfd', - padding_free = None, - pad_to_multiple_of = None, - eval_packing = None, - completion_only_loss = None, - assistant_only_loss = False, - loss_type = 'nll', - activation_offloading = False, - pad_token = None, - vllm_sampling_params = None, - unsloth_num_chunks = -1, - unsloth_logit_chunk_multiplier = None, - unsloth_grpo_mini_batch = None, - max_seq_length = None,context_parallel_size = 1,context_parallel_backend = 'ring',context_parallel_rotate_method = 'allgather', - **kwargs, - ): - if learning_rate < 1e-7: print(f'Unsloth: Your learning rate of `{learning_rate}` is too small and less than 1e-7! Consider increasing it, otherwise gradient updates will be close to 0!') - if learning_rate > 1: print(f'Unsloth: Your learning rate of `{learning_rate}` is way too larger > 1! Consider decreasing it to 1e-1, otherwise gradient updates will explode!') - _unsloth_fsdp2 = os.environ.get('FSDP_VERSION', '').strip() == '2' - _unsloth_fsdp_cfg = fsdp_config - if fsdp and isinstance(_unsloth_fsdp_cfg, str) and os.path.isfile(_unsloth_fsdp_cfg): - try: - import json as _unsloth_json - with open(_unsloth_fsdp_cfg, encoding = 'utf-8') as _unsloth_f: - _unsloth_fsdp_cfg = _unsloth_json.load(_unsloth_f) - except Exception: - _unsloth_fsdp_cfg = None - if fsdp and _unsloth_fsdp_cfg is None: - _unsloth_fsdp_cfg = {} - if fsdp and isinstance(_unsloth_fsdp_cfg, dict): - _unsloth_fsdp2 = _unsloth_fsdp2 or str(_unsloth_fsdp_cfg.get('fsdp_version', _unsloth_fsdp_cfg.get('version', 1))) == '2' - _unsloth_flash_adamw = str(getattr(optim, 'value', optim)).strip().lower() == 'adamw_flash' - if _unsloth_flash_adamw: optim = 'adamw_torch' - if _unsloth_fsdp2 and str(getattr(optim, 'value', optim)) == 'adamw_8bit': - print("Unsloth: FSDP2 shards parameters as DTensors, which bitsandbytes' 8-bit optimizers cannot step. Switching optim = 'adamw_8bit' to 'adamw_torch_fused'.") - optim = 'adamw_torch_fused' - if num_train_epochs is None: - num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override - if output_dir is None and save_strategy == 'steps' and save_steps == 500: - output_dir = 'unsloth_training_checkpoints' - save_strategy = 'no' - try: - from unsloth_zoo.dataset_num_proc import get_dataset_num_proc as _unsloth_get_dataset_num_proc - except Exception: - try: - from unsloth.dataset_num_proc import get_dataset_num_proc as _unsloth_get_dataset_num_proc - except Exception: - _unsloth_get_dataset_num_proc = None - if _unsloth_get_dataset_num_proc is not None: - dataset_num_proc = _unsloth_get_dataset_num_proc(dataset_num_proc, serial_as_none = False) - if os.environ.get('UNSLOTH_ENABLE_FLEX_ATTENTION', '0') == '1': - from unsloth_zoo.flex_attention import HAS_FLEX_ATTENTION - if HAS_FLEX_ATTENTION and pad_to_multiple_of is None: - from unsloth_zoo.flex_attention import FLEX_ATTENTION_BLOCK_SIZE - pad_to_multiple_of = FLEX_ATTENTION_BLOCK_SIZE - - if eval_steps is not None and eval_strategy != 'steps': - print(f'Unsloth: `eval_steps = {eval_steps}` is ignored because `eval_strategy` is {getattr(eval_strategy, "value", eval_strategy)!r}. Set `eval_strategy = "steps"` to evaluate every `eval_steps` steps.') - - - # One dict so the filter sees the mirrored parameters AND `**kwargs`: - # filtering kwargs alone would double-bind any argument TRL renamed, - # since the new name is itself a mirrored parameter. - _unsloth_config_arguments = dict( - output_dir = output_dir, - per_device_train_batch_size = per_device_train_batch_size, - num_train_epochs = num_train_epochs, - max_steps = max_steps, - learning_rate = learning_rate, - lr_scheduler_type = lr_scheduler_type, - lr_scheduler_kwargs = lr_scheduler_kwargs, - warmup_steps = warmup_steps, - optim = optim, - optim_args = optim_args, - weight_decay = weight_decay, - adam_beta1 = adam_beta1, - adam_beta2 = adam_beta2, - adam_epsilon = adam_epsilon, - optim_target_modules = optim_target_modules, - gradient_accumulation_steps = gradient_accumulation_steps, - average_tokens_across_devices = average_tokens_across_devices, - max_grad_norm = max_grad_norm, - label_smoothing_factor = label_smoothing_factor, - bf16 = bf16, - fp16 = fp16, - bf16_full_eval = bf16_full_eval, - fp16_full_eval = fp16_full_eval, - tf32 = tf32, - gradient_checkpointing = gradient_checkpointing, - gradient_checkpointing_kwargs = gradient_checkpointing_kwargs, - torch_compile = torch_compile, - torch_compile_backend = torch_compile_backend, - torch_compile_mode = torch_compile_mode, - use_liger_kernel = use_liger_kernel, - liger_kernel_config = liger_kernel_config, - use_cache = use_cache, - neftune_noise_alpha = neftune_noise_alpha, - torch_empty_cache_steps = torch_empty_cache_steps, - auto_find_batch_size = auto_find_batch_size, - logging_strategy = logging_strategy, - logging_steps = logging_steps, - logging_first_step = logging_first_step, - log_on_each_node = log_on_each_node, - logging_nan_inf_filter = logging_nan_inf_filter, - include_num_input_tokens_seen = include_num_input_tokens_seen, - log_level = log_level, - log_level_replica = log_level_replica, - disable_tqdm = disable_tqdm, - report_to = report_to, - run_name = run_name, - project = project, - trackio_space_id = trackio_space_id, - eval_strategy = eval_strategy, - eval_steps = eval_steps, - eval_delay = eval_delay, - per_device_eval_batch_size = per_device_eval_batch_size, - prediction_loss_only = prediction_loss_only, - eval_on_start = eval_on_start, - eval_do_concat_batches = eval_do_concat_batches, - eval_use_gather_object = eval_use_gather_object, - eval_accumulation_steps = eval_accumulation_steps, - batch_eval_metrics = batch_eval_metrics, - save_only_model = save_only_model, - save_strategy = save_strategy, - save_steps = save_steps, - save_on_each_node = save_on_each_node, - save_total_limit = save_total_limit, - enable_jit_checkpoint = enable_jit_checkpoint, - push_to_hub = push_to_hub, - hub_token = hub_token, - hub_private_repo = hub_private_repo, - hub_model_id = hub_model_id, - hub_strategy = hub_strategy, - hub_always_push = hub_always_push, - hub_revision = hub_revision, - load_best_model_at_end = load_best_model_at_end, - metric_for_best_model = metric_for_best_model, - greater_is_better = greater_is_better, - ignore_data_skip = ignore_data_skip, - restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint, - full_determinism = full_determinism, - seed = seed, - data_seed = data_seed, - use_cpu = use_cpu, - accelerator_config = accelerator_config, - parallelism_config = parallelism_config, - dataloader_drop_last = dataloader_drop_last, - dataloader_num_workers = dataloader_num_workers, - dataloader_pin_memory = dataloader_pin_memory, - dataloader_persistent_workers = dataloader_persistent_workers, - dataloader_prefetch_factor = dataloader_prefetch_factor, - remove_unused_columns = remove_unused_columns, - label_names = label_names, - train_sampling_strategy = train_sampling_strategy, - length_column_name = length_column_name, - ddp_find_unused_parameters = ddp_find_unused_parameters, - ddp_bucket_cap_mb = ddp_bucket_cap_mb, - ddp_broadcast_buffers = ddp_broadcast_buffers, - ddp_backend = ddp_backend, - ddp_timeout = ddp_timeout, - fsdp = fsdp, - fsdp_config = fsdp_config, - deepspeed = deepspeed, - debug = debug, - skip_memory_metrics = skip_memory_metrics, - do_train = do_train, - do_eval = do_eval, - do_predict = do_predict, - resume_from_checkpoint = resume_from_checkpoint, - warmup_ratio = warmup_ratio, - logging_dir = logging_dir, - local_rank = local_rank, - model_init_kwargs = model_init_kwargs, - trust_remote_code = trust_remote_code, - router_aux_loss_coef = router_aux_loss_coef, - chat_template_path = chat_template_path, - dataset_text_field = dataset_text_field, - dataset_kwargs = dataset_kwargs, - dataset_num_proc = dataset_num_proc, - eos_token = eos_token, - max_length = max_length, - truncation_mode = truncation_mode, - shuffle_dataset = shuffle_dataset, - packing = packing, - packing_strategy = packing_strategy, - padding_free = padding_free, - pad_to_multiple_of = pad_to_multiple_of, - eval_packing = eval_packing, - completion_only_loss = completion_only_loss, - assistant_only_loss = assistant_only_loss, - loss_type = loss_type, - activation_offloading = activation_offloading, - pad_token = pad_token,**kwargs) - super().__init__(**_unsloth_filter_config_init_kwargs(SFTConfig, _unsloth_config_arguments, mirrored_from = __class__)) - self.vllm_sampling_params = vllm_sampling_params - self.unsloth_num_chunks = unsloth_num_chunks - if unsloth_grpo_mini_batch is not None: - if self.generation_batch_size >= unsloth_grpo_mini_batch: - self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch - else: - raise ValueError( - f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, " - f"which is self.per_device_train_batch_size * gradient_accumulation_steps." - ) - self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier - self.max_seq_length = max_seq_length - self.context_parallel_size = context_parallel_size - self.context_parallel_backend = context_parallel_backend - self.context_parallel_rotate_method = context_parallel_rotate_method - # Unsloth: keep the reentrant checkpoint path - if getattr(self, 'gradient_checkpointing', False): - _gc_kwargs = getattr(self, 'gradient_checkpointing_kwargs', None) or {} - if _gc_kwargs.get('context_fn') is None and not _gc_kwargs.get('debug', False): - _gc_kwargs['use_reentrant'] = True - self.gradient_checkpointing_kwargs = _gc_kwargs - from unsloth.optimizers.flash_adamw import is_flash_adamw_marked, mark_flash_adamw, unmark_flash_adamw - if _unsloth_flash_adamw or is_flash_adamw_marked(self): - mark_flash_adamw(self) - else: - unmark_flash_adamw(self) - -pass - -class _UnslothSFTTrainer(_BaseTrainer): - """""" - - _tag_names = ["trl", "sft"] - _name = "SFT" - - def __init__( - self, - model: "str | PreTrainedModel | PeftModel", - args: SFTConfig | TrainingArguments | None = None, - data_collator: DataCollator | None = None, - train_dataset: Dataset | IterableDataset | None = None, - eval_dataset: Dataset - | IterableDataset - | DatasetDict - | IterableDatasetDict - | dict[str, Dataset | IterableDataset] - | None = None, - processing_class: PreTrainedTokenizerBase | ProcessorMixin | None = None, - compute_loss_func: Callable | None = None, - compute_metrics: Callable[[EvalPrediction], dict] | None = None, - callbacks: list[TrainerCallback] | None = None, - optimizers: tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None] = (None, None), - optimizer_cls_and_kwargs: tuple[type[torch.optim.Optimizer], dict[str, Any]] | None = None, - preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, - quantization_config: "BitsAndBytesConfig | None" = None, - peft_config: "PeftConfig | None" = None, - formatting_func: Callable[[dict], str] | None = None, - ): - # Args - if args is None: - model_name = model if isinstance(model, str) else get_config_model_id(model.config) - model_name = model_name.split("/")[-1] - args = SFTConfig(f"{model_name}-SFT") - elif isinstance(args, TrainingArguments) and not isinstance(args, SFTConfig): - dict_args = args.to_dict() - dict_args["hub_token"] = args.hub_token # to_dict hides the hub_token - if Version(transformers.__version__) < Version("5.0.0"): - dict_args.pop("push_to_hub_token", None) - args = UnslothSFTConfig(**dict_args) - - if train_dataset is None: - raise ValueError("`train_dataset` is required") - elif isinstance(train_dataset, IterableDataset): - # IterableDataset requires dispatch_batches=False because Accelerate's dispatch mode may try to concatenate - # batches from multiple processes, leading to mismatch errors. - if args.accelerator_config.dispatch_batches is True: - logger.warning( - "You are using an `IterableDataset` for training with `dispatch_batches=True`. `dispatch_batches` " - "is forced to `False` when using an `IterableDataset`. To remove this warning, unset " - "`dispatch_batches` in `SFTConfig` or set it to `False`." - ) - args.accelerator_config.dispatch_batches = False - elif not isinstance(train_dataset, (Dataset, list, tuple)): - raise TypeError( - f"`train_dataset` must be a `Dataset` or `IterableDataset`, got `{type(train_dataset).__name__}`." - ) - - # Model - # PEFT initializes the adapter weights randomly, so set_seed must be done before creating the model to ensure - # reproducibility. - set_seed(args.seed) - if isinstance(model, str): - model_init_kwargs = dict(args.model_init_kwargs or {}) # copy to avoid mutating model_init_kwargs - if quantization_config is not None: - if "quantization_config" in model_init_kwargs: - raise ValueError( - "You set `quantization_config` both as a trainer argument and in `args.model_init_kwargs`. " - "Please set it in only one place, preferably as a trainer argument." - ) - model_init_kwargs["quantization_config"] = quantization_config - # Distributed training requires device_map=None ["auto" fails] - if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]: - model_init_kwargs["device_map"] = None - model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) - model_revision = model_init_kwargs.get("revision") - model = create_model_from_path(model, **model_init_kwargs) - else: - model_revision = None - if args.model_init_kwargs is not None: - logger.warning( - "You passed `model_init_kwargs` to the `SFTConfig`, but your model is already instantiated. " - "The `model_init_kwargs` will be ignored." - ) - if quantization_config is not None: - logger.warning( - "You passed `quantization_config` to the trainer, but your model is already instantiated. The " - "`quantization_config` will be ignored." - ) - # Non-quantized models do not have the `is_loaded_in_{8,4}bit` attributes, whereas quantized models do - _is_quantized_model = getattr(model, "is_loaded_in_4bit", False) or getattr(model, "is_loaded_in_8bit", False) - - # Processing class - if processing_class is None: - processing_class = AutoProcessor.from_pretrained( - get_config_model_id(model.config), revision=model_revision, trust_remote_code=args.trust_remote_code - ) - - # Handle pad token for processors or tokenizers - if isinstance(processing_class, ProcessorMixin): - self._tokenizer = processing_class.tokenizer - self._is_vlm = True - elif isinstance(processing_class, PreTrainedTokenizerBase): - self._tokenizer = processing_class - self._is_vlm = False - else: - raise TypeError("The `processing_class` must be either a `PreTrainedTokenizerBase` or a `ProcessorMixin`") - - if args.eos_token is not None: - if args.eos_token not in self._tokenizer.get_vocab(): - raise ValueError( - f"The specified `eos_token` ('{args.eos_token}') is not found in the vocabulary of the given " - f"`processing_class` ({processing_class.__class__.__name__}). Ensure that the `eos_token` exists " - "in the vocabulary before using it as an EOS token." - ) - self._tokenizer.eos_token = args.eos_token - # The model must agree with the tokenizer on the eos token from construction, so mirror it onto the model - # configs. The generation config may hold several eos tokens, any of which halts generation, so the new - # one is added to the existing ones instead of replacing them. - model.config.get_text_config().eos_token_id = self._tokenizer.eos_token_id - eos_token_ids = model.generation_config.eos_token_id - if eos_token_ids is None: - eos_token_ids = [] - elif isinstance(eos_token_ids, int): - eos_token_ids = [eos_token_ids] - if self._tokenizer.eos_token_id not in eos_token_ids: - model.generation_config.eos_token_id = [self._tokenizer.eos_token_id, *eos_token_ids] - - if args.chat_template_path is not None: - if os.path.isfile(args.chat_template_path) and args.chat_template_path.endswith((".jinja", ".j2")): - with open(args.chat_template_path, encoding="utf-8") as chat_template_file: - processing_class.chat_template = chat_template_file.read() - added_tokens = [] - else: - model, processing_class, added_tokens = clone_chat_template( - model, processing_class, args.chat_template_path - ) - else: - added_tokens = [] - - # Vision dataset detection - dataset_sample = next(iter(train_dataset)) - self._is_vision_dataset = "image" in dataset_sample or "images" in dataset_sample - # Unsloth: override _is_vlm for VLM models that pass a bare tokenizer - if not self._is_vlm and self._is_vision_dataset: - _m = model - if hasattr(_m, "model"): _m = _m.model - if hasattr(getattr(_m, "config", None), "vision_config") or\ - _m.__class__.__name__.endswith("ForConditionalGeneration"): - self._is_vlm = True - if self._is_vision_dataset and not self._is_vlm: - raise ValueError( - "The dataset appears to be vision-related (contains 'image' or 'images' keys), but the provided " - "model does not seem to be a vision-language model. Please check your model and dataset." - ) - - if self._is_vision_dataset and args.packing: - raise ValueError( - "Packing is not supported for vision datasets. Please set `packing=False` in the SFTConfig." - ) - if self._is_vision_dataset and args.padding_free: - raise ValueError( - "Padding-free training is yet not supported for vision datasets. Please set `padding_free=False` in " - "the `SFTConfig`." - ) - if ( - self._is_vision_dataset - and args.assistant_only_loss - and Version(transformers.__version__) < Version("5.18.0.dev0") - ): - raise ValueError( - "Assistant-only loss for vision datasets requires transformers>=5.18.0. Please upgrade transformers " - "or set `assistant_only_loss=False` in the `SFTConfig`." - ) - if self._is_vision_dataset and args.max_length is not None and args.truncation_mode == "keep_end": - raise ValueError( - "truncation_mode='keep_end' is not supported for vision datasets. Image tokens reside " - "inside the prompt portion of the sequence; depending on the example, keep_end may silently " - "drop them, causing pixel_values to be forwarded to the model with no corresponding visual " - "tokens in input_ids. Use truncation_mode='keep_start' (the default) or set max_length=None." - ) - - # PEFT - if False: - if not is_peft_available(): - raise ImportError( - "You passed `peft_config` but the `peft` library is not installed. " - "Install it with `pip install trl[peft]`." - ) - if not isinstance(peft_config, PeftConfig): - raise TypeError( - f"`peft_config` must be a `peft.PeftConfig` instance (e.g. `peft.LoraConfig`), " - f"got {type(peft_config).__name__}." - ) - if is_peft_model(model): - raise ValueError( - "You passed a `PeftModel` instance together with a `peft_config` to the trainer. Please first merge " - "and unload the existing adapter, save the resulting base model, and then pass that base model along " - "with the new `peft_config` to the trainer." - ) - if added_tokens: - # Ensure that the added tokens are trainable - if peft_config.trainable_token_indices is None: - peft_config.trainable_token_indices = {"embed_tokens": added_tokens} - elif "embed_tokens" not in peft_config.trainable_token_indices: - peft_config.trainable_token_indices["embed_tokens"] = added_tokens - else: - peft_config.trainable_token_indices["embed_tokens"].extend(added_tokens) - # Ensure that the lm_head is trainable - if peft_config.modules_to_save is None or "lm_head" not in peft_config.modules_to_save: - logger.warning( - "Cloning chat template added new tokens to the tokenizer, but 'lm_head' is not in PEFT's " - "`modules_to_save`. As a result, the model may not learn to generate outputs with these new " - "tokens, leading to degraded generation quality. To fix this, add " - "`modules_to_save=['lm_head']` to your PEFT configuration." - ) - - if peft_config.modules_to_save is None: - peft_config.modules_to_save = ["lm_head"] - else: - peft_config.modules_to_save.append("lm_head") - # Create PEFT model - # ZeRO-3 + PEFT for non-quantized models: - # - PEFT's default autocast_adapter_dtype=True upcasts LoRA adapter params to fp32 even when the base model is bf16. - # - ZeRO-3's _allgather_params_coalesced allocates output buffers using the dtype of the first persistent parameter, - # so mixed-dtype persistent_parameters [bf16 base + fp32 LoRA] cause a TypeError on the first optimizer step. - # - Passing autocast_adapter_dtype=False keeps adapter params in the base model dtype [bf16], fixing the mismatch. - # - This is safe: the fp32 upcast is a QLoRA-specific concern [low-bit quantized base models], not needed for - # non-quantized bf16 training. - # - See: - # - TRL issue: https://github.com/huggingface/trl/issues/6089 - # - Upstream issue: https://github.com/deepspeedai/DeepSpeed/issues/8072 - get_peft_model_kwargs = {} - if args.deepspeed_plugin is not None and args.deepspeed_plugin.zero_stage == 3 and not _is_quantized_model: - get_peft_model_kwargs["autocast_adapter_dtype"] = False - model = get_peft_model(model, peft_config, **get_peft_model_kwargs) - - # PEFT + DeepSpeed ZeRO-3 requires reentrant checkpointing. For more details, see - # https://github.com/huggingface/trl/issues/2514#issuecomment-2692152703. - # Can be removed once https://github.com/deepspeedai/DeepSpeed/pull/8130 is merged and released. - if ( - is_peft_model(model) - and args.deepspeed_plugin is not None - and args.deepspeed_plugin.zero_stage == 3 - and args.gradient_checkpointing - ): - if (args.gradient_checkpointing_kwargs or {}).get("selective"): - raise ValueError( - "`gradient_checkpointing_kwargs={'selective': True}` is not supported with PEFT + DeepSpeed " - "ZeRO-3, which requires reentrant gradient checkpointing while SAC requires non-reentrant " - "checkpointing." - ) - args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} - use_reentrant = args.gradient_checkpointing_kwargs.get("use_reentrant") - if use_reentrant is False: - logger.warning( - "You are using PEFT with DeepSpeed ZeRO-3 and gradient checkpointing with `use_reentrant=False`. " - "`use_reentrant` is forced to `True` in this configuration to ensure correct training. To remove " - "this warning, unset `use_reentrant` in `gradient_checkpointing_kwargs` or set it to `True`." - ) - args.gradient_checkpointing_kwargs["use_reentrant"] = True - - # When using gradient checkpointing with PEFT, we need to enable input gradients. transformers.Trainer normally - # handles this, but a bug currently prevents it; see https://github.com/huggingface/transformers/issues/42489 - if is_peft_model(model) and args.gradient_checkpointing: - model.enable_input_require_grads() - - # When using QLoRA, the PEFT adapter weights are converted to bf16 to follow the recommendations from the - # original paper [see https://huggingface.co/papers/2305.14314, paragraph 3]. Normally, this can be done by - # passing `autocast_adapter_dtype=False` to `get_peft_model`, but this option is not yet supported for - # quantized models. See: https://github.com/huggingface/peft/issues/2889 - # The DoRA magnitude vector is excluded: unlike LoRA A/B, its optimizer updates can be smaller than bf16 can - # represent, silently freezing it, see #7268. - if False: - for name, param in model.named_parameters(): - if param.requires_grad and "lora_magnitude_vector" not in name: - param.data = param.data.to(torch.bfloat16) - - # Data collator - # BFD packing requires padding-free mode; otherwise, the collator outputs padded attention masks, causing - # FlashAttention to ignore position_ids and recompute them incorrectly from the padded attention mask. - self.padding_free = args.padding_free or (args.packing and args.packing_strategy in {"bfd", "bfd_split"}) - # A hub kernel can be requested with a revision and/or a kernel name [`repo_id@revision:kernel_name`], while - # the variants above are bare repo ids, so compare against the repo id alone. - attn_implementation = model.config._attn_implementation.split("@")[0].split(":")[0] - use_flash_attention = attn_implementation in FLASH_ATTENTION_VARIANTS - if self.padding_free: - if data_collator is not None: - raise ValueError("Passing a custom data collator is not supported when using padding-free.") - if args.packing and args.packing_strategy == "wrapped": - logger.warning( - "You are passing `padding_free=True` with the 'wrapped' packing strategy, which is not " - "recommended. Please refer to the documentation to understand why this is not recommended." - ) - if not use_flash_attention: - logger.warning( - "Padding-free training is enabled, but the attention implementation is not set to a supported " - "Flash Attention variant. Padding-free training flattens batches into a single sequence, and only " - "the following implementations are known to reliably support this: " - f"{', '.join(sorted(FLASH_ATTENTION_VARIANTS))}. Using other implementations may lead to " - "unexpected behavior. To ensure compatibility, set `attn_implementation` in the model " - "configuration to one of these supported options or verify that your attention mechanism can " - "handle flattened sequences." - ) - if args.assistant_only_loss and not is_conversational(dataset_sample): - raise ValueError( - "You set `assistant_only_loss=True`, but the dataset is not conversational. This option is only " - "supported for conversational datasets." - ) - - # When assistant_only_loss is enabled, swap in a training chat template with {% generation %} markers - # if the current template doesn't already have them. - if args.assistant_only_loss and not has_generation_markers(processing_class.chat_template): - try: - self.chat_template = get_training_chat_template(processing_class) - except ValueError: - if getattr(self, '_is_vision_dataset', False) or (getattr(args, 'dataset_kwargs', None) or {}).get('skip_prepare_dataset'): raise - self.chat_template = None - self._unsloth_assistant_mask_fallback = True - print('Unsloth: TRL has no training chat template for this tokenizer, so assistant_only_loss masks non-assistant tokens with train_on_responses_only markers.') - else: - self.chat_template = None - - # A template can define generation markers and still attribute the assistant's end-of-turn token to the next - # message, leaving it out of the assistant mask so the model is never trained to stop. - if args.assistant_only_loss and not getattr(self, '_unsloth_assistant_mask_fallback', False) and not is_chat_template_stop_token_trained( - processing_class, chat_template=self.chat_template - ): - logger.warning( - "The chat template does not include the assistant turn's end-of-turn token in the loss mask; " - "the model may not learn to stop. The training loss still looks healthy, so this usually only " - "surfaces at inference. Either set `assistant_only_loss=False` to train on the full sequence, " - "or edit the chat template so the end-of-turn token falls inside " - "`{% generation %}...{% endgeneration %}`." - ) - - # Decide whether to use completion-only loss: if not specified, then it is set to True if the dataset format - # is prompt-completion, and False if the dataset format is language modeling. - if args.completion_only_loss is None: - self.completion_only_loss = "prompt" in dataset_sample and "completion" in dataset_sample - else: - self.completion_only_loss = args.completion_only_loss - - if data_collator is None: - # Get the pad token: if not provided, use the one from the processing class or the eos token - # if the processing class does not have a pad token. - pad_token = args.pad_token or self._tokenizer.pad_token or self._tokenizer.eos_token - if pad_token not in self._tokenizer.get_vocab(): - raise ValueError( - f"The specified `pad_token` ('{pad_token}') is not found in the vocabulary of the given " - f"`processing_class` ({processing_class.__class__.__name__}). Ensure that the `pad_token` exists " - "in the vocabulary before using it as a padding token." - ) - self._tokenizer.pad_token = pad_token - # The model must agree with the tokenizer on the pad token from construction, so mirror it onto the model - # configs. - model.config.get_text_config().pad_token_id = self._tokenizer.pad_token_id - model.generation_config.pad_token_id = self._tokenizer.pad_token_id - if self._is_vision_dataset: - data_collator = DataCollatorForVisionLanguageModeling( - processor=processing_class, - max_length=args.max_length, - completion_only_loss=self.completion_only_loss, - pad_to_multiple_of=args.pad_to_multiple_of, - dataset_text_field=args.dataset_text_field, - assistant_only_loss=args.assistant_only_loss, - chat_template=self.chat_template, - ) - else: - data_collator = DataCollatorForLanguageModeling( - pad_token_id=self._tokenizer.pad_token_id, - padding_free=self.padding_free, - pad_to_multiple_of=args.pad_to_multiple_of, - ) - - if args.packing and args.packing_strategy in {"bfd", "bfd_split"} and not use_flash_attention: - logger.warning( - "You are using packing, but the attention implementation is not set to a supported Flash Attention " - "variant. Packing gathers multiple samples into a single sequence, and only the following " - f"implementations are known to reliably support this: {', '.join(sorted(FLASH_ATTENTION_VARIANTS))}. " - "Using other implementations may lead to cross-contamination between samples. To avoid this, either " - "disable packing by setting `packing=False`, or set `attn_implementation` in the model configuration " - "to one of these supported options." - ) - - # Dataset - if self.padding_free and not args.packing and args.max_length is not None and not self._is_vision_dataset: - raise ValueError( - "When `padding_free=True` without packing, `max_length` is not enforced. Either enable packing " - "(e.g., `packing=True, packing_strategy='bfd'`), provide already truncated inputs, or set " - "`max_length=None`." - ) - # Skip dataset preparation if `skip_prepare_dataset=True` in `dataset_kwargs`, or if it's a VLM, where - # preprocessing [e.g., image-to-pixel conversion] is too costly and done on the fly instead. - self._skip_prepare_dataset = ( - args.dataset_kwargs is not None - and args.dataset_kwargs.get("skip_prepare_dataset", False) - or self._is_vision_dataset - ) - # Kept on the instance so that `evaluate` can preprocess freshly-passed eval datasets the same way. - self._formatting_func = formatting_func - eval_datasets = ( - eval_dataset if isinstance(eval_dataset, dict) else {"eval": eval_dataset} if eval_dataset else {} - ) - self._reject_skip_prepare_without_labels({"train": train_dataset, **eval_datasets}, data_collator) - if not self._skip_prepare_dataset: - if self.completion_only_loss and formatting_func: - raise ValueError( - "A formatting function was provided while `completion_only_loss=True`, which is incompatible. " - "Using a formatter converts the dataset to a language modeling type, conflicting with " - "completion-only loss. To resolve this, apply your formatting function before passing the " - "dataset, or disable `completion_only_loss` in `SFTConfig`." - ) - self._unsloth_model_ref = model - train_dataset = self._prepare_dataset( - train_dataset, processing_class, args, args.packing, formatting_func, "train" - ) - if eval_dataset is not None: - packing = args.packing if args.eval_packing is None else args.eval_packing - if isinstance(eval_dataset, dict): - eval_dataset = { - key: self._prepare_dataset(dataset, processing_class, args, packing, formatting_func, key) - for key, dataset in eval_dataset.items() - } - else: - eval_dataset = self._prepare_dataset( - eval_dataset, processing_class, args, packing, formatting_func, "eval" - ) - - # Loss function - if args.loss_type not in ("nll", "dft"): - raise ValueError(f"Invalid `loss_type` {args.loss_type} passed. Supported values are 'nll' and 'dft'.") - if args.loss_type == "dft" and compute_loss_func is not None: - raise ValueError( - "You passed a `compute_loss_func` together with `loss_type='dft'` to the `SFTTrainer`. When using " - "`loss_type='dft'`, the loss function is internally set to the DFT loss, so passing a " - "`compute_loss_func` is not allowed." - ) - if compute_loss_func is not None: - warnings.warn( - "`compute_loss_func` is deprecated and will be removed in v2.0.0. Subclass `SFTTrainer` and override " - "`compute_loss` instead. Expect slower training meanwhile: the loss needs the full logits, so each " - "step runs an extra forward pass.", - FutureWarning, - stacklevel=2, - ) - if args.label_smoothing_factor > 0 and args.loss_type == "nll": - warnings.warn( - "`label_smoothing_factor` is deprecated in `SFTTrainer` and will be removed in v2.0.0. Label smoothing " - "needs the full logits, so the loss is computed from a forward pass that does not use the fused LM " - "head. Expect slower training meanwhile: each step runs an extra forward pass.", - FutureWarning, - stacklevel=2, - ) - if is_peft_model(model) and isinstance(model.get_output_embeddings(), BaseTunerLayer): - # The log-probabilities are computed by multiplying the hidden states by `lm_head.weight` directly, so an - # adapter on the LM head would be ignored and never trained. - raise ValueError( - "Applying a PEFT adapter to `lm_head` is not supported: the log-probabilities are computed from " - "`lm_head.weight` directly, so the adapter would never be trained. Remove `'lm_head'` from your " - "`target_modules`." - ) - - # Transformers explicitly set use_reentrant=True in the past to silence a PyTorch warning, but the default was - # never updated once PyTorch switched to recommending use_reentrant=False. Until that change lands upstream - # [see https://github.com/huggingface/transformers/pull/43203] and is released [most likely in 5.0.0], we - # default to the recommended non-reentrant behavior here, while preserving any user-provided value. - if args.gradient_checkpointing and Version(transformers.__version__) < Version("5.0.0"): - args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} - args.gradient_checkpointing_kwargs.setdefault("use_reentrant", False) - - # `selective` is a TRL-only key: when set, the SAC wrapper strips it before it reaches - # `torch.utils.checkpoint`, otherwise it is dropped here [a missing key already means SAC is off] - if (args.gradient_checkpointing_kwargs or {}).get("selective"): - if not args.gradient_checkpointing: - raise ValueError( - "`gradient_checkpointing_kwargs={'selective': True}` requires `gradient_checkpointing=True`." - ) - if (args.gradient_checkpointing_kwargs or {}).get("use_reentrant"): - raise ValueError( - "`gradient_checkpointing_kwargs={'selective': True}` requires non-reentrant gradient " - "checkpointing. Set `use_reentrant` to `False`, or leave it unset." - ) - # Wrap before `super[].__init__[]` so the context function is set when the Trainer enables checkpointing. - enable_selective_activation_checkpointing(model) - elif args.gradient_checkpointing_kwargs: - args.gradient_checkpointing_kwargs.pop("selective", None) - - super().__init__( - model=model, - args=args, - data_collator=data_collator, - train_dataset=train_dataset, - eval_dataset=eval_dataset, - processing_class=processing_class, - compute_loss_func=compute_loss_func, - compute_metrics=compute_metrics, - callbacks=callbacks, - optimizers=optimizers, - optimizer_cls_and_kwargs=optimizer_cls_and_kwargs, - preprocess_logits_for_metrics=preprocess_logits_for_metrics, - ) - - # With several GPUs visible and no distributed launcher, `Trainer` wraps the model in `nn.DataParallel`, whose - # replicas would all run the forward bound to the original model. Checked after `Trainer.__init__`, which sets - # `n_gpu` to 1 for a model split across devices with `device_map`. - if args.n_gpu > 1: - raise ValueError( - f"{type(self).__name__} does not support `nn.DataParallel`, which `Trainer` uses when several GPUs are " - "visible to a single process. Launch the script with `accelerate launch` or `torchrun`, or make a " - "single GPU visible with `CUDA_VISIBLE_DEVICES`." - ) - - # Liger's fused linear cross-entropy replaces `model.forward` when training starts, which would drop the fused - # LM head, so only its layer kernels are applied - if args.use_liger_kernel: - warnings.warn( - "`use_liger_kernel=True` is deprecated and will be removed in v2.0.0. Use the Hub kernels instead, " - 'with `model_init_kwargs={"use_kernels": True}`.', - FutureWarning, - stacklevel=2, - ) - liger_kernel_config = args.liger_kernel_config or {} - if liger_kernel_config.get("fused_linear_cross_entropy"): - raise ValueError( - '`liger_kernel_config={"fused_linear_cross_entropy": True}` is not supported: it replaces the ' - "model's forward, which this trainer patches with a fused LM head." - ) - args.liger_kernel_config = {**liger_kernel_config, "fused_linear_cross_entropy": False} - - # Compute the per-token log-probabilities in chunks, without materializing the full logits - # `is_top1` feeds the `mean_token_accuracy` metric - add_fused_lm_head( - self.model.get_base_model() if is_peft_model(self.model) else self.model, - outputs=("log_probs", "entropy", "is_top1"), - ) - - # Context parallelism can only express full causal attention: the per-layer attention mask is dropped - # and replaced by `is_causal=True`. Packed sequences rely on a block-diagonal mask to keep documents - # from attending to each other, so a packed batch would silently train with documents attending across - # their boundaries. Read the parallelism config off the accelerator rather than off `args`: when - # context parallelism is configured through an accelerate YAML, `args.parallelism_config` stays None. - if not _is_package_version_below("accelerate", "1.10.1"): - parallelism_config = self.accelerator.parallelism_config - if ( - parallelism_config is not None - and parallelism_config.cp_enabled - and (args.packing or args.eval_packing) - ): - raise ValueError( - "Packing is not compatible with context parallelism (`cp_size > 1`). Packing relies on a " - "block-diagonal attention mask to keep packed documents separate, and context parallelism " - "drops that mask, so documents would attend across their boundaries. Set `packing=False` " - "and `eval_packing=False`, or disable context parallelism." - ) - - # Sequence parallelism [Ulysses/ALST] shards batches along the sequence dimension and requires - # `position_ids` in every batch to preserve each token's global position. Sequence parallelism was added in - # accelerate 1.12.0. - if ( - Version(accelerate.__version__) >= Version("1.12.0") - and self.accelerator.parallelism_config is not None - and self.accelerator.parallelism_config.sp_enabled - and isinstance(self.data_collator, DataCollatorForLanguageModeling) - ): - self.data_collator.return_position_ids = True - - # Under DeepSpeed sequence parallelism [Ulysses/ALST], `Trainer` reduces the model's own loss across ranks, - # which is the fused LM head's negative log-likelihood - if ( - Version(accelerate.__version__) >= Version("1.12.0") - and self.accelerator.parallelism_config is not None - and self.accelerator.parallelism_config.sp_backend == "deepspeed" - and self.accelerator.parallelism_config.sp_enabled - and (args.loss_type == "dft" or self.compute_loss_func is not None or args.label_smoothing_factor > 0) - ): - raise ValueError( - "`loss_type='dft'`, `compute_loss_func` and `label_smoothing_factor` are not supported with DeepSpeed " - "sequence parallelism (Ulysses/ALST): the loss is the model's own negative log-likelihood, reduced " - "across ranks by `Trainer`." - ) - - # Initialize activation offloading context - if self.args.activation_offloading: - self.maybe_activation_offload_context = get_act_offloading_ctx_manager(model=self.model) - else: - self.maybe_activation_offload_context = contextlib.nullcontext() - - # MoE load-balancing auxiliary loss. `output_router_logits` in the config means the model returns its router - # logits; `router_aux_loss_coef` is the architecture's own coefficient, 0.0 when the config declares none. - text_config = model.config.get_text_config() - coef = args.router_aux_loss_coef - self.router_aux_loss_coef = getattr(text_config, "router_aux_loss_coef", 0.0) if coef is None else coef - if coef and not hasattr(text_config, "output_router_logits"): - raise ValueError( - f"`router_aux_loss_coef` is set to {coef} but {type(model).__name__} is not a Mixture-of-Experts model " - f"that returns its router logits, so there is no auxiliary loss to weight." - ) - self.aux_loss_enabled = hasattr(text_config, "output_router_logits") and self.router_aux_loss_coef != 0.0 - if hasattr(text_config, "router_aux_loss_coef"): - # The native and chunked forwards add the aux loss from the model config, so keep the config in sync with - # the coef: enable it [and propagate the coef] when non-zero, disable it otherwise. This overrides any - # `output_router_logits` the model was loaded with, so `router_aux_loss_coef=0.0` reliably turns it off. - text_config.output_router_logits = self.aux_loss_enabled - text_config.router_aux_loss_coef = self.router_aux_loss_coef - - # Initialize the metrics - self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} - self._total_train_tokens = 0 - - # Tensor parallel ranks are given the same batch, so a token count gathered across all processes repeats - # every token once per rank in the group. Context and sequence parallelism shard the batch before the loss is - # computed, so they need no such correction. `ParallelismConfig.tp_size` requires accelerate 1.10.0. - if Version(accelerate.__version__) >= Version("1.10.0") and self.accelerator.parallelism_config is not None: - self._tp_size = self.accelerator.parallelism_config.tp_size - else: - self._tp_size = 1 - - # Add tags to the model - self.model.add_model_tags(self._tag_names) - - def _prepare_dataset( - self, - dataset: Union[Dataset, IterableDataset], - processing_class, - args, - packing: bool, - formatting_func: Optional[Callable[[dict], str]], - dataset_name: str, - ) -> Union[Dataset, IterableDataset]: - import inspect as _inspect - try: - _unsloth_pack_has_strategy = "strategy" in _inspect.signature(pack_dataset).parameters - except Exception: - _unsloth_pack_has_strategy = True - _unsloth_wrapped_packing = packing and ( - getattr(args, "packing_strategy", None) in ("wrapped", "bfd_split", "bfd-requeue") - or not _unsloth_pack_has_strategy - ) - # All Unsloth Zoo code licensed under LGPLv3 - try: - if isinstance(dataset, ConstantLengthDataset): return dataset - except: - pass - - map_kwargs = {} - use_desc = isinstance(dataset, Dataset) - is_vlm = hasattr(processing_class, "tokenizer") - tokenizer = processing_class - if is_vlm: tokenizer = processing_class.tokenizer - - # Detect whether the model's module needs token_type_ids when training - import sys as _sys - _needs_token_type_ids = False - # Split to avoid compiler substring match on masking_utils names - _ccm = 'create_' + 'causal_mask_mapping' - _model = getattr(self, '_unsloth_model_ref', None) or getattr(self, 'model', None) - if _model is not None: - for _m in (_model, getattr(_model, 'model', None)): - if _m is None: continue - _mod = _sys.modules.get(type(_m).__module__) - if _mod is not None and hasattr(_mod, _ccm): - _needs_token_type_ids = True - break - - if not _needs_token_type_ids: - # Fallback: model not yet available, check processor class MRO - for _base in type(processing_class).__mro__: - _base_mod = getattr(_base, '__module__', '') - if 'transformers.models.' in _base_mod: - _modeling_mod = _base_mod.replace('.processing_', '.modeling_') - _mod = _sys.modules.get(_modeling_mod) - if _mod is not None and hasattr(_mod, _ccm): - _needs_token_type_ids = True - break - if _needs_token_type_ids and hasattr(args, 'remove_unused_columns'): - args.remove_unused_columns = False - - # Get max length - max_seq_length = getattr(args, "max_length", 0) or 0 - if max_seq_length == 0: max_seq_length = getattr(args, "max_seq_length", 0) - if max_seq_length == 0: max_seq_length = getattr(self, "max_seq_length", 0) - if max_seq_length == 0: max_seq_length = getattr(self, "max_seq", 0) - if max_seq_length == 0: raise RuntimeError("Unsloth: max_seq_length is 0! Please specify one!") - dataset_text_field = getattr(args, "dataset_text_field", "text") - do_truncation = max_seq_length != 0 - do_formatting_func = False - do_tokenize = True - do_prompt_completion = False - do_messages = False - - # Get correct column names - column_names = set(next(iter(dataset)).keys()) - used_column_names = ["input_ids"] - if "attention_mask" in column_names: - used_column_names.append("attention_mask") - if _needs_token_type_ids: - used_column_names.append("token_type_ids") - - # Skip tokenization if already tokenized; just set the data collator - from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling - if "labels" in column_names: - # Most likely forgot data collator! - if is_vlm and not hasattr(tokenizer, "pad"): - raise RuntimeError(f"Unsloth: {processing_class.__class__} does not have .pad!") - self.data_collator = DataCollatorForSeq2Seq(tokenizer) - used_column_names.append("labels") - do_tokenize = False - elif "input_ids" in column_names: - if is_vlm and not hasattr(tokenizer, "pad"): - raise RuntimeError(f"Unsloth: {processing_class.__class__} does not have .pad!") - mask_columns = [x for x in ("completion_mask", "assistant_masks") if x in column_names] - if mask_columns: - used_column_names += mask_columns - else: - self.data_collator = DataCollatorForLanguageModeling(tokenizer, mlm = False) - do_tokenize = False - elif "prompt" in column_names and "completion" in column_names: - # Prompt/completion dataset (used with completion_only_loss). - # TRL's __init__ already set self.data_collator for completion_only_loss - # before calling us -- we must NOT overwrite it here. - do_prompt_completion = True - used_column_names.append("completion_mask") - elif (dataset_text_field not in column_names or getattr(args, "assistant_only_loss", False)) and\ - formatting_func is None and "messages" in column_names and\ - isinstance(_first_messages := next(iter(dataset))["messages"], list) and\ - len(_first_messages) != 0 and isinstance(_first_messages[0], dict) and\ - "role" in _first_messages[0]: - # Text columns carry no assistant masks, so assistant_only_loss prefers "messages" even if text exists. - do_messages = True - elif dataset_text_field not in column_names: - do_formatting_func = True - if formatting_func is None: - raise RuntimeError("Unsloth: You must specify a `formatting_func`") - pass - if do_tokenize and not do_messages and not do_prompt_completion and getattr(args, "assistant_only_loss", False): - raise ValueError( - "Unsloth: assistant_only_loss needs a conversational `messages` dataset and no formatting_func. " - "Or use `train_on_responses_only` on your formatted text." - ) - - if do_tokenize: - if do_formatting_func: - test_text = formatting_func(next(iter(dataset))) - if not isinstance(test_text, list): - raise ValueError( - "Unsloth: The `formatting_func` should return a list of processed strings." - ) - test_text = test_text[0] - elif do_prompt_completion: - _first_ex = next(iter(dataset)) - try: - from trl import is_conversational as _sft_is_conversational - except ImportError: - def _sft_is_conversational(example): - for key in ("prompt", "completion", "messages"): - val = example.get(key) - if isinstance(val, list) and val and isinstance(val[0], dict): - if "role" in val[0] and "content" in val[0]: - return True - return False - _is_conv = _sft_is_conversational(_first_ex) - if not _is_conv: - if getattr(args, "assistant_only_loss", False): - raise ValueError( - "Unsloth: assistant_only_loss needs conversational prompt/completion messages, not plain strings." - ) - test_text = _first_ex["prompt"] - else: - test_text = None # chat template handles BOS - elif do_messages: - test_text = None # chat template handles BOS - else: - # No [0] on a str: that is the first CHARACTER, so startswith(bos_token) - # below could never match. Only unwrap when the field really is a list. - test_text = next(iter(dataset))[dataset_text_field] - if isinstance(test_text, (list, tuple)): - test_text = test_text[0] if len(test_text) != 0 else None - - chat_template = getattr(processing_class, 'chat_template', '') - if chat_template == '' and is_vlm: - chat_template = getattr(tokenizer, 'chat_template', '') - if chat_template is None: - chat_template = '' - - # Detect double BOS so we can drop the duplicate - add_special_tokens = True - bos_token_1 = getattr(processing_class, 'bos_token', None) - bos_token_2 = getattr(tokenizer, 'bos_token', None) - bos_token = bos_token_1 or bos_token_2 - - if bos_token is not None: - if (test_text is not None and test_text.startswith(bos_token)) or bos_token in chat_template: - add_special_tokens = False - print("Unsloth: We found double BOS tokens - we shall remove one automatically.") - pass - - def _tokenize(example): - return tokenizer( - example[dataset_text_field] if not do_formatting_func else formatting_func(example), - truncation = do_truncation and not _unsloth_wrapped_packing, - max_length = max_seq_length, - return_token_type_ids = _needs_token_type_ids, - add_special_tokens = add_special_tokens, - ) - pass - - # Worker count via the shared policy; spelled exactly as the edit unsloth applies here, so unsloth sees it as done. - if not isinstance(dataset, IterableDataset): - try: - from unsloth_zoo.dataset_num_proc import get_dataset_num_proc as _unsloth_get_dataset_num_proc - except ImportError: - from unsloth.dataset_num_proc import get_dataset_num_proc as _unsloth_get_dataset_num_proc - map_kwargs["num_proc"] = _unsloth_get_dataset_num_proc( - getattr(args, "dataset_num_proc", None) - ) - else: - map_kwargs["batch_size"] = _iterable_batch_size(dataset) - - if do_messages or do_prompt_completion: - assistant_only_loss = bool(getattr(args, "assistant_only_loss", False)) - training_chat_template = getattr(self, "chat_template", None) - # No {% generation %} markers (or no TRL training template): assistant masks would be all zeros, - # so mask with train_on_responses_only markers instead. - import re as _re - def _template_text(template): - # Named template sets are picked per call: check them all. - if isinstance(template, dict): - return "\n".join(v for v in template.values() if isinstance(v, str)) - return template or "" - marker_masks = assistant_only_loss and ( - getattr(self, "_unsloth_assistant_mask_fallback", False) or - _re.search(r"\{%-?\s*generation\s*-?%\}", _template_text(training_chat_template or chat_template)) is None - ) - response_labels = None - if marker_masks: - from unsloth_zoo.dataset_utils import train_on_responses_only as _response_labels - marker_kwargs = {} - # Harmony (gpt-oss): auto-detection would keep only the final channel, not analysis/tool calls. - if "<|channel|>" in _template_text(getattr(tokenizer, "chat_template", None)) and\ - not hasattr(tokenizer, "_unsloth_output_part"): - marker_kwargs = dict( - instruction_part = "<|start|>user<|message|>", - response_part = "<|start|>assistant", - ) - response_labels = _response_labels( - None, tokenizer = processing_class, return_function = True, **marker_kwargs, - ) - if assistant_only_loss: - used_column_names.append("assistant_masks") - - if do_prompt_completion: - _eos_token = getattr(tokenizer, 'eos_token', None) - - def _tokenize_pc(example): - if _is_conv: - template_kwargs = dict(example.get("chat_template_kwargs") or {}) - if training_chat_template is not None: - template_kwargs["chat_template"] = training_chat_template - prompt_ids = processing_class.apply_chat_template( - example["prompt"], tokenize=True, - add_generation_prompt=True, return_dict=False, - tools=example.get("tools"), - **template_kwargs, - ) - if prompt_ids and isinstance(prompt_ids[0], list): - prompt_ids = prompt_ids[0] - pc_processed = processing_class.apply_chat_template( - example["prompt"] + example["completion"], - return_dict=True, tokenize=True, - return_assistant_tokens_mask=assistant_only_loss and not marker_masks, - tools=example.get("tools"), - **template_kwargs, - ) - if isinstance(pc_processed.get("input_ids", [None])[0], list): - pc_processed = {k: v[0] for k, v in pc_processed.items()} - pc_ids = pc_processed["input_ids"] - else: - _completion = example["completion"] - if _eos_token and not _completion.endswith(_eos_token): - _completion = _completion + _eos_token - prompt_ids = tokenizer( - example["prompt"], add_special_tokens=add_special_tokens, - )["input_ids"] - pc_ids = tokenizer( - example["prompt"] + _completion, - add_special_tokens=add_special_tokens, - )["input_ids"] - n_prompt = min(len(prompt_ids), len(pc_ids)) - completion_mask = [0] * n_prompt + [1] * (len(pc_ids) - n_prompt) - result = {"input_ids": pc_ids, "completion_mask": completion_mask} - if assistant_only_loss: - if marker_masks: - labels = response_labels({"input_ids": [pc_ids]})["labels"][0] - result["assistant_masks"] = [int(x != -100) for x in labels] - else: - result["assistant_masks"] = list(pc_processed["assistant_masks"]) - if do_truncation and not _unsloth_wrapped_packing and max_seq_length > 0: - result = {k: v[:max_seq_length] for k, v in result.items()} - if _needs_token_type_ids: - result["token_type_ids"] = [0] * len(result["input_ids"]) - return result - - if use_desc: - map_kwargs["desc"] = 'Unsloth: Tokenizing ["prompt"+"completion"]' - import warnings as _w - try: - from unsloth_zoo.dataset_num_proc import map_failure_diagnostics as _unsloth_map_diagnostics - except ImportError: - from unsloth.dataset_num_proc import map_failure_diagnostics as _unsloth_map_diagnostics - with _w.catch_warnings(), _unsloth_map_diagnostics(map_kwargs.get("num_proc")): - _w.filterwarnings("ignore", message=".*couldn't be hashed properly.*") - dataset = dataset.map( - _tokenize_pc, batched=False, - remove_columns=list(column_names), **map_kwargs, - ) - if assistant_only_loss: - # Prompt-history assistant turns are no target while completion_mask applies. - pc_completion_only = getattr(self, "completion_only_loss", None) - if pc_completion_only is None: - pc_completion_only = getattr(args, "completion_only_loss", None) - if pc_completion_only is None or pc_completion_only: - dataset = dataset.filter( - lambda row: any(a and c for a, c in zip(row["assistant_masks"], row["completion_mask"])) - ) - else: - dataset = dataset.filter(lambda row: any(row["assistant_masks"])) - elif do_messages: - truncate_messages = do_truncation and isinstance(max_seq_length, int) and max_seq_length > 0 and\ - not locals().get("_unsloth_wrapped_packing", False) - message_cut = slice(-max_seq_length, None) if truncate_messages and\ - getattr(args, "truncation_mode", "keep_start") == "keep_end" else slice(None, max_seq_length) - def _tokenize_messages(example): - template_kwargs = dict(example.get("chat_template_kwargs") or {}) - if training_chat_template is not None: - template_kwargs["chat_template"] = training_chat_template - # Arrow fills missing keys with None, which `is defined` template checks then index. - messages = [ - {k: v for k, v in message.items() if v is not None} - if isinstance(message, dict) else message - for message in example["messages"] - ] - if is_vlm: - # Processor templates index content parts, as TRL's prepare_multimodal_messages builds them. - messages = [ - {**message, "content": [{"type": "text", "text": message["content"]}]} - if isinstance(message, dict) and isinstance(message.get("content"), str) else message - for message in messages - ] - tools = example.get("tools") - if isinstance(tools, str): - import json as _json - tools = _json.loads(tools) - processed = processing_class.apply_chat_template( - messages, tools = tools, - tokenize = True, return_dict = True, - return_assistant_tokens_mask = assistant_only_loss and not marker_masks, - **template_kwargs, - ) - input_ids = processed["input_ids"] - if input_ids and isinstance(input_ids[0], list): - input_ids = input_ids[0] - result = {"input_ids": list(input_ids)} - if assistant_only_loss: - if marker_masks: - labels = response_labels({"input_ids": [result["input_ids"]]})["labels"][0] - result["assistant_masks"] = [int(x != -100) for x in labels] - else: - masks = processed["assistant_masks"] - if masks and isinstance(masks[0], list): - masks = masks[0] - result["assistant_masks"] = list(masks) - if truncate_messages: - result = {k: v[message_cut] for k, v in result.items()} - if _needs_token_type_ids: - result["token_type_ids"] = [0] * len(result["input_ids"]) - return result - - if use_desc: map_kwargs["desc"] = 'Unsloth: Tokenizing ["messages"]' - import warnings as _w - with _w.catch_warnings(): - _w.filterwarnings("ignore", message = ".*couldn't be hashed properly.*") - dataset = dataset.map( - _tokenize_messages, batched=False, - remove_columns=list(column_names), **map_kwargs, - ) - if assistant_only_loss: - # Truncation can drop every assistant token; such a row has no target. - dataset = dataset.filter(lambda row: any(row["assistant_masks"])) - else: - if use_desc: map_kwargs["desc"] = f'Unsloth: Tokenizing ["{dataset_text_field}"]' - import warnings as _w - try: - from unsloth_zoo.dataset_num_proc import map_failure_diagnostics as _unsloth_map_diagnostics - except ImportError: - from unsloth.dataset_num_proc import map_failure_diagnostics as _unsloth_map_diagnostics - with _w.catch_warnings(), _unsloth_map_diagnostics(map_kwargs.get("num_proc")): - _w.filterwarnings("ignore", message=".*couldn't be hashed properly.*") - dataset = dataset.map(_tokenize, batched = True, remove_columns = list(column_names), **map_kwargs) - - # VLMs need .pad; switch the data collator - # transformers' collator rebuilds labels from input_ids, dropping assistant masks: keep TRL's for messages. - if is_vlm and not hasattr(processing_class, "pad") and not do_prompt_completion and not do_messages: - data_collator = DataCollatorForLanguageModeling(tokenizer, mlm = False) - self.data_collator = data_collator - pass - pass - - # Imported here: this function's source is copied into the compiled trainer module. - from unsloth_zoo.dataset_utils import trl_collator_applies_masks - column_names = set(next(iter(dataset)).keys()) - if "labels" not in column_names and not trl_collator_applies_masks(): - completion_only_loss = getattr(self, "completion_only_loss", None) - if completion_only_loss is None: - completion_only_loss = getattr(args, "completion_only_loss", None) - if completion_only_loss is None: - completion_only_loss = do_prompt_completion - mask_columns = [x for x in ("completion_mask", "assistant_masks") if x in column_names] - if not completion_only_loss and "completion_mask" in mask_columns: - mask_columns.remove("completion_mask") - if mask_columns: - def _build_labels(example): - masks = [example[x] for x in mask_columns] - return {"labels": [t if all(m) else -100 for t, *m in zip(example["input_ids"], *masks)]} - label_kwargs = {} - if isinstance(dataset, Dataset): - label_kwargs["num_proc"] = map_kwargs.get("num_proc", None) - label_kwargs["desc"] = f"Unsloth: Building labels for {dataset_name} dataset" - dataset = dataset.map(_build_labels, remove_columns = mask_columns, **label_kwargs) - used_column_names = [x for x in used_column_names if x not in mask_columns] + ["labels"] - pass - if packing: - # Use TRL's pack_dataset if available - try: - # A presence probe, not a use: TRL exports pack_dataset only on some - # versions, and the bare name is what the except below is for. - pack_dataset # noqa: F821 - except: - print("Unsloth: Hugging Face's packing is currently buggy - we're disabling it for now!") - return dataset - - if max_seq_length == 0: - raise ValueError("When packing is enabled, `max_seq_length` can't be `None`.") - - if use_desc: map_kwargs["desc"] = f"Unsloth: Packing {dataset_name} dataset" - _pack_kwargs = {"map_kwargs": map_kwargs} - if _unsloth_pack_has_strategy: - _pack_kwargs["strategy"] = getattr(args, "packing_strategy", "bfd") - dataset = pack_dataset( - dataset.select_columns(used_column_names), - max_seq_length, - **_pack_kwargs, - ) - pass - return dataset - - def _set_signature_columns_if_needed(self): - # If `self.args.remove_unused_columns` is True, non-signature columns are removed. - # By default, this method sets `self._signature_columns` to the model's expected inputs (usually, "input_ids", - # "attention_mask" and "labels"). Dataset preparation also produces a "seq_lengths" column (for packing / - # padding-free), so we override the default signature columns to keep it alongside the model inputs. - if self._signature_columns is None: - if self._is_vision_dataset: - self._signature_columns = ["messages", "prompt", "completion", "image", "images"] - else: - self._signature_columns = ["input_ids", "labels", "seq_lengths"] - - def _reject_skip_prepare_without_labels(self, datasets: dict[str, Dataset], data_collator) -> None: - # This guard may look defensive, but it covers a behavior change introduced when label building moved from - # the collator to dataset preparation: the collator used to consume the mask columns directly, so a - # skipped-preparation dataset carrying masks trained correctly. Now labels are built during preparation, which - # is skipped here, and the collator ignores the mask columns. Without a "labels" column, such a dataset would - # silently optimize the loss over the full sequence, so we fail loudly instead. Checked both at init and in - # `evaluate`, since a dataset passed directly to `evaluate` also skips preparation. - if not ( - self._skip_prepare_dataset - and not self._is_vision_dataset - and isinstance(data_collator, DataCollatorForLanguageModeling) - ): - return - for name, dataset in datasets.items(): - cols = _unsloth_dataset_column_names(dataset) - if "labels" not in cols and ("completion_mask" in cols or "assistant_masks" in cols): - raise ValueError( - f"The {name} dataset has mask columns but no 'labels', and `skip_prepare_dataset=True` skips " - "label building, so it would train on the full sequence. Add a 'labels' column (-100 for " - "non-loss tokens) or drop `skip_prepare_dataset`." - ) - - def evaluate( - self, - eval_dataset: Dataset - | IterableDataset - | DatasetDict - | IterableDatasetDict - | dict[str, Dataset | IterableDataset] - | None = None, - ignore_keys: list[str] | None = None, - metric_key_prefix: str = "eval", - ) -> dict[str, float]: - # When a dataset is passed directly to `evaluate` (e.g. a held-out test set), preprocess it the same way - # `__init__` does, so that `evaluate` accepts the same dataset types as the trainer (language modeling, - # prompt-completion, etc.). `_prepare_dataset` is idempotent: it skips datasets that are already tokenized. A - # `str` selects a dataset that was already prepared at init time, so it's left untouched. - if not self._skip_prepare_dataset and eval_dataset is not None and not isinstance(eval_dataset, str): - packing = self.args.packing if self.args.eval_packing is None else self.args.eval_packing - if isinstance(eval_dataset, dict): - eval_dataset = { - key: self._prepare_dataset( - dataset, self.processing_class, self.args, packing, self._formatting_func, key - ) - for key, dataset in eval_dataset.items() - } - else: - eval_dataset = self._prepare_dataset( - eval_dataset, self.processing_class, self.args, packing, self._formatting_func, "eval" - ) - eval_datasets = ( - eval_dataset - if isinstance(eval_dataset, dict) - else {"eval": eval_dataset} - if eval_dataset is not None and not isinstance(eval_dataset, str) - else {} - ) - self._reject_skip_prepare_without_labels(eval_datasets, self.data_collator) - return super().evaluate( - eval_dataset=eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix - ) - - def compute_loss( - self, - model, - inputs, - return_outputs = False, - num_items_in_batch = None, - ): - # TRL >= 1.15 computes loss_type="dft" only in the compute_loss this replaces (else it trains as NLL): - # run it on TRL's fused LM head. Older TRL routes DFT through compute_loss_func. - if ( - getattr(self.args, "loss_type", None) == "dft" - and self.compute_loss_func is None - and getattr(self, "label_smoother", None) is None - ): - import trl.trainer.utils as _trl_utils - if hasattr(_trl_utils, "add_fused_lm_head"): - # TRL's fused head projects with lm_head.weight directly, skipping QAT fake quantization. - try: - head = model.get_output_embeddings() - except Exception: - head = None - try: - from unsloth.models.llama import lm_head_needs_forward as _needs_forward - except Exception: - _needs_forward = lambda module: False - if head is not None and ( - _needs_forward(head) - or any( - getattr(m, "activation_fake_quantizer", None) is not None - or getattr(m, "weight_fake_quantizer", None) is not None - for m in head.modules() - ) - ): - raise NotImplementedError( - "Unsloth: loss_type = 'dft' runs on TRL's fused LM head, which skips the QAT fake " - "quantization of lm_head. Leave lm_head out of QAT or use loss_type = 'nll'." - ) - inputs = {**inputs, "use_cache": False} - outputs = model(**inputs, fused_lm_head = True) - if getattr(outputs, "log_probs", None) is None: - raise RuntimeError( - "Unsloth: loss_type = 'dft' needs TRL's fused LM head, but the model returned no " - "per-token log-probabilities. Use loss_type = 'nll'." - ) - per_token_loss = -outputs.log_probs * outputs.log_probs.exp().detach() - # Trainer divides by gradient accumulation itself when the model does not take num_items_in_batch. - if not getattr(self, "model_accepts_loss_kwargs", True): - num_items_in_batch = None - num_tokens = ( - outputs.label_mask.sum() if num_items_in_batch is None else num_items_in_batch - ) - loss = per_token_loss.sum() / torch.as_tensor(num_tokens).clamp(min = 1) - if self.args.average_tokens_across_devices and num_items_in_batch is not None: - # As TRL / Trainer: tensor parallel ranks share one batch, so they do not count. - tp_size = ( - self.get_tp_size() if callable(getattr(self, "get_tp_size", None)) else 1 - ) - loss = loss * (self.accelerator.num_processes // max(int(tp_size or 1), 1)) - aux_loss = getattr(outputs, "aux_loss", None) - if getattr(self, "aux_loss_enabled", False) and aux_loss is not None: - steps = ( - self.current_gradient_accumulation_steps - if num_items_in_batch is not None - else 1 - ) - loss = loss + self.router_aux_loss_coef * aux_loss / steps - if return_outputs: - # prediction_step / compute_metrics read logits, which the fused head never builds (as TRL does). - with torch.no_grad(): - outputs = model(**{k: v for k, v in inputs.items() if k != "labels"}) - return (loss, outputs) if return_outputs else loss - outputs = super().compute_loss( - model, - inputs, - return_outputs = return_outputs, - num_items_in_batch = num_items_in_batch, - ) - return outputs - - # During eval, Trainer calls prediction_step, which asks `compute_loss` for the outputs even when only the loss is - # needed. Ask for them only when the logits are, since they take an extra forward pass. - def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys: list[str] | None = None): - inputs = self._prepare_inputs(inputs) - with torch.no_grad(), self.compute_loss_context_manager(): - if prediction_loss_only: - loss = self.compute_loss(model, inputs, return_outputs=False) - logits, labels = None, None - else: - loss, outputs = self.compute_loss(model, inputs, return_outputs=True) - logits, labels = outputs.logits, inputs.get("labels") - return loss, logits, labels - - # Override training step to add activation offloading context. - def training_step(self, *args, **kwargs): - with self.maybe_activation_offload_context: - return super().training_step(*args, **kwargs) - - def log(self, logs: dict[str, float], start_time: float | None = None) -> None: - mode = "train" if self.model.training else "eval" - metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()} # average the metrics - - # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` - # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. - if mode == "eval": - metrics = {f"eval_{key}": val for key, val in metrics.items()} - - logs.update(metrics) - super().log(logs, start_time) - self._metrics[mode].clear() - - # Ensure the model card is saved along with the checkpoint - def _save_checkpoint(self, model, trial): - if self.args.hub_model_id is None: - model_name = Path(self.args.output_dir).name - else: - model_name = self.args.hub_model_id.split("/")[-1] - self.create_model_card(model_name=model_name) - super()._save_checkpoint(model, trial) -class UnslothSFTTrainer(_UnslothSFTTrainer): - """ - - Trainer for Supervised Fine-Tuning (SFT) method. - - This class is a wrapper around the [`~transformers.Trainer`] class and inherits all of its attributes and methods. - - Example: - - ```python - >>> from trl import SFTTrainer - >>> from datasets import load_dataset - - >>> dataset = load_dataset("roneneldan/TinyStories", split="train[:1%]") - - >>> trainer = SFTTrainer( - ... model="Qwen/Qwen2.5-0.5B-Instruct", - ... train_dataset=dataset, - ... ) - >>> trainer.train() - ``` - - Args: - model (`str` or [`~transformers.PreTrainedModel`] or [`~peft.PeftModel`]): - Model to be trained. Can be either: - - - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a - path to a *directory* containing model weights saved using - [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded - using `.from_pretrained` (where `` is derived from the model - config) with the keyword arguments in `args.model_init_kwargs`. If `dtype` is not specified in - `args.model_init_kwargs`, it defaults to `float32`. This differs from - [`~transformers.PreTrainedModel.from_pretrained`], where (since Transformers v5) the dtype is inferred - from the model config. - - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. - - A [`~peft.PeftModel`] object. Only causal language models are supported. - args ([`SFTConfig`], *optional*): - Configuration for this trainer. If `None`, a default configuration is used. - data_collator ([`~transformers.DataCollator`], *optional*): - Function to use to form a batch from a list of elements of the processed `train_dataset` or `eval_dataset`. - Will default to [`~trainer.sft_trainer.DataCollatorForLanguageModeling`], or to - [`~trainer.sft_trainer.DataCollatorForVisionLanguageModeling`] if the dataset contains images. - train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): - Dataset to use for training. This trainer supports both [language modeling](#language-modeling) type and - [prompt-completion](#prompt-completion) type. The format of the samples can be either: - - - [Standard](dataset_formats#standard): Each sample contains plain text. - - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role - and content). - - The trainer also supports pre-tokenized datasets, recognized by a required `input_ids` column. An optional - `labels` column (`-100` on tokens excluded from the loss) is used as is if present; otherwise labels are - built from the optional `assistant_masks` / `completion_mask` columns (which are folded in then dropped, - `completion_mask` only when `completion_only_loss=True`), or default to a copy of `input_ids`. Sequences - are truncated to `max_length` during preparation. With `skip_prepare_dataset=True`, preparation is skipped - and the collator is expected to handle the dataset as is. - - When `train_dataset` is an [`~datasets.IterableDataset`] (e.g. a streaming dataset), `max_steps` must be - set in the training arguments, since its length cannot be inferred and the total number of training steps - is required to bound the training loop and configure the learning rate scheduler. - eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`], [`~datasets.DatasetDict`], [`~datasets.IterableDatasetDict`] or `dict[str, Dataset | IterableDataset]`): - Dataset to use for evaluation. It must meet the same requirements as `train_dataset`. - processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.ProcessorMixin`], *optional*): - Processing class used to process the data. If `None`, the processing class is loaded from the model's name - with [`~transformers.AutoProcessor.from_pretrained`]. A padding token, `tokenizer.pad_token`, must be set. - If the processing class has not set a padding token, `tokenizer.eos_token` will be used as the default. - compute_loss_func (`Callable`, *optional*): - A function that accepts the model outputs, the labels, and the number of items in the entire accumulated - batch (batch_size * gradient_accumulation_steps) and returns the loss. This is deprecated and will be - removed in v2.0.0: subclass [`SFTTrainer`] and override `compute_loss` instead. - compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): - The function that will be used to compute metrics at evaluation. Must take a - [`~transformers.EvalPrediction`] and return a dictionary string to metric values. When passing - [`SFTConfig`] with `batch_eval_metrics` set to `True`, your `compute_metrics` function must take a boolean - `compute_result` argument. This will be triggered after the last eval batch to signal that the function - needs to calculate and return the global summary statistics rather than accumulating the batch-level - statistics. The logits it receives are recomputed with an extra forward pass. This is deprecated and will - be removed in v2.0.0. - callbacks (list of [`~transformers.TrainerCallback`], *optional*): - List of callbacks to customize the training loop. Will add those to the list of default callbacks detailed - in [here](https://huggingface.co/docs/transformers/main_classes/callback). - - If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`] - method. - optimizers (`tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None]`, *optional*, defaults to `(None, None)`): - A tuple containing the optimizer and the scheduler to use. Will default to an instance of `AdamW` on your - model and a scheduler given by [`~transformers.get_linear_schedule_with_warmup`] controlled by `args`. - optimizer_cls_and_kwargs (`tuple[Type[torch.optim.Optimizer], Dict[str, Any]]`, *optional*): - A tuple containing the optimizer class and keyword arguments to use. Overrides `optim` and `optim_args` in - `args`. Incompatible with the `optimizers` argument. - - Unlike `optimizers`, this argument avoids the need to place model parameters on the correct devices before - initializing the Trainer. - preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`, *optional*): - A function that preprocess the logits right before caching them at each evaluation step. Must take two - tensors, the logits and the labels, and return the logits once processed as desired. The modifications made - by this function will be reflected in the predictions received by `compute_metrics`. - - Note that the labels (second parameter) will be `None` if the dataset does not have them. - quantization_config ([`~transformers.BitsAndBytesConfig`], *optional*): - Quantization configuration used when loading the model from a model identifier. Combine with `peft_config` - for QLoRA training. Ignored if the model is already instantiated. - peft_config ([`~peft.PeftConfig`], *optional*): - PEFT configuration used to wrap the model. If `None`, the model is not wrapped. - formatting_func (`Callable`, *optional*): - Formatting function applied to the dataset before tokenization. Applying the formatting function explicitly - converts the dataset into a [language modeling](#language-modeling) type. - - """ - def __init__( - self, - model, - args = None, - data_collator = None, - train_dataset = None, - eval_dataset = None, - processing_class = None, - compute_loss_func = None, - compute_metrics = None, - callbacks = None, - optimizer_cls_and_kwargs = None, - preprocess_logits_for_metrics = None, - quantization_config = None, - peft_config = None, - formatting_func = None, - **kwargs - ): - if args is None: args = UnslothSFTConfig() - _unsloth_precision_requested = tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) - use_bf16 = getattr(args, 'bf16', False) - if type(use_bf16) is not bool: use_bf16 = False - use_fp16 = getattr(args, 'fp16', False) - if type(use_fp16) is not bool: use_fp16 = False - force_float32 = False - try: - from unsloth_zoo.device_type import device_is_bf16_supported as _bf16_supported - except Exception: - _bf16_supported = torch.cuda.is_bf16_supported - try: - from unsloth_zoo.device_type import DEVICE_TYPE as _unsloth_device_type - except Exception: - _unsloth_device_type = None - full_finetuning = getattr(model, '_unsloth_full_finetuning', None) - if full_finetuning is None: full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1' - model_forced_float32 = getattr(model, '_unsloth_forced_float32', None) - if model_forced_float32 is None: model_forced_float32 = os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1' - if model_forced_float32 and not (full_finetuning and _bf16_supported()): - print('Unsloth: Switching to float32 training since model cannot work with float16') - force_float32 = True - mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') - dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None) - if dtype is None: dtype = model.get_input_embeddings().weight.dtype - from unsloth_zoo.utils import _get_dtype - dtype = _get_dtype(dtype) - float16 = dtype == torch.float16 - bfloat16 = dtype == torch.bfloat16 - float32 = dtype == torch.float32 - user_float32 = bool(getattr(model, '_unsloth_user_float32', False)) - if full_finetuning: - if bfloat16 and use_fp16: use_fp16 = False - if float16 and use_bf16: use_bf16 = False - if not force_float32 and bfloat16 and use_fp16 and _unsloth_device_type == 'xpu' and not _bf16_supported(): - print('Unsloth: The model was loaded in bfloat16 on purpose, so training uses bfloat16 mixed precision instead of float16.') - use_fp16 = False - use_bf16 = True - args.fp16 = False - args.bf16 = True - if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`') - if not force_float32 and (bfloat16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`') - if force_float32: - # Forced float32 training - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32' and float32 and user_float32 and not _bf16_supported(): - print('Unsloth: Model is in float32 and this GPU has no bfloat16 support, so training stays in float32. Pass fp16 = True to force float16 mixed precision instead.') - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32': - # Mixed precision training. bf16 only if the GPU supports it; V100/T4 use fp16. - use_bf16_amp = (not float16) and (_bf16_supported() or (bfloat16 and _unsloth_device_type == 'xpu')) - args.fp16 = not use_bf16_amp - args.bf16 = use_bf16_amp - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16_amp else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16_amp else 'fp16' - # args.mixed_precision is a new argument which needs to be set now - elif mixed_precision_dtype == 'bfloat16': - # Both False since bfloat16 full finetuning doesn't do any autocasting. - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif use_bf16 or use_fp16: - # transformers <5 exported this itself from fp16/bf16; 5.x dropped the write, so an - # explicit flag left it unset and GRPO readers defaulted to 'fp16', wrapping a - # bfloat16 model in a float16 autocast. See unslothai/unsloth#4891. - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16 else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16 else 'fp16' - - if getattr(args, 'eval_dataset', None) is not None and getattr(args, 'eval_strategy', 'no') == 'no': - args.eval_strategy = 'steps' - if getattr(args, 'eval_steps', None) is None: args.eval_steps = 0.1 - ga_steps = getattr(args, 'gradient_accumulation_steps', None) - if ga_steps is not None and ga_steps > 1: - from transformers import __version__ as transformers_version - if Version(transformers_version) <= Version('4.45.2'): - print('**** Unsloth: Please use our fixed gradient_accumulation_steps by updating transformers, TRL and Unsloth!\n' - '`pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth transformers trl unsloth_zoo`') - if getattr(args, 'eval_strategy', 'no') != 'no': - eval_bsz = getattr(args, 'per_device_eval_batch_size', 8) - eval_generations = getattr(args, 'num_generations_eval', None) or getattr(args, 'num_generations', None) or 1 - if eval_bsz == 8 and args.per_device_train_batch_size < eval_bsz and (args.per_device_train_batch_size * args.world_size) % eval_generations == 0: args.per_device_eval_batch_size = args.per_device_train_batch_size - if getattr(args, 'eval_accumulation_steps', None) is None and ga_steps is not None: args.eval_accumulation_steps = ga_steps - fp16_full_eval = getattr(args, 'fp16_full_eval', False) - if type(fp16_full_eval) is not bool: fp16_full_eval = False - bf16_full_eval = getattr(args, 'bf16_full_eval', False) - if type(bf16_full_eval) is not bool: bf16_full_eval = False - if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True - if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False - if force_float32: - args.bf16_full_eval = False - args.fp16_full_eval = False - elif os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16': - args.bf16_full_eval = True - args.fp16_full_eval = False - elif not bf16_full_eval and not fp16_full_eval: - args.bf16_full_eval = args.bf16 - args.fp16_full_eval = args.fp16 - if tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) != _unsloth_precision_requested: - from unsloth.models.loader_utils import sync_deepspeed_mixed_precision - sync_deepspeed_mixed_precision(args) - if model is not None: - from unsloth.models.loader_utils import apply_deepspeed_zero3_model_setup - apply_deepspeed_zero3_model_setup(model) - if model is not None: - _warnings_issued = getattr(model, 'warnings_issued', None) - if _warnings_issued is None: - model.warnings_issued = {} - elif not isinstance(_warnings_issued, dict): - try: - model.warnings_issued = dict(_warnings_issued) - except Exception: - model.warnings_issued = {} - _unsloth_explicit_max_length = None - try: - import dataclasses as _unsloth_dc - _unsloth_cfg_cls = type(args) - while '_unsloth_patched_rl_config' in _unsloth_cfg_cls.__dict__ or _unsloth_cfg_cls.__name__.startswith('Unsloth'): - _unsloth_cfg_cls = _unsloth_cfg_cls.__bases__[0] - _unsloth_default_max_length = None - for _unsloth_field in _unsloth_dc.fields(_unsloth_cfg_cls): - if _unsloth_field.name == 'max_length': _unsloth_default_max_length = _unsloth_field.default - _unsloth_given_max_length = getattr(args, 'max_length', None) - if (_unsloth_given_max_length or 0) > 0 and _unsloth_given_max_length != _unsloth_default_max_length: - _unsloth_explicit_max_length = _unsloth_given_max_length - except Exception: - _unsloth_explicit_max_length = None - if 'max_seq_length' not in locals() and not hasattr(args, 'max_seq_length'): - pass - else: - model_max_seq_length = getattr(model, 'max_seq_length', None) - args_max_seq_length = getattr(args, 'max_seq_length', None) - if args_max_seq_length is None and model_max_seq_length is not None: - max_seq_length = min(model.max_seq_length, _unsloth_explicit_max_length) if _unsloth_explicit_max_length else model.max_seq_length - if hasattr(args, 'max_seq_length'): args.max_seq_length = max_seq_length - elif args_max_seq_length is not None and model_max_seq_length is not None: - if args_max_seq_length > model_max_seq_length: - print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but ' - 'the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.') - args.max_seq_length = model_max_seq_length - if 'max_length' not in locals() and not hasattr(args, 'max_length'): - pass - else: - if hasattr(args, 'max_seq_length') and args.max_seq_length is not None and args.max_seq_length > 0: - if hasattr(args, 'max_length'): - args.max_length = args.max_seq_length - max_length = args.max_length - else: - model_max_length = getattr(model, 'max_seq_length', None) - if model_max_length is None: model_max_length = getattr(model, 'max_length', None) - if model_max_length is not None: - args.max_length = min(_unsloth_explicit_max_length, model_max_length) if _unsloth_explicit_max_length else model_max_length - max_length = args.max_length - elif hasattr(args, 'max_length') and args.max_length is not None: - max_length = args.max_length - # if we are here, then we are in a weird case where max_length is set but max_seq_length is not set - setattr(model, 'max_seq_length', max_length) - else: - print('Unsloth: We did not find `max_seq_length` or `max_length` in the model or args. We will set it to 1024.') - args.max_length = 1024 - if getattr(args, 'padding_free', False) is True and not getattr(args, 'packing', False) and getattr(args, 'max_length', None) is not None: - _unsloth_prep_truncates = True - _unsloth_skip_prepare = False - try: - _unsloth_dataset_kwargs = getattr(args, 'dataset_kwargs', None) - if _unsloth_dataset_kwargs is not None and _unsloth_dataset_kwargs.get('skip_prepare_dataset', False): - _unsloth_prep_truncates = False - _unsloth_skip_prepare = True - except Exception: - pass - def _unsloth_is_transformed(_ds): - _f = getattr(_ds, 'format', None) - _f = _f.get('type') if isinstance(_f, dict) else None - return (_f or getattr(_ds, '_format_type', None)) == 'custom' - try: - _unsloth_transformed = _unsloth_is_transformed(train_dataset) - _unsloth_columns = None if _unsloth_transformed else getattr(train_dataset, 'column_names', None) - if _unsloth_columns is None and train_dataset is not None: - _unsloth_probe_cols = iter(train_dataset) - if _unsloth_probe_cols is train_dataset or _unsloth_probe_cols is iter(train_dataset): - _unsloth_prep_truncates = False - else: - _unsloth_first_row = next(_unsloth_probe_cols, None) - if isinstance(_unsloth_first_row, dict): _unsloth_columns = list(_unsloth_first_row.keys()) - if _unsloth_columns is None and not _unsloth_transformed: - _unsloth_columns = getattr(train_dataset, 'column_names', None) - if isinstance(_unsloth_columns, dict): - _unsloth_columns = [_c for _v in _unsloth_columns.values() for _c in (_v or [])] - if _unsloth_columns is not None and ('input_ids' in _unsloth_columns or 'labels' in _unsloth_columns): - _unsloth_prep_truncates = False - except Exception: - _unsloth_prep_truncates = False - def _unsloth_truncatable(_ds): - if _ds is None or not hasattr(_ds, 'map'): return False - try: - if str((getattr(_ds, 'format', None) or {}).get('type', '')).lower() == 'custom': - return False - except Exception: - return False - try: - _cols = getattr(_ds, 'column_names', None) - if isinstance(_cols, dict): - _cols = [_c for _v in _cols.values() for _c in (_v or [])] - if 'seq_lengths' in (_cols or ()): return False - return bool(_cols) and 'input_ids' in _cols - except Exception: - return False - _unsloth_cap = args.max_length - _unsloth_truncation_mode = getattr(args, 'truncation_mode', 'keep_start') or 'keep_start' - _unsloth_keep_end = _unsloth_truncation_mode == 'keep_end' - _unsloth_known_mode = _unsloth_truncation_mode in ('keep_start', 'keep_end') - _unsloth_slice = slice(-_unsloth_cap, None) if _unsloth_keep_end else slice(None, _unsloth_cap) - _unsloth_eval_packing = getattr(args, 'packing', False) if getattr(args, 'eval_packing', None) is None else getattr(args, 'eval_packing') - _unsloth_completion_only = getattr(args, 'completion_only_loss', None) - if _unsloth_completion_only is None: - try: - _unsloth_fmt = getattr(train_dataset, 'format', None) - _unsloth_fmt = _unsloth_fmt if isinstance(_unsloth_fmt, dict) else {} - _unsloth_shown = _unsloth_fmt.get('columns') - if _unsloth_fmt.get('output_all_columns') or not _unsloth_shown: - _unsloth_shown = getattr(train_dataset, 'column_names', None) - _unsloth_train_sample = {} if _unsloth_is_transformed(train_dataset) else dict.fromkeys( - _unsloth_shown or []) - if not _unsloth_train_sample: - _unsloth_probe = iter(train_dataset) - if _unsloth_probe is train_dataset or _unsloth_probe is iter(train_dataset): - _unsloth_train_sample = {} - else: - _unsloth_train_sample = next(_unsloth_probe, None) or {} - except Exception: - _unsloth_train_sample = {} - _unsloth_completion_only = ('prompt' in _unsloth_train_sample and 'completion' in _unsloth_train_sample) - args._unsloth_completion_only_loss = _unsloth_completion_only - _UNSLOTH_SCAN_ROWS = 1024 - def _unsloth_within_cap(_ds): - if _ds is None: return True - _unsloth_own = getattr(_ds, '__dict__', None) - if isinstance(_unsloth_own, dict): - _unsloth_claim = _unsloth_own.get('_unsloth_truncated_to') - if isinstance(_unsloth_claim, int) and not isinstance(_unsloth_claim, bool): - return _unsloth_claim <= _unsloth_cap - try: - try: _n = len(_ds) - except Exception: _n = None - _unsloth_rows = iter(_ds) - if _unsloth_rows is iter(_ds): return False - _seen = 0 - for _row in _unsloth_rows: - if 'input_ids' not in _row: return True - if len(_row['input_ids']) > _unsloth_cap: return False - _seen += 1 - if _n is None and _seen >= _UNSLOTH_SCAN_ROWS: return False - except Exception: - return False - return True - def _unsloth_splits_within_cap(_ev): - _splits = list(_ev.values()) if isinstance(_ev, dict) else [_ev] - return all(_unsloth_within_cap(_s) for _s in _splits) - if not _unsloth_skip_prepare: - _unsloth_map_kw = {} - try: - try: - from unsloth_zoo.dataset_num_proc import get_dataset_num_proc as _unsloth_get_nproc - except ImportError: - from unsloth.dataset_num_proc import get_dataset_num_proc as _unsloth_get_nproc - _unsloth_nproc = _unsloth_get_nproc(getattr(args, 'dataset_num_proc', None)) - except Exception: - _unsloth_nproc = None - if _unsloth_nproc: _unsloth_map_kw['num_proc'] = _unsloth_nproc - def _unsloth_is_sequence_column(_col): - try: - if len(_col) == 0: return False - _first = _col[0] - except Exception: - return False - if isinstance(_first, (str, bytes)): return False - try: len(_first) - except Exception: return False - return True - def _unsloth_cut_value(_v, _r): - try: - if len(_v) != len(_r): return _v - except Exception: return _v - return _v[_unsloth_slice] - def _unsloth_truncate_rows(_batch): - _ids = _batch.get('input_ids') - _out = {} - for _k, _col in _batch.items(): - if not _unsloth_is_sequence_column(_col): - _out[_k] = _col - elif _ids is None: - _out[_k] = [_v[_unsloth_slice] for _v in _col] - else: - _out[_k] = [_unsloth_cut_value(_v, _r) for _v, _r in zip(_col, _ids)] - return _out - def _unsloth_is_stream(_ds): - try: return not hasattr(_ds, '__len__') - except Exception: return True - def _unsloth_pretokenized(_ds): - try: - _cols = None if _unsloth_is_transformed(_ds) else getattr(_ds, 'column_names', None) - if isinstance(_cols, dict): - _cols = [_c for _v in _cols.values() for _c in (_v or [])] - if _cols is not None: return 'input_ids' in _cols - _probe = iter(_ds) - if _probe is _ds or _probe is iter(_ds): return True - _row = next(_probe, None) - except Exception: return True - return isinstance(_row, dict) and 'input_ids' in _row - def _unsloth_rank_first(): - try: - from accelerate import PartialState - return PartialState().main_process_first() - except Exception: - import contextlib - return contextlib.nullcontext() - def _unsloth_cap_split(_ds): - with _unsloth_rank_first(): - return _unsloth_cap_one(_ds) - def _unsloth_cap_one(_ds): - if _ds is None: return _ds, True - if not _unsloth_truncatable(_ds): return _ds, not _unsloth_pretokenized(_ds) - _kw = {} if _unsloth_is_stream(_ds) else _unsloth_map_kw - _new = _ds.map(_unsloth_truncate_rows, batched = True, **_kw) - _unsloth_cols = getattr(_new, 'column_names', None) or () - _unsloth_masks = [] - if _unsloth_completion_only and 'completion_mask' in _unsloth_cols: - _unsloth_masks.append('completion_mask') - if 'assistant_masks' in _unsloth_cols: - _unsloth_masks.append('assistant_masks') - try: - _unsloth_supervision = (['labels'] if 'labels' in _unsloth_cols else []) + _unsloth_masks - if _unsloth_supervision: - _new = _new.filter(lambda _e, _c = tuple(_unsloth_supervision): any(all((_x != -100) if _n == 'labels' else _x for _n, _x in zip(_c, _v)) for _v in zip(*[_e[_n] for _n in _c])), **_kw) - except Exception: - pass - try: - if _unsloth_supervision and len(_new) == 0: _unsloth_emptied.append(1) - except TypeError: - pass - return _new, (True if _unsloth_is_stream(_new) else _unsloth_within_cap(_new)) - _unsloth_emptied = [] - _unsloth_orig_train = train_dataset - _unsloth_orig_eval = eval_dataset if 'eval_dataset' in locals() else None - try: - _unsloth_capped = _unsloth_known_mode - if not _unsloth_known_mode: - print('Unsloth: `truncation_mode = ' + str(_unsloth_truncation_mode) + '` is not one of keep_start / keep_end, so `max_length` is not being enforced here.') - if _unsloth_known_mode and not _unsloth_prep_truncates: - train_dataset, _unsloth_split_ok = _unsloth_cap_split(train_dataset) - _unsloth_capped = _unsloth_capped and _unsloth_split_ok - if _unsloth_eval_packing or not _unsloth_known_mode: - _unsloth_capped = False - elif 'eval_dataset' in locals() and eval_dataset is not None: - if isinstance(eval_dataset, dict): - _unsloth_new_eval = {} - for _k, _v in eval_dataset.items(): - _unsloth_new_eval[_k], _unsloth_split_ok = _unsloth_cap_split(_v) - _unsloth_capped = _unsloth_capped and _unsloth_split_ok - eval_dataset = _unsloth_new_eval - else: - eval_dataset, _unsloth_split_ok = _unsloth_cap_split(eval_dataset) - _unsloth_capped = _unsloth_capped and _unsloth_split_ok - _unsloth_prep_truncates = _unsloth_capped - if not _unsloth_capped: - print('Unsloth: `max_length` cannot be enforced for every split here, so padding-free batching is being turned off instead.') - except Exception as _unsloth_truncate_error: - train_dataset = _unsloth_orig_train - if 'eval_dataset' in locals(): eval_dataset = _unsloth_orig_eval - _unsloth_prep_truncates = False - print('Unsloth: could not truncate the pre-tokenized dataset to `max_length` (' + str(_unsloth_truncate_error) + ').') - if _unsloth_emptied: - raise ValueError('Unsloth: truncating to `max_length = ' + str(args.max_length) + '` left every row with no supervised token, so there is nothing to train on. The supervised part of your rows starts past that length: raise `max_length`, or set `truncation_mode = "keep_end"` if the completion sits at the end of each row.') - if not _unsloth_prep_truncates and not _unsloth_eval_packing: - def _unsloth_attests(_ds): - if _ds is None: return True - _own = getattr(_ds, '__dict__', None) - if not isinstance(_own, dict): return False - _claim = _own.get('_unsloth_truncated_to') - if not isinstance(_claim, int) or isinstance(_claim, bool): return False - return _claim <= _unsloth_cap - _unsloth_attest_eval = eval_dataset if 'eval_dataset' in locals() else None - _unsloth_attest_splits = list(_unsloth_attest_eval.values()) if isinstance(_unsloth_attest_eval, dict) else [_unsloth_attest_eval] - if _unsloth_attests(train_dataset) and all(_unsloth_attests(_s) for _s in _unsloth_attest_splits): - _unsloth_prep_truncates = True - if _unsloth_prep_truncates: - args.max_seq_length = args.max_length - args.max_length = None - max_length = None - else: - _unsloth_scan_eval = None if _unsloth_eval_packing else (eval_dataset if 'eval_dataset' in locals() else None) - if not (_unsloth_within_cap(train_dataset) and _unsloth_splits_within_cap(_unsloth_scan_eval)): - raise ValueError('Unsloth: `max_length = ' + str(args.max_length) + '` cannot be enforced. Your dataset already carries `input_ids` and holds rows longer than that, and nothing downstream truncates pre-tokenized rows. Truncate it yourself before passing it in, or drop `max_length`.') - print('Unsloth: Turning padding-free batching off, since your dataset is already tokenized and cannot be truncated here. Padding-free batching cannot enforce a `max_length` of ' + str(args.max_length) + '.') - args.padding_free = False - if model is not None and hasattr(model, 'for_training'): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right' - if 'processing_class' in locals(): - if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right' - if hasattr(processing_class, 'tokenizer') and hasattr(processing_class.tokenizer, 'padding_side'): processing_class.tokenizer.padding_side = 'right' - __tokenizer = processing_class if 'processing_class' in locals() else tokenizer - from unsloth_zoo.vision_utils import UnslothVisionDataCollator - if not isinstance(data_collator, UnslothVisionDataCollator): - if isinstance(data_collator, DataCollatorForSeq2Seq) and 'labels' not in _unsloth_dataset_column_names(train_dataset): - data_collator = TransformersDataCollatorForLanguageModeling( - __tokenizer, - mlm = False, - mlm_probability = 0.0, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling) and 'labels' in _unsloth_dataset_column_names(train_dataset): - data_collator = DataCollatorForSeq2Seq( - __tokenizer, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - else: - if hasattr(args, 'remove_unused_columns'): args.remove_unused_columns = False - if hasattr(args, 'dataset_text_field'): args.dataset_text_field = '' - if hasattr(args, 'dataset_kwargs'): args.dataset_kwargs = {'skip_prepare_dataset': True} - if not isinstance(data_collator, UnslothVisionDataCollator): - if not hasattr(__tokenizer, 'pad') and hasattr(__tokenizer, 'tokenizer'): - if isinstance(data_collator, DataCollatorForSeq2Seq): - data_collator = DataCollatorForSeq2Seq( - __tokenizer.tokenizer, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling): - data_collator = TransformersDataCollatorForLanguageModeling( - __tokenizer.tokenizer, - mlm = False, - mlm_probability = 0.0, - pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None), - ) - other_metrics = [] - - from unsloth_zoo.logging_utils import PatchRLStatistics - PatchRLStatistics('sft_trainer', other_metrics) - IGNORED_TOKENIZER_NAMES = os.environ.get('UNSLOTH_IGNORED_TOKENIZER_NAMES', '').split('\n') - from unsloth_zoo.tokenizer_utils import fix_untrained_tokens - from unsloth_zoo.training_utils import fix_zero_training_loss - if 'tokenizer' not in locals(): tokenizer = processing_class - fix_untrained_tokens(model, tokenizer, train_dataset, IGNORED_TOKENIZER_NAMES, eps = 1e-16) - fix_zero_training_loss(model, tokenizer, train_dataset) - - # [TODO] Fix up DataParallel multiplying batch sizes - # [TODO] DDP works, but DP seems to not work? [TODO] - if getattr(args, "parallel_mode", None) == ParallelMode.NOT_DISTRIBUTED and args.n_gpu > 1: - if getattr(args, "_n_gpu", 1) != 1: - args._n_gpu = 1 - if "model" in locals(): - try: - from unsloth.models._gc_config_default import honour_gradient_checkpointing_off - honour_gradient_checkpointing_off(args, model) - except Exception: - pass - if "model" in locals() and hasattr(model, "for_training"): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if "model" in locals(): - try: - from unsloth.models._ddp_defaults import unsloth_default_ddp_find_unused_parameters - unsloth_default_ddp_find_unused_parameters(args, model) - except Exception: - pass - super().__init__( - model = model, - args = args, - data_collator = data_collator, - train_dataset = train_dataset, - eval_dataset = eval_dataset, - processing_class = processing_class, - compute_loss_func = compute_loss_func, - compute_metrics = compute_metrics, - callbacks = callbacks, - optimizer_cls_and_kwargs = optimizer_cls_and_kwargs, - preprocess_logits_for_metrics = preprocess_logits_for_metrics, - quantization_config = quantization_config, - peft_config = peft_config, - formatting_func = formatting_func,**kwargs) - # A tensor parallel model (tp_native.tp_trainer_setup) stays on its training path: for_inference would - # route the decoder layers onto decode paths outside the tensor parallel regions. - if "model" in locals() and hasattr(model, "for_inference") and not getattr(self, "_unsloth_tensor_parallel", False): - model.for_inference() - if hasattr(self, 'neftune_hook_handle'): - self.neftune_hook_handle.remove() - if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle - if getattr(args, 'neftune_noise_alpha', None) is not None: - model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha - pass - if hasattr(self, 'accelerator'): - scaler = self.accelerator.scaler - current_model = model - while hasattr(current_model, 'model'): - current_model.accelerator_scaler = scaler - current_model = current_model.model - current_model.accelerator_scaler = scaler - pass - if hasattr(self, 'train'): - self.train = MethodType(prepare_for_training_mode(self.__class__.train), self) - pass - if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'): - _vllm_tok = self.llm.get_tokenizer() - _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None) - if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None: - _vllm_tok.chat_template = _pc.chat_template - pass - if getattr(self, 'aux_loss_enabled', False) and hasattr(getattr(self, 'model', None), 'config'): - _text_config = self.model.config - if hasattr(_text_config, 'get_text_config'): _text_config = _text_config.get_text_config() - _n_experts = [getattr(_text_config, _k) for _k in ('num_local_experts', 'num_experts', 'n_routed_experts', 'moe_num_experts') if isinstance(getattr(_text_config, _k, None), int)] - if _n_experts and all(_n <= 1 for _n in _n_experts): - self.aux_loss_enabled = False - _text_config.output_router_logits = False - pass - if hasattr(self, 'aux_loss_enabled') and hasattr(getattr(self, 'model', None), 'modules'): - for _module in self.model.modules(): - _config = getattr(_module, 'config', None) - if 'router_aux_loss_coef' in vars(_module) and hasattr(_config, 'get_text_config'): - _coef = getattr(_config.get_text_config(), 'router_aux_loss_coef', None) - if _coef is not None: - _module.router_aux_loss_coef = _coef - pass - -pass - - -if hasattr(logger, "addFilter"): - import logging - class HideLoggingMessage(logging.Filter): - def __init__(self, text): self.text = text - def filter(self, x): return not (self.text in x.getMessage()) - pass - logger.addFilter(HideLoggingMessage("`use_cache=True`")) - diff --git a/unsloth_compiled_cache/UnslothServerDistillationTrainer.py b/unsloth_compiled_cache/UnslothServerDistillationTrainer.py deleted file mode 100644 index ef36d91910b28fce94f2ca06f0a8fc61682187ea..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/UnslothServerDistillationTrainer.py +++ /dev/null @@ -1,1576 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from unsloth_zoo.temporary_patches.common import torch_compile -from unsloth_zoo.temporary_patches.common import _maybe_compile -import functools -import transformers -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from trl.experimental.server_distillation.server_distillation_trainer import (Any, DistillationTrainer, F, ServerDistillationConfig, ServerDistillationTrainer, _add_tail_bucket, _jsd_divergence, build_teacher_request_inputs, profiling_decorator, torch) - - -import os -import math -import logging -from typing import * -from dataclasses import dataclass, field -from packaging.version import Version -import torch -import numpy as np -from contextlib import nullcontext -from torch.nn import functional as F -import inspect -from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling -from transformers.training_args import ParallelMode -from unsloth_zoo.device_type import DEVICE_TYPE, DEVICE_TYPE_TORCH, device_synchronize - -# Wrap trainer with padding to right and enable training mode -import functools -from types import MethodType -try: - from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers -except: - def reset_unsloth_gradient_checkpointing_buffers(): pass -# Canonical reset lives in unsloth.models._utils so the SFT auto-packing wrapper and the plain -# Trainer loop can import the same helper; fall back to a no-op only if it can't be imported. -try: - from unsloth.models._utils import _unsloth_reset_stray_compile_cache -except Exception: - def _unsloth_reset_stray_compile_cache(self): pass -try: - from unsloth.models._utils import _unsloth_dataset_column_names -except Exception: - def _unsloth_dataset_column_names(dataset): return dataset.column_names -# Drops/renames config arguments the installed TRL no longer accepts, so a -# script pinned to an older TRL keeps working after an upgrade. Falls back to -# the historical raw passthrough so this can never break trainer construction. -try: - from unsloth.models.rl_config_compat import filter_config_init_kwargs as _unsloth_filter_config_init_kwargs - # A cache file generated here can be imported by an older Unsloth whose filter - # predates `mirrored_from`, so drop the argument rather than raise TypeError. - if "mirrored_from" not in inspect.signature(_unsloth_filter_config_init_kwargs).parameters: - _unsloth_filter_config_init_kwargs_old = _unsloth_filter_config_init_kwargs - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): - return _unsloth_filter_config_init_kwargs_old(config_class, kwargs) -except Exception: - def _unsloth_filter_config_init_kwargs(config_class, kwargs, **kw): return kwargs -def prepare_for_training_mode(f): - @functools.wraps(f) - def wrapper(self, *args, **kwargs): - # Drop any torch.compile graph cache poisoned by a stray pre-train forward. - try: - _unsloth_reset_stray_compile_cache(self) - except Exception: - pass - # Finish the previous W&B run if this is a subsequent train() call. - # We do this at the START of train() (not the end) so that - # evaluate() / log() still work after train() completes. - # HF's WandbCallback.setup() will call wandb.init() for the new run. - # See: https://github.com/unslothai/unsloth/issues/3954 - if getattr(self, '_unsloth_training_completed', False): - try: - import wandb - if wandb.run is not None: - wandb.finish() - # Reset HF's WandbCallback so it calls wandb.init() for the new run - for cb in self.callback_handler.callbacks: - if type(cb).__name__ == 'WandbCallback': - cb._initialized = False - break - except: - pass - # Enable training mode - _was_training = None - # Restore the GC mode the model was configured with at setup; fall back to - # the training args only when it wasn't recorded (issue #4735). Use hasattr, - # not a None sentinel, so a deliberately-recorded None is restored verbatim. - _model = getattr(self, 'model', None) - if hasattr(_model, '_unsloth_gradient_checkpointing'): - use_gc = _model._unsloth_gradient_checkpointing - else: - use_gc = getattr(self.args, 'gradient_checkpointing', True) - if hasattr(self, 'model') and hasattr(self.model, "training"): - _was_training = self.model.training - if hasattr(self, 'model') and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - output = f(self, *args, **kwargs) - # Restore previous mode when possible - if hasattr(self, 'model') and hasattr(self.model, "for_inference"): - if _was_training is False: - self.model.for_inference() - elif _was_training is True and hasattr(self.model, "for_training"): - self.model.for_training(use_gradient_checkpointing=use_gc) - # Reset gradient checkpointing buffers to free memory while staying ready for next run - try: - reset_unsloth_gradient_checkpointing_buffers() - except: - pass - # Mark that training completed so the next train() call can - # finish this W&B run before starting a new one - self._unsloth_training_completed = True - return output - return wrapper -pass - -torch_compile_options = { - "epilogue_fusion" : True, - "max_autotune" : False, - "shape_padding" : True, - "trace.enabled" : False, - "triton.cudagraphs" : False, -} - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_hidden_states_selective_log_softmax( - hidden_states: torch.Tensor, - lm_head: torch.Tensor, - index: torch.Tensor, - chunks: int = 4, - logit_scale_multiply: float = 0.0, - logit_scale_divide: float = 0.0, - logit_softcapping: float = 0.0, - temperature: float = 1.0, - # Rows per chunk cap. Read HERE, in the default, not in the body: the body - # is traced with fullgraph = True, and `os.environ` is an unsupported op - # there on torch 2.4 -- Dynamo raises - # torch._dynamo.exc.Unsupported: const method call bytes.decode - # from os._Environ.__getitem__, which the eager fallback does not catch - # (it only catches recompile-limit and disabled-hook breaks), so the very - # first call would die even with the variable unset. A default is evaluated - # once when the def runs, which is import time, outside any traced region. - # It is also a plain int argument, so Dynamo guards on it instead of - # constant-folding an unguarded read (2.7+ never notice a later change). - # A non-numeric value is ignored rather than raised on, so a typo cannot - # break the import. 0 keeps the previous chunk boundaries exactly. - max_rows_per_chunk: int = ( - int(os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip()) - if os.environ.get("UNSLOTH_GRPO_MAX_ROWS_PER_CHUNK", "0").strip().isdigit() - else 0 - ), - lm_head_bias: Optional[torch.Tensor] = None, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, - # Read in the default for the same Dynamo reason as max_rows_per_chunk. Pairs with vLLM's - # head_dtype=float32 that load_vllm requests under the same flag. - fp32_lm_head: bool = ( - os.environ.get("UNSLOTH_VLLM_FP32_LM_HEAD", "0").strip().lower() in ("1", "true", "on", "yes") - ), - # Older Unsloth scores old / ref log-probs here and the policy with the fused kernel: one scorer keeps ratio at 1. - fused_linear: bool = ( - os.environ.get("UNSLOTH_FUSED_LINEAR", "1") != "0" - and os.environ.get("UNSLOTH_GRPO_CHUNKED_LOGPROBS", "1") != "0" - ), -) -> torch.Tensor: - # All Unsloth Zoo code licensed under AGPL3 - from unsloth_zoo.fused_losses import chunked_logprob as _chunked_logprob - if (fused_linear and _chunked_logprob._FUSED_LINEAR and not fp32_lm_head - and support_top_k == 0 and support_top_p == 1.0 and support_min_p == 0.0 - and support_excluded_ids is None): - return _chunked_logprob.unsloth_chunked_logprobs( - hidden_states, lm_head, index, - temperature = temperature, - logit_scale_multiply = logit_scale_multiply, - logit_scale_divide = logit_scale_divide, - logit_softcapping = logit_softcapping, - lm_head_bias = lm_head_bias, - max_rows_per_chunk = max_rows_per_chunk if max_rows_per_chunk > 0 else None, - ).reshape((hidden_states.shape[0], hidden_states.shape[1])) - # Reshape on this tensor's own last dim: a no-op, so a wrong-width caller - # cannot have its row count silently rewritten and instead fails at the - # matmul below, which prints both operands. Do not swap in a bare - # torch._check: it reports only "Expected cond to be True", naming neither - # operand, and Dynamo rejects a message-carrying one. Callers dispatch on - # the width first -- see `compute_logprobs_chunk`, the packed path and - # `_pg_grad_forward`. - flat_hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) - flat_index = index.reshape(-1) - - # Each chunk materialises rows x vocab logits and then a float32 copy of - # them, all on the device holding the output head. With a large vocabulary - # and a fixed chunk count that grows with the batch, so the peak scales with - # the batch rather than staying bounded. max_rows_per_chunk caps the rows - # per chunk instead, which is pure loop splitting: more, smaller chunks, - # same concatenated result. 0 (the default) keeps the previous chunk - # boundaries exactly. - if max_rows_per_chunk > 0: - n_rows = flat_hidden_states.shape[0] - chunks = max(chunks, -(-n_rows // max_rows_per_chunk)) - chunks = min(chunks, max(n_rows, 1)) - - # A one-row chunk's head gradient is a K=1 matmul Inductor mis-lowers under dynamic shapes (illegal - # memory access or NaN head gradients, torch 2.13), so keep chunks at 2+ rows when the head trains. - # Exactly `chunks` row ranges, so the unrolled loop length never depends on T: torch.chunk - # returns ceil(T / ceil(T / chunks)) pieces and recompiled per new count; tensor_split guards - # on T % chunks. At most rows // 2 chunks when the head trains keeps every piece at 2+ rows. - lm_head_grad = torch.is_grad_enabled() and lm_head.requires_grad - if lm_head_grad and flat_hidden_states.shape[0] < 2 * chunks: - chunks = max(flat_hidden_states.shape[0] // 2, 1) - - n_rows = flat_hidden_states.shape[0] - chunks = max(1, min(chunks, n_rows)) - bounds = [(i * n_rows) // chunks for i in range(chunks + 1)] - chunked_hidden_states = [flat_hidden_states[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - - # Cast once, outside the loop, so every chunk shares one fp32 head copy (also for autograd). - if fp32_lm_head: - lm_head = lm_head.to(torch.float32) - if lm_head_bias is not None: - lm_head_bias = lm_head_bias.to(torch.float32) - - all_per_token_logps = [] - - for chunk_hidden_states, chunk_index in zip(chunked_hidden_states, chunked_index): - # When the model is dispatched over several devices, the output head can - # sit on a different one from the hidden states, because accelerate - # places the tail of the model on the last device it fills. Co-locate on - # the head's device before the matmul, otherwise this raises - # Unhandled FakeTensor Device Propagation for aten.mm.default, - # found two different devices cuda:0, cuda:1 - # On a single device every .to() here is a no-op and the result is - # bit-identical to before. - chunk_hidden_states = chunk_hidden_states.to(device = lm_head.device, dtype = lm_head.dtype) - chunk_index = chunk_index.to(lm_head.device) - if fp32_lm_head: - # The GRPO loss runs this under bf16 / fp16 autocast, which would downcast the matmul. - with torch.autocast(device_type = chunk_hidden_states.device.type, enabled = False): - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - elif lm_head_bias is None: - chunk_logits = chunk_hidden_states @ lm_head.t() - else: - chunk_logits = torch.nn.functional.linear(chunk_hidden_states, lm_head, lm_head_bias) - - if logit_scale_multiply != 0.0: - chunk_logits = chunk_logits * logit_scale_multiply - if logit_scale_divide != 0.0: - chunk_logits = chunk_logits / logit_scale_divide - if logit_softcapping != 0.0: - chunk_logits = logit_softcapping * torch.tanh(chunk_logits / logit_softcapping) - - chunk_logits = chunk_logits.to(torch.float32) - - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - per_token_logps = _truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - ) - else: - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - # Return to the caller's device so the concatenation below and every - # downstream consumer see the device they started on. - all_per_token_logps.append(per_token_logps.to(hidden_states.device)) - - all_per_token_logps = torch.concat(all_per_token_logps) - - all_per_token_logps = all_per_token_logps.reshape((hidden_states.shape[0], hidden_states.shape[1])) - return all_per_token_logps - -@_maybe_compile(dynamic = True, fullgraph = True, options = torch_compile_options,) -def chunked_selective_log_softmax( - logits, - index, - temperature: float = 1.0, - chunks: int = 4, - support_top_k: int = 0, - support_top_p: float = 1.0, - support_min_p: float = 0.0, - support_triton_sampler: bool = False, - support_excluded_ids: Optional[torch.Tensor] = None, -): - # Fixed `chunks` row ranges; see chunked_hidden_states_selective_log_softmax. - flat_logits = logits.reshape(-1, logits.shape[-1]) - flat_index = index.reshape(-1) - # Fewer rows than chunks: one row per range, as torch.chunk did, never empty ranges. - chunks = max(1, min(chunks, flat_logits.shape[0])) - bounds = [(i * flat_logits.shape[0]) // chunks for i in range(chunks + 1)] - chunked_logits = [flat_logits[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - chunked_index = [flat_index[a:b] for a, b in zip(bounds[:-1], bounds[1:])] - all_per_token_logps = [] - # Per-chunk selective_log_softmax. - for chunk_logits, chunk_index in zip(chunked_logits, chunked_index): - chunk_logits = chunk_logits.to(torch.float32) - if temperature != 1.0: - chunk_logits = chunk_logits / temperature - if support_top_k > 0 or support_top_p < 1.0 or support_min_p > 0.0 or support_excluded_ids is not None: - all_per_token_logps.append(_truncated_support_logps( - chunk_logits, chunk_index, support_top_k, support_top_p, support_min_p, - support_triton_sampler, support_excluded_ids, - )) - continue - # Target excluded from the normaliser: x_i - logsumexp(x) drops confident tokens (softmax collapse). - is_target = torch.arange(chunk_logits.shape[-1], device = chunk_logits.device) == chunk_index.unsqueeze(-1) - selected_logits = torch.where(is_target, chunk_logits, 0.0).sum(-1) - others = torch.where(is_target, torch.finfo(chunk_logits.dtype).min, chunk_logits) - per_token_logps = -torch.nn.functional.softplus(torch.logsumexp(others, dim = -1) - selected_logits) - all_per_token_logps.append(per_token_logps) - pass - all_per_token_logps = torch.concat(all_per_token_logps) - all_per_token_logps = all_per_token_logps.reshape((logits.shape[0], logits.shape[1])) - return all_per_token_logps - -def calculate_pad_tokens_in_prompt( - input_ids: torch.Tensor, - logits_to_keep: int, - pad_token_id: int -) -> torch.Tensor: - """Count left-padded tokens per sequence, e.g. [pad, pad, pad, cat] -> 3.""" - if logits_to_keep >= input_ids.shape[1]: - raise ValueError("logits_to_keep must be smaller than the sequence length.") - - prompt_section = input_ids[:, :-logits_to_keep] - - padding_mask = (prompt_section == pad_token_id) - - pad_token_counts = padding_mask.sum(dim=1) - - return pad_token_counts - -def create_completion_attention_mask( - completion_input_ids: torch.Tensor, - left_pad_tokens_per_prompt: torch.Tensor, - max_left_pad: int, - pad_token_id: int -) -> torch.Tensor: - """Build a completion mask that zeros leading prompt and trailing pad tokens. - - For [p,p,p,c,c,c,pad,pad,pad] (p=sliced prompt, c=completion, pad=padding) - this returns [0,0,0,1,1,1,0,0,0]. - """ - batch_size, completion_len = completion_input_ids.shape - device = completion_input_ids.device - - num_tokens_to_mask = max_left_pad - left_pad_tokens_per_prompt - - indices = torch.arange(completion_len, device=device).unsqueeze(0) - shift_mask = indices >= num_tokens_to_mask.unsqueeze(1) - - non_padding_mask = (completion_input_ids != pad_token_id) - - final_mask = shift_mask & non_padding_mask - - return final_mask - -def left_pack_padding(tensor: torch.Tensor, pad_id: int) -> torch.Tensor: - """Move all padding tokens in each sequence to the right.""" - mask = (tensor != pad_id) - # stable=True since the binary mask is unordered. - sorted_indices = torch.argsort(mask, dim=1, descending=True, stable=True) - packed_tensor = torch.gather(tensor, 1, sorted_indices) - return packed_tensor - -def align_logprobs_with_mask( - logprob_tensor: torch.Tensor, - attention_mask: torch.Tensor, - pad_value: float = 0.0 -) -> torch.Tensor: - """Align a log probability tensor with a given attention mask.""" - - device = logprob_tensor.device - batch_size, logprob_seq_len = logprob_tensor.shape - mask_seq_len = attention_mask.shape[1] - - padded_logprobs = torch.full( - attention_mask.shape, - fill_value=pad_value, - dtype=logprob_tensor.dtype, - device=device - ) - - left_pad_counts = torch.argmax(attention_mask, dim=1) - - cols = torch.arange(logprob_seq_len, device=device) - dest_indices = left_pad_counts.unsqueeze(1) + cols - - # Destination row indices, shape [batch_size, logprob_seq_len]. - row_indices = torch.arange(batch_size, device=device).unsqueeze(1).expand_as(dest_indices) - - # Keep only in-bounds destinations, then scatter via advanced indexing. - valid_mask = dest_indices < mask_seq_len - valid_rows = row_indices[valid_mask] - valid_cols = dest_indices[valid_mask] - valid_vals = logprob_tensor[valid_mask] - padded_logprobs[valid_rows, valid_cols] = valid_vals - - return padded_logprobs - -def align_completion_tool_mask( - tool_mask: torch.Tensor, - completion_mask: torch.Tensor, -) -> torch.Tensor: - """Align a raw completion-length tool/env mask with Unsloth's repacked loss mask.""" - if tool_mask is None: - return completion_mask - if tool_mask.shape[0] != completion_mask.shape[0]: - raise ValueError("tool_mask batch size must match completion_mask batch size.") - - tool_mask = tool_mask.to(device=completion_mask.device) - if tool_mask.shape == completion_mask.shape: - aligned_tool_mask = tool_mask - else: - aligned_tool_mask = align_logprobs_with_mask( - tool_mask, - completion_mask, - pad_value=0, - ) - return completion_mask * aligned_tool_mask.to(dtype=completion_mask.dtype) - -def autotune_batch_and_chunks( - total_input_rows, - seq_len, - hidden_size, - vocab_size, - dtype_bytes=16, - multiplier=None -): - if multiplier is None: - final_m = max(4, seq_len // 4096) - else: - final_m = multiplier - - if torch.cuda.is_available(): - free_bytes, _ = torch.cuda.mem_get_info() - limit_gb = (free_bytes / (1024**3))*.80 - elif hasattr(torch, "xpu") and torch.xpu.is_available(): - # XPU: estimate free memory as total - reserved. - total_mem = torch.xpu.get_device_properties(0).total_memory - reserved_mem = torch.xpu.memory_reserved() - free_bytes = total_mem - reserved_mem - limit_gb = (free_bytes / (1024**3)) * 0.80 - else: - # Fallback: assume 8GB available. - limit_gb = 8.0 - - bytes_to_gb = 1024**3 - - b_vals = torch.arange(total_input_rows, 0, -1, device='cpu', dtype=torch.float32) - - hidden_gb = (b_vals * seq_len * hidden_size * dtype_bytes) / bytes_to_gb - - base_logits = ((b_vals/total_input_rows) * b_vals * seq_len * vocab_size * dtype_bytes) / bytes_to_gb - logits_gb = base_logits / final_m - - total_mem_gb = hidden_gb + logits_gb - - valid_mask = total_mem_gb <= limit_gb - valid_indices = torch.nonzero(valid_mask, as_tuple=False) - - if valid_indices.shape[0] == 0: - #This means your GPU will OOM - # Capped at the row count: unsloth's no-grad pass divides rows by this without max(1, ...). - return max(1, min(4, total_input_rows)), final_m - - best_idx = valid_indices[0].item() - final_b = int(b_vals[best_idx].item()) - - return final_b, final_m - -def sanitize_logprob(logprob): - """Local port of trl.scripts.vllm_serve.sanitize_logprob. - Filters NaN logprobs from vLLM outputs.""" - value = logprob.logprob - if math.isnan(value): - logging.getLogger(__name__).warning( - f"Generated NaN logprob, token logprob '{logprob}' will be ignored" - ) - return None - return value -@dataclass -class UnslothServerDistillationConfig(ServerDistillationConfig): - """ - - Configuration class for the [`experimental.server_distillation.ServerDistillationTrainer`]. - - Extends [`DistillationConfig`] with the address of an external vLLM server that scores the student's completions. - The teacher is never held locally: instead of a local forward pass, per-token teacher logprobs are fetched from the - server, so only the teacher's top-k logprobs are available and the loss is restricted to a sparse support. - - Parameters: - teacher_model_server_url (`str` or `None`, *optional*): - Base URL of a vLLM server hosting the teacher model (e.g., `"http://localhost:8000"`). Required. - vllm_server_port (`int`, *optional*, defaults to `8001`): - Port of the student vLLM server. Overrides the base default of `8000` to stay distinct from the teacher - server, which commonly runs on `8000`. - loss_top_k (`int`, *optional*, defaults to `1`): - Number of top tokens the teacher server returns and over which the divergence is computed. For `beta == 0` - (pure forward KL) this must be positive; for `beta > 0` it must be `1` (top-1 support). `0` (full - vocabulary) is not available with a server teacher, which only exposes top-k logprobs. - loss_add_tail (`bool`, *optional*, defaults to `True`): - Whether to append a tail bucket representing the remaining probability mass outside the selected top-k - support when computing the loss. - reverse_kl_top_1_mode (`str`, *optional*, defaults to `"sampled"`): - Selection rule for the reverse-KL top-1 token when `beta > 0` and `loss_top_k == 1`. `"sampled"` uses the - actual completion token in the batch. `"argmax"` uses the student's highest-probability token (not - supported by the server, which cannot score arbitrary student-selected tokens). This setting does not - affect the forward-KL support, which always uses the teacher's top-1 token. - - """ - _unsloth_accepts_adamw_flash = True - vllm_sampling_params: Optional[Any] = field( - default = None, - metadata = {'help': 'vLLM SamplingParams'}, - ) - unsloth_num_chunks : Optional[int] = field( - default = -1, - metadata = {'help': 'Chunk size to reduce memory usage. -1 is most efficient.'}, - ) - unsloth_logit_chunk_multiplier : Optional[int] = field( - default = None, - metadata = {'help': 'Multiplier for chunked logit computations.'}, - ) - unsloth_grpo_mini_batch : Optional[int] = field( - default = None, - metadata = {'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'}, - ) - - def __init__( - self, - output_dir = None, - per_device_train_batch_size = 4, - num_train_epochs = 3.0, - max_steps = -1, - learning_rate = 5e-05, - lr_scheduler_type = 'linear', - lr_scheduler_kwargs = None, - warmup_steps = 0.1, - optim = 'adamw_8bit', - optim_args = None, - weight_decay = 0.001, - adam_beta1 = 0.9, - adam_beta2 = 0.999, - adam_epsilon = 1e-08, - optim_target_modules = None, - gradient_accumulation_steps = 2, - average_tokens_across_devices = True, - max_grad_norm = 1.0, - label_smoothing_factor = 0.0, - bf16 = False, - fp16 = False, - bf16_full_eval = False, - fp16_full_eval = False, - tf32 = None, - gradient_checkpointing = True, - gradient_checkpointing_kwargs = None, - torch_compile = False, - torch_compile_backend = None, - torch_compile_mode = None, - use_liger_kernel = False, - liger_kernel_config = None, - use_cache = False, - neftune_noise_alpha = None, - torch_empty_cache_steps = 250, - auto_find_batch_size = False, - logging_strategy = 'steps', - logging_steps = 1, - logging_first_step = False, - log_on_each_node = True, - logging_nan_inf_filter = False, - include_num_input_tokens_seen = False, - log_level = 'passive', - log_level_replica = 'warning', - disable_tqdm = None, - report_to = 'none', - run_name = None, - project = 'huggingface', - trackio_space_id = 'trackio', - eval_strategy = 'no', - eval_steps = None, - eval_delay = 0, - per_device_eval_batch_size = 4, - prediction_loss_only = False, - eval_on_start = False, - eval_do_concat_batches = True, - eval_use_gather_object = False, - eval_accumulation_steps = 2, - batch_eval_metrics = False, - save_only_model = False, - save_strategy = 'steps', - save_steps = 500, - save_on_each_node = False, - save_total_limit = None, - enable_jit_checkpoint = False, - push_to_hub = False, - hub_token = None, - hub_private_repo = None, - hub_model_id = None, - hub_strategy = 'every_save', - hub_always_push = False, - hub_revision = None, - load_best_model_at_end = False, - metric_for_best_model = None, - greater_is_better = None, - ignore_data_skip = False, - restore_callback_states_from_checkpoint = False, - full_determinism = False, - seed = 3407, - data_seed = None, - use_cpu = False, - accelerator_config = None, - parallelism_config = None, - dataloader_drop_last = False, - dataloader_num_workers = 0, - dataloader_pin_memory = True, - dataloader_persistent_workers = False, - dataloader_prefetch_factor = None, - remove_unused_columns = False, - label_names = None, - train_sampling_strategy = 'random', - length_column_name = 'length', - ddp_find_unused_parameters = None, - ddp_bucket_cap_mb = None, - ddp_broadcast_buffers = None, - ddp_backend = None, - ddp_timeout = 1800, - fsdp = None, - fsdp_config = None, - deepspeed = None, - debug = '', - skip_memory_metrics = True, - do_train = False, - do_eval = False, - do_predict = False, - resume_from_checkpoint = None, - warmup_ratio = None, - logging_dir = None, - local_rank = -1, - model_init_kwargs = None, - trust_remote_code = False, - teacher_model_name_or_path = None, - teacher_model_revision = None, - teacher_model_init_kwargs = None, - disable_dropout = False, - max_completion_length = 512, - ds3_gather_for_generation = True, - shuffle_dataset = True, - pad_to_multiple_of = None, - temperature = 1.0, - top_p = 1.0, - top_k = 0, - min_p = None, - generation_kwargs = {}, - chat_template_kwargs = None, - repetition_penalty = 1.0, - cache_implementation = None, - use_vllm = False, - vllm_mode = 'colocate', - vllm_model_impl = 'vllm', - vllm_enable_sleep_mode = False, - vllm_structured_outputs_regex = None, - vllm_server_base_url = None, - vllm_server_host = '0.0.0.0', - vllm_server_port = 8001, - vllm_server_timeout = 240.0, - vllm_group_port = 51216, - vllm_gpu_memory_utilization = 0.3, - vllm_max_model_length = None, - vllm_tensor_parallel_size = 1, - beta = 1.0, - max_tool_calling_iterations = None, - log_completions = False, - num_completions_to_print = None, - log_unique_prompts = False, - teacher_model_server_url = None, - loss_top_k = 1, - loss_add_tail = True, - reverse_kl_top_1_mode = 'sampled', - vllm_sampling_params = None, - unsloth_num_chunks = -1, - unsloth_logit_chunk_multiplier = None, - unsloth_grpo_mini_batch = None, - - **kwargs, - ): - if learning_rate < 1e-7: print(f'Unsloth: Your learning rate of `{learning_rate}` is too small and less than 1e-7! Consider increasing it, otherwise gradient updates will be close to 0!') - if learning_rate > 1: print(f'Unsloth: Your learning rate of `{learning_rate}` is way too larger > 1! Consider decreasing it to 1e-1, otherwise gradient updates will explode!') - _unsloth_fsdp2 = os.environ.get('FSDP_VERSION', '').strip() == '2' - _unsloth_fsdp_cfg = fsdp_config - if fsdp and isinstance(_unsloth_fsdp_cfg, str) and os.path.isfile(_unsloth_fsdp_cfg): - try: - import json as _unsloth_json - with open(_unsloth_fsdp_cfg, encoding = 'utf-8') as _unsloth_f: - _unsloth_fsdp_cfg = _unsloth_json.load(_unsloth_f) - except Exception: - _unsloth_fsdp_cfg = None - if fsdp and _unsloth_fsdp_cfg is None: - _unsloth_fsdp_cfg = {} - if fsdp and isinstance(_unsloth_fsdp_cfg, dict): - _unsloth_fsdp2 = _unsloth_fsdp2 or str(_unsloth_fsdp_cfg.get('fsdp_version', _unsloth_fsdp_cfg.get('version', 1))) == '2' - _unsloth_flash_adamw = str(getattr(optim, 'value', optim)).strip().lower() == 'adamw_flash' - if _unsloth_flash_adamw: optim = 'adamw_torch' - if _unsloth_fsdp2 and str(getattr(optim, 'value', optim)) == 'adamw_8bit': - print("Unsloth: FSDP2 shards parameters as DTensors, which bitsandbytes' 8-bit optimizers cannot step. Switching optim = 'adamw_8bit' to 'adamw_torch_fused'.") - optim = 'adamw_torch_fused' - if num_train_epochs is None: - num_train_epochs = 3.0 # Default to 3 epochs if None, max_steps will override - if output_dir is None and save_strategy == 'steps' and save_steps == 500: - output_dir = 'unsloth_training_checkpoints' - save_strategy = 'no' - if os.environ.get('UNSLOTH_ENABLE_FLEX_ATTENTION', '0') == '1': - from unsloth_zoo.flex_attention import HAS_FLEX_ATTENTION - if HAS_FLEX_ATTENTION and pad_to_multiple_of is None: - from unsloth_zoo.flex_attention import FLEX_ATTENTION_BLOCK_SIZE - pad_to_multiple_of = FLEX_ATTENTION_BLOCK_SIZE - - if eval_steps is not None and eval_strategy != 'steps': - print(f'Unsloth: `eval_steps = {eval_steps}` is ignored because `eval_strategy` is {getattr(eval_strategy, "value", eval_strategy)!r}. Set `eval_strategy = "steps"` to evaluate every `eval_steps` steps.') - - if temperature <= 0: - raise ValueError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.') - elif temperature >= 10: - raise ValueError('Unsloth: Please set a positive non-zero temperature less than 10, since sampling will be quite erratic.') - - - # One dict so the filter sees the mirrored parameters AND `**kwargs`: - # filtering kwargs alone would double-bind any argument TRL renamed, - # since the new name is itself a mirrored parameter. - _unsloth_config_arguments = dict( - output_dir = output_dir, - per_device_train_batch_size = per_device_train_batch_size, - num_train_epochs = num_train_epochs, - max_steps = max_steps, - learning_rate = learning_rate, - lr_scheduler_type = lr_scheduler_type, - lr_scheduler_kwargs = lr_scheduler_kwargs, - warmup_steps = warmup_steps, - optim = optim, - optim_args = optim_args, - weight_decay = weight_decay, - adam_beta1 = adam_beta1, - adam_beta2 = adam_beta2, - adam_epsilon = adam_epsilon, - optim_target_modules = optim_target_modules, - gradient_accumulation_steps = gradient_accumulation_steps, - average_tokens_across_devices = average_tokens_across_devices, - max_grad_norm = max_grad_norm, - label_smoothing_factor = label_smoothing_factor, - bf16 = bf16, - fp16 = fp16, - bf16_full_eval = bf16_full_eval, - fp16_full_eval = fp16_full_eval, - tf32 = tf32, - gradient_checkpointing = gradient_checkpointing, - gradient_checkpointing_kwargs = gradient_checkpointing_kwargs, - torch_compile = torch_compile, - torch_compile_backend = torch_compile_backend, - torch_compile_mode = torch_compile_mode, - use_liger_kernel = use_liger_kernel, - liger_kernel_config = liger_kernel_config, - use_cache = use_cache, - neftune_noise_alpha = neftune_noise_alpha, - torch_empty_cache_steps = torch_empty_cache_steps, - auto_find_batch_size = auto_find_batch_size, - logging_strategy = logging_strategy, - logging_steps = logging_steps, - logging_first_step = logging_first_step, - log_on_each_node = log_on_each_node, - logging_nan_inf_filter = logging_nan_inf_filter, - include_num_input_tokens_seen = include_num_input_tokens_seen, - log_level = log_level, - log_level_replica = log_level_replica, - disable_tqdm = disable_tqdm, - report_to = report_to, - run_name = run_name, - project = project, - trackio_space_id = trackio_space_id, - eval_strategy = eval_strategy, - eval_steps = eval_steps, - eval_delay = eval_delay, - per_device_eval_batch_size = per_device_eval_batch_size, - prediction_loss_only = prediction_loss_only, - eval_on_start = eval_on_start, - eval_do_concat_batches = eval_do_concat_batches, - eval_use_gather_object = eval_use_gather_object, - eval_accumulation_steps = eval_accumulation_steps, - batch_eval_metrics = batch_eval_metrics, - save_only_model = save_only_model, - save_strategy = save_strategy, - save_steps = save_steps, - save_on_each_node = save_on_each_node, - save_total_limit = save_total_limit, - enable_jit_checkpoint = enable_jit_checkpoint, - push_to_hub = push_to_hub, - hub_token = hub_token, - hub_private_repo = hub_private_repo, - hub_model_id = hub_model_id, - hub_strategy = hub_strategy, - hub_always_push = hub_always_push, - hub_revision = hub_revision, - load_best_model_at_end = load_best_model_at_end, - metric_for_best_model = metric_for_best_model, - greater_is_better = greater_is_better, - ignore_data_skip = ignore_data_skip, - restore_callback_states_from_checkpoint = restore_callback_states_from_checkpoint, - full_determinism = full_determinism, - seed = seed, - data_seed = data_seed, - use_cpu = use_cpu, - accelerator_config = accelerator_config, - parallelism_config = parallelism_config, - dataloader_drop_last = dataloader_drop_last, - dataloader_num_workers = dataloader_num_workers, - dataloader_pin_memory = dataloader_pin_memory, - dataloader_persistent_workers = dataloader_persistent_workers, - dataloader_prefetch_factor = dataloader_prefetch_factor, - remove_unused_columns = remove_unused_columns, - label_names = label_names, - train_sampling_strategy = train_sampling_strategy, - length_column_name = length_column_name, - ddp_find_unused_parameters = ddp_find_unused_parameters, - ddp_bucket_cap_mb = ddp_bucket_cap_mb, - ddp_broadcast_buffers = ddp_broadcast_buffers, - ddp_backend = ddp_backend, - ddp_timeout = ddp_timeout, - fsdp = fsdp, - fsdp_config = fsdp_config, - deepspeed = deepspeed, - debug = debug, - skip_memory_metrics = skip_memory_metrics, - do_train = do_train, - do_eval = do_eval, - do_predict = do_predict, - resume_from_checkpoint = resume_from_checkpoint, - warmup_ratio = warmup_ratio, - logging_dir = logging_dir, - local_rank = local_rank, - model_init_kwargs = model_init_kwargs, - trust_remote_code = trust_remote_code, - teacher_model_name_or_path = teacher_model_name_or_path, - teacher_model_revision = teacher_model_revision, - teacher_model_init_kwargs = teacher_model_init_kwargs, - disable_dropout = disable_dropout, - max_completion_length = max_completion_length, - ds3_gather_for_generation = ds3_gather_for_generation, - shuffle_dataset = shuffle_dataset, - pad_to_multiple_of = pad_to_multiple_of, - temperature = temperature, - top_p = top_p, - top_k = top_k, - min_p = min_p, - generation_kwargs = generation_kwargs, - chat_template_kwargs = chat_template_kwargs, - repetition_penalty = repetition_penalty, - cache_implementation = cache_implementation, - use_vllm = use_vllm, - vllm_mode = vllm_mode, - vllm_model_impl = vllm_model_impl, - vllm_enable_sleep_mode = vllm_enable_sleep_mode, - vllm_structured_outputs_regex = vllm_structured_outputs_regex, - vllm_server_base_url = vllm_server_base_url, - vllm_server_host = vllm_server_host, - vllm_server_port = vllm_server_port, - vllm_server_timeout = vllm_server_timeout, - vllm_group_port = vllm_group_port, - vllm_gpu_memory_utilization = vllm_gpu_memory_utilization, - vllm_max_model_length = vllm_max_model_length, - vllm_tensor_parallel_size = vllm_tensor_parallel_size, - beta = beta, - max_tool_calling_iterations = max_tool_calling_iterations, - log_completions = log_completions, - num_completions_to_print = num_completions_to_print, - log_unique_prompts = log_unique_prompts, - teacher_model_server_url = teacher_model_server_url, - loss_top_k = loss_top_k, - loss_add_tail = loss_add_tail, - reverse_kl_top_1_mode = reverse_kl_top_1_mode,**kwargs) - super().__init__(**_unsloth_filter_config_init_kwargs(ServerDistillationConfig, _unsloth_config_arguments, mirrored_from = __class__)) - self.vllm_sampling_params = vllm_sampling_params - self.unsloth_num_chunks = unsloth_num_chunks - if unsloth_grpo_mini_batch is not None: - if self.generation_batch_size >= unsloth_grpo_mini_batch: - self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch - else: - raise ValueError( - f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, " - f"which is self.per_device_train_batch_size * gradient_accumulation_steps." - ) - self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier - - # Unsloth: keep the reentrant checkpoint path - if getattr(self, 'gradient_checkpointing', False): - _gc_kwargs = getattr(self, 'gradient_checkpointing_kwargs', None) or {} - if _gc_kwargs.get('context_fn') is None and not _gc_kwargs.get('debug', False): - _gc_kwargs['use_reentrant'] = True - self.gradient_checkpointing_kwargs = _gc_kwargs - from unsloth.optimizers.flash_adamw import is_flash_adamw_marked, mark_flash_adamw, unmark_flash_adamw - if _unsloth_flash_adamw or is_flash_adamw_marked(self): - mark_flash_adamw(self) - else: - unmark_flash_adamw(self) - -pass - -class _UnslothServerDistillationTrainer(DistillationTrainer): - """""" - - _tag_names = ["trl", "server-distillation"] - _name = "Server Distillation" - - def __init__( - self, - model, - args: ServerDistillationConfig | None = None, - train_dataset=None, - eval_dataset=None, - processing_class=None, - callbacks=None, - optimizers=(None, None), - quantization_config=None, - peft_config=None, - ): - if args is None: - args = ServerDistillationConfig(output_dir="tmp_server_distillation") - - # No local teacher: the base sets `self.teacher_model = None`, then we attach the server client below. - super().__init__( - model=model, - teacher_model=None, - args=args, - train_dataset=train_dataset, - eval_dataset=eval_dataset, - processing_class=processing_class, - callbacks=callbacks, - optimizers=optimizers, - quantization_config=quantization_config, - peft_config=peft_config, - ) - - self.reverse_kl_top_1_mode = args.reverse_kl_top_1_mode - self.loss_top_k = args.loss_top_k - self.loss_add_tail = args.loss_add_tail - - from trl.generation.vllm_client import VLLMClient - - self.teacher_client = VLLMClient(base_url=args.teacher_model_server_url, connection_timeout=60.0) - - def _get_reverse_kl_top_1_tokens( - self, student_scores: torch.Tensor, completion_tokens: torch.Tensor - ) -> torch.Tensor: - """Return the reverse-KL top-1 token IDs for the mixed top-1 loss path. - - Args: - student_scores: Any (B, T, V) tensor whose argmax selects the student's top token - (logits or log-probs — both are order-preserving). - completion_tokens: (B, T) actual token IDs in the completion. - """ - if self.reverse_kl_top_1_mode == "argmax": - return student_scores.argmax(dim=-1) - return completion_tokens - - def _compute_sparse_top_1_divergence_loss( - self, - student_log_probs: torch.Tensor, - teacher_top1_token_ids: torch.Tensor, - teacher_top1_logprobs: torch.Tensor, - reverse_token_ids: torch.Tensor, - reverse_teacher_logprobs: torch.Tensor, - labels: torch.Tensor, - num_items_in_batch=None, - ) -> torch.Tensor: - """Compute exact generalized JSD/KL on top-1 support for the mixed beta>0 path.""" - neg_inf = torch.full((), float("-inf"), dtype=student_log_probs.dtype, device=student_log_probs.device) - - if self.beta == 1: - support = reverse_token_ids.unsqueeze(-1) - support_mask = torch.ones_like(support, dtype=torch.bool) - teacher_support_logprobs = reverse_teacher_logprobs.unsqueeze(-1) - else: - teacher_support = teacher_top1_token_ids.unsqueeze(-1) - reverse_support = reverse_token_ids.unsqueeze(-1) - support = torch.cat([teacher_support, reverse_support], dim=-1) - support_mask = torch.ones_like(support, dtype=torch.bool) - support_mask[..., 1] = support[..., 1] != support[..., 0] - teacher_support_logprobs = torch.stack([teacher_top1_logprobs, reverse_teacher_logprobs], dim=-1) - support = torch.where(support_mask, support, torch.zeros_like(support)) - - student_support_logprobs = student_log_probs.gather(-1, support) - student_support_logprobs = torch.where(support_mask, student_support_logprobs, neg_inf) - teacher_support_logprobs = torch.where(support_mask, teacher_support_logprobs, neg_inf) - - if self.loss_add_tail: - base_support_mask = support_mask - student_sparse_log_probs, support_mask = _add_tail_bucket(student_support_logprobs, base_support_mask) - teacher_sparse_log_probs, _ = _add_tail_bucket(teacher_support_logprobs, base_support_mask) - else: - student_sparse_log_probs = student_support_logprobs - torch.logsumexp( - student_support_logprobs, dim=-1, keepdim=True - ) - teacher_sparse_log_probs = teacher_support_logprobs - torch.logsumexp( - teacher_support_logprobs, dim=-1, keepdim=True - ) - - jsd = _jsd_divergence(student_sparse_log_probs, teacher_sparse_log_probs, self.beta, support_mask) - return self._reduce_divergence_loss( - jsd, completion_mask=labels != -100, reduction="batchmean", num_items_in_batch=num_items_in_batch - ) - - def _compute_prompt_length(self, inputs: dict[str, torch.Tensor | Any]) -> int: - """Compute the earliest prompt boundary that still includes every completion token in the batch.""" - if inputs.get("labels") is not None: - attention_mask = inputs["attention_mask"] - labels = inputs["labels"] - full_lengths = attention_mask.sum(dim=1) - completion_lengths = (labels != -100).sum(dim=1) - return int((full_lengths - completion_lengths).min().item()) - return inputs["prompts"].shape[1] - - @staticmethod - def _reduce_divergence_loss(jsd, completion_mask=None, reduction="batchmean", num_items_in_batch=None): - """Reduce a per-token divergence tensor over the valid completion tokens. - - When `num_items_in_batch` is provided (as under gradient accumulation), the divergence is reduced as `sum / - num_items_in_batch`, matching the gradient-accumulation-correct behavior of HF's default cross-entropy. - Otherwise it falls back to the local `reduction` (default `batchmean`). See issue #4719. - """ - mask = None - if completion_mask is not None: - mask = completion_mask.bool() - jsd = jsd[mask] - - if num_items_in_batch is not None: - # Normalize by the global number of valid tokens for gradient-accumulation-correct loss. - jsd_sum = jsd.sum() - if isinstance(num_items_in_batch, torch.Tensor): - num_items_in_batch = num_items_in_batch.to(jsd_sum.device) - return jsd_sum / num_items_in_batch - if reduction == "batchmean": - # clamp_min(1) avoids 0/0 -> nan when a sample has no unmasked positions - # (e.g. completion fully truncated). jsd[mask] is empty -> jsd.sum() == 0, - # so 0/1 == 0 with a valid grad path. - denom = mask.sum().clamp_min(1) if completion_mask is not None else max(jsd.size(0), 1) - return jsd.sum() / denom - elif reduction == "sum": - return jsd.sum() - elif reduction == "mean": - return jsd.mean() - else: - return jsd - - @profiling_decorator - def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): - # Reconstruct the old `input_ids` / `attention_mask` / `labels` layout this server path is built on from GRPO's - # keys, so it rides on the base's (evolving) generation. The teacher-server loss is migrated onto the stable - # base wholesale later (issue #6449, item 58). The first four lines mirror `DistillationTrainer.compute_loss`. - prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"] - completion_ids, completion_mask = inputs["completion_ids"], inputs["completion_mask"] - input_ids = torch.cat([prompt_ids, completion_ids], dim=1) - attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) - # Server-only: the teacher-server loss masks on `labels`/`prompt_attention_mask`, so rebuild them too. - completion_labels = torch.where(completion_mask.bool(), completion_ids, torch.full_like(completion_ids, -100)) - labels = torch.cat([torch.full_like(prompt_ids, -100), completion_labels], dim=1) - inputs = { - **inputs, - "input_ids": input_ids, - "attention_mask": attention_mask, - "labels": labels, - "prompt_attention_mask": prompt_mask, - } - - # Student forward pass - student_outputs = model( - input_ids=inputs["input_ids"], - attention_mask=inputs["attention_mask"], - ) - prompt_length = self._compute_prompt_length(inputs) - labels = inputs["labels"][:, prompt_length:] - completion_tokens = inputs["input_ids"][:, prompt_length:] - - # Server path: token-level divergence using teacher logprobs. - # The server returns: - # actual_logprobs – (B, T) teacher log p(x_actual) (for reverse KL) - # topk_logprobs – (B, T, K) teacher top-k sorted logprobs (for forward KL) - # topk_token_ids – (B, T, K) corresponding token IDs - teacher_result = self._get_teacher_token_logprobs_from_server(inputs, prompt_length) - - student_logits = student_outputs.logits[:, prompt_length - 1 : -1, :] - student_log_probs = F.log_softmax(student_logits / self.temperature, dim=-1) - - comp_len = teacher_result["actual_logprobs"].shape[1] - completion_tokens = completion_tokens[:, :comp_len] - trimmed_labels = labels[:, :comp_len] - - if self.beta > 0: - loss = self._compute_server_sparse_top_1_divergence_loss( - teacher_result=teacher_result, - student_log_probs=student_log_probs[:, :comp_len, :], - completion_tokens=completion_tokens, - labels=trimmed_labels, - ) - else: - loss = self._compute_server_forward_kl_loss( - teacher_result=teacher_result, - student_log_probs=student_log_probs[:, :comp_len, :], - labels=trimmed_labels, - ) - - # The base trainer sets `loss_is_scaled_for_ga = True` because it normalizes by the global completion-token - # count. The server path normalizes locally with `batchmean` and does not consume `num_items_in_batch`, so it - # must re-apply that scaling itself. - if self.model.training: - loss = loss / self.current_gradient_accumulation_steps - - return (loss, student_outputs) if return_outputs else loss - - def _get_teacher_token_logprobs_from_server( - self, - inputs: dict[str, torch.Tensor | Any], - aligned_prompt_length: int, - ) -> dict[str, torch.Tensor]: - """Fetch per-token teacher logprobs from an external vLLM server. - - Returns a dict with: - ``actual_logprobs`` - (batch, completion_length) teacher log-prob for the actual - token at each position (for reverse KL). - ``topk_logprobs`` - (batch, completion_length, K) teacher top-k sorted logprobs - (for forward KL). - ``topk_token_ids`` - (batch, completion_length, K) corresponding token IDs. - """ - import numpy as np - - input_ids = inputs["input_ids"] - batch_size = input_ids.shape[0] - sequences, prompt_lengths, completion_lengths = build_teacher_request_inputs( - input_ids, - inputs["attention_mask"], - prompt_attention_mask=inputs.get("prompt_attention_mask"), - labels=inputs.get("labels"), - ) - - # The pure forward server path can use the requested teacher top-k support. - # When beta > 0, config validation restricts the server-backed path to top-1. - requested_top_k = self.loss_top_k - result = self.teacher_client.get_sequence_logprobs( - sequences=sequences, - prompt_lengths=prompt_lengths, - top_logprobs=requested_top_k, - temperature=self.temperature, - ) - K = requested_top_k - - device = input_ids.device - labels = inputs.get("labels") - if labels is None: - raise ValueError("labels are required to align teacher-server logprobs with the student loss tensors.") - - # The student loss slices tensors in padded-sequence coordinates starting at `aligned_prompt_length`. - # Place each teacher completion into that same coordinate system by locating the first non-masked completion - # token in `labels`. This works for both left-padded off-policy batches and on-policy batches where - # completions are right-padded after a fixed-width prompt block. - completion_offsets = [] - label_mask = labels != -100 - for sample_mask, comp_len in zip(label_mask, completion_lengths, strict=True): - if comp_len == 0: - completion_offsets.append(0) - continue - completion_start = int(torch.nonzero(sample_mask, as_tuple=False)[0].item()) - completion_offsets.append(completion_start - aligned_prompt_length) - - # Size the output tensors to tightly fit the teacher logprobs. Using the full padded - # sequence length would include padding positions with -inf teacher logprobs, producing - # +inf in the forward pass and NaN gradients in the backward pass (0 * inf = NaN). - # Shorter samples in variable-length batches still need the -inf sentinel at the tail; - # downstream loss consumers (_compute_server_sparse_top_1_divergence_loss, - # _compute_server_forward_kl_loss) neutralise those positions before the divergence - # math runs. - completion_length = max( - (offset + len(lps) for offset, lps in zip(completion_offsets, result["logprobs"], strict=True)), - default=0, - ) - - # actual_logprobs: (B, T) - teacher logprob for the actual token - def _actual_to_tensor(key): - arr = np.full((batch_size, completion_length), float("-inf"), dtype=np.float32) - for i, (offset, seq_lps) in enumerate(zip(completion_offsets, result[key], strict=True)): - if seq_lps: - vals = np.array(seq_lps, dtype=np.float32) # (comp_len_i, 1) - arr[i, offset : offset + vals.shape[0]] = vals[:, 0] - return torch.from_numpy(arr).to(device) - - # topk: (B, T, K) - def _topk_to_tensor(key, k, np_dtype, fill): - arr = np.full((batch_size, completion_length, k), fill, dtype=np_dtype) - for i, (offset, seq_vals) in enumerate(zip(completion_offsets, result[key], strict=True)): - if seq_vals: - vals = np.array(seq_vals, dtype=np_dtype) # (comp_len_i, k) - arr[i, offset : offset + vals.shape[0], :] = vals - return torch.from_numpy(arr).to(device) - - return { - "actual_logprobs": _actual_to_tensor("actual_logprobs"), - "topk_logprobs": _topk_to_tensor("logprobs", K, np.float32, float("-inf")), - "topk_token_ids": _topk_to_tensor("logprob_token_ids", K, np.int64, 0), - } - - def _compute_server_sparse_top_1_divergence_loss( - self, - teacher_result: dict[str, torch.Tensor], - student_log_probs: torch.Tensor, - completion_tokens: torch.Tensor, - labels: torch.Tensor, - ) -> torch.Tensor: - """Compute exact sparse top-1 generalized JSD/KL from server-provided teacher logprobs. - - Args: - teacher_result: dict with ``actual_logprobs`` (B, T), ``topk_logprobs`` (B, T, K), - ``topk_token_ids`` (B, T, K). - student_log_probs: (B, T, V) student log-softmax over vocabulary. - completion_tokens: (B, T) actual token IDs in the completion. - labels: (B, T) with -100 for positions to ignore. - """ - topk_teacher_lps = teacher_result["topk_logprobs"] # (B, T, 1) - topk_token_ids = teacher_result["topk_token_ids"] # (B, T, 1) - actual_teacher_lps = teacher_result["actual_logprobs"] # (B, T) - required = labels != -100 - - missing_actual = required & ~torch.isfinite(actual_teacher_lps) - if missing_actual.any(): - missing_count = int(missing_actual.sum().item()) - total_required = int(required.sum().item()) - raise ValueError( - "Teacher server is missing actual-token logprobs for required reverse-KL positions: " - f"{missing_count}/{total_required}." - ) - if self.beta < 1: - teacher_top1_logprobs = topk_teacher_lps.squeeze(-1) - missing_top1 = required & ~torch.isfinite(teacher_top1_logprobs) - if missing_top1.any(): - missing_count = int(missing_top1.sum().item()) - total_required = int(required.sum().item()) - raise ValueError( - "Teacher server is missing top-1 logprobs for required forward-KL positions: " - f"{missing_count}/{total_required}." - ) - - # Replace -inf teacher logprobs at intra-batch padding (labels == -100) with 0 so - # reverse-KL's student_probs·(log_s - log_t) does not leak +inf into the backward pass. - pad_mask_2d = ~required - pad_mask_3d = pad_mask_2d.unsqueeze(-1) - topk_teacher_lps = torch.where(pad_mask_3d, 0.0, topk_teacher_lps) - actual_teacher_lps = torch.where(pad_mask_2d, 0.0, actual_teacher_lps) - - # Server path only supports "sampled" mode — config validation enforces this, but we guard - # explicitly so future relaxations of the config check don't silently change behaviour. - reverse_token_ids = self._get_reverse_kl_top_1_tokens(student_log_probs, completion_tokens) - # The server path normalizes locally (batchmean), not by num_items_in_batch: teacher logprobs may not cover - # every student completion token (the loss is summed over the trimmed teacher window), so the global token - # count would be the wrong denominator. Gradient-accumulation normalization for the server path is left as a - # follow-up. - return self._compute_sparse_top_1_divergence_loss( - student_log_probs=student_log_probs, - teacher_top1_token_ids=topk_token_ids.squeeze(-1), - teacher_top1_logprobs=topk_teacher_lps.squeeze(-1), - reverse_token_ids=reverse_token_ids, - reverse_teacher_logprobs=actual_teacher_lps, - labels=labels, - ) - - def _compute_server_forward_kl_loss( - self, - teacher_result: dict[str, torch.Tensor], - student_log_probs: torch.Tensor, - labels: torch.Tensor, - ) -> torch.Tensor: - """Compute sparse forward KL from server-provided teacher top-k logprobs (beta==0 path). - - Args: - teacher_result: dict with ``topk_logprobs`` (B, T, K) and ``topk_token_ids`` (B, T, K). - student_log_probs: (B, T, V) student log-softmax over vocabulary. - labels: (B, T) with -100 for positions to ignore. - """ - teacher_topk_logprobs = teacher_result["topk_logprobs"] - teacher_topk_token_ids = teacher_result["topk_token_ids"] - valid = teacher_topk_logprobs > float("-inf") - neg_inf = torch.full((), float("-inf"), dtype=student_log_probs.dtype, device=student_log_probs.device) - student_topk_logprobs = student_log_probs.gather(dim=-1, index=teacher_topk_token_ids) - student_topk_logprobs = torch.where(valid, student_topk_logprobs, neg_inf) - teacher_topk_logprobs = torch.where(valid, teacher_topk_logprobs, neg_inf) - - if self.loss_add_tail: - base_support_mask = valid - student_sparse_log_probs, support_mask = _add_tail_bucket(student_topk_logprobs, base_support_mask) - teacher_sparse_log_probs, _ = _add_tail_bucket(teacher_topk_logprobs, base_support_mask) - else: - support_mask = valid - student_sparse_log_probs = student_topk_logprobs - torch.logsumexp( - student_topk_logprobs, dim=-1, keepdim=True - ) - teacher_sparse_log_probs = teacher_topk_logprobs - torch.logsumexp( - teacher_topk_logprobs, dim=-1, keepdim=True - ) - - jsd = _jsd_divergence( - student_sparse_log_probs, - teacher_sparse_log_probs, - beta=0.0, - support_mask=support_mask, - ) - # See `_compute_server_sparse_top_1_divergence_loss`: the server path normalizes locally, not by - # num_items_in_batch, because the teacher window may not cover every student completion token. - return self._reduce_divergence_loss(jsd, completion_mask=labels != -100, reduction="batchmean") -class UnslothServerDistillationTrainer(_UnslothServerDistillationTrainer): - """ - Distillation from a teacher hosted on an external vLLM server. - - Instead of running a local teacher forward pass, per-token teacher logprobs are fetched from a vLLM server via - [`~generation.vllm_client.VLLMClient`]. The server only returns the teacher's top-k logprobs, so the divergence is - restricted to a sparse support (top-1 for `beta > 0`, top-k for the pure forward path `beta = 0`). Everything else - — the student forward, generation, buffering, metrics — is inherited from [`DistillationTrainer`]. - - """ - def __init__( - self, - model, - args = None, - train_dataset = None, - eval_dataset = None, - processing_class = None, - callbacks = None, - quantization_config = None, - peft_config = None, - **kwargs - ): - if args is None: args = UnslothServerDistillationConfig() - _unsloth_precision_requested = tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) - use_bf16 = getattr(args, 'bf16', False) - if type(use_bf16) is not bool: use_bf16 = False - use_fp16 = getattr(args, 'fp16', False) - if type(use_fp16) is not bool: use_fp16 = False - force_float32 = False - try: - from unsloth_zoo.device_type import device_is_bf16_supported as _bf16_supported - except Exception: - _bf16_supported = torch.cuda.is_bf16_supported - try: - from unsloth_zoo.device_type import DEVICE_TYPE as _unsloth_device_type - except Exception: - _unsloth_device_type = None - full_finetuning = getattr(model, '_unsloth_full_finetuning', None) - if full_finetuning is None: full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1' - model_forced_float32 = getattr(model, '_unsloth_forced_float32', None) - if model_forced_float32 is None: model_forced_float32 = os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1' - if model_forced_float32 and not (full_finetuning and _bf16_supported()): - print('Unsloth: Switching to float32 training since model cannot work with float16') - force_float32 = True - mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') - dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None) - if dtype is None: dtype = model.get_input_embeddings().weight.dtype - from unsloth_zoo.utils import _get_dtype - dtype = _get_dtype(dtype) - float16 = dtype == torch.float16 - bfloat16 = dtype == torch.bfloat16 - float32 = dtype == torch.float32 - user_float32 = bool(getattr(model, '_unsloth_user_float32', False)) - if full_finetuning: - if bfloat16 and use_fp16: use_fp16 = False - if float16 and use_bf16: use_bf16 = False - if not force_float32 and bfloat16 and use_fp16 and _unsloth_device_type == 'xpu' and not _bf16_supported(): - print('Unsloth: The model was loaded in bfloat16 on purpose, so training uses bfloat16 mixed precision instead of float16.') - use_fp16 = False - use_bf16 = True - args.fp16 = False - args.bf16 = True - if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`') - if not force_float32 and (bfloat16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`') - if force_float32: - # Forced float32 training - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32' and float32 and user_float32 and not _bf16_supported(): - print('Unsloth: Model is in float32 and this GPU has no bfloat16 support, so training stays in float32. Pass fp16 = True to force float16 mixed precision instead.') - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32': - # Mixed precision training. bf16 only if the GPU supports it; V100/T4 use fp16. - use_bf16_amp = (not float16) and (_bf16_supported() or (bfloat16 and _unsloth_device_type == 'xpu')) - args.fp16 = not use_bf16_amp - args.bf16 = use_bf16_amp - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16_amp else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16_amp else 'fp16' - # args.mixed_precision is a new argument which needs to be set now - elif mixed_precision_dtype == 'bfloat16': - # Both False since bfloat16 full finetuning doesn't do any autocasting. - args.fp16 = False - args.bf16 = False - os.environ['ACCELERATE_MIXED_PRECISION'] = 'no' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no' - # args.mixed_precision is a new argument which needs to be set now - elif use_bf16 or use_fp16: - # transformers <5 exported this itself from fp16/bf16; 5.x dropped the write, so an - # explicit flag left it unset and GRPO readers defaulted to 'fp16', wrapping a - # bfloat16 model in a float16 autocast. See unslothai/unsloth#4891. - os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16 else 'fp16' - if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16 else 'fp16' - - if getattr(args, 'eval_dataset', None) is not None and getattr(args, 'eval_strategy', 'no') == 'no': - args.eval_strategy = 'steps' - if getattr(args, 'eval_steps', None) is None: args.eval_steps = 0.1 - ga_steps = getattr(args, 'gradient_accumulation_steps', None) - if ga_steps is not None and ga_steps > 1: - from transformers import __version__ as transformers_version - if Version(transformers_version) <= Version('4.45.2'): - print('**** Unsloth: Please use our fixed gradient_accumulation_steps by updating transformers, TRL and Unsloth!\n' - '`pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth transformers trl unsloth_zoo`') - if getattr(args, 'eval_strategy', 'no') != 'no': - eval_bsz = getattr(args, 'per_device_eval_batch_size', 8) - eval_generations = getattr(args, 'num_generations_eval', None) or getattr(args, 'num_generations', None) or 1 - if eval_bsz == 8 and args.per_device_train_batch_size < eval_bsz and (args.per_device_train_batch_size * args.world_size) % eval_generations == 0: args.per_device_eval_batch_size = args.per_device_train_batch_size - if getattr(args, 'eval_accumulation_steps', None) is None and ga_steps is not None: args.eval_accumulation_steps = ga_steps - fp16_full_eval = getattr(args, 'fp16_full_eval', False) - if type(fp16_full_eval) is not bool: fp16_full_eval = False - bf16_full_eval = getattr(args, 'bf16_full_eval', False) - if type(bf16_full_eval) is not bool: bf16_full_eval = False - if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True - if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False - if force_float32: - args.bf16_full_eval = False - args.fp16_full_eval = False - elif os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16': - args.bf16_full_eval = True - args.fp16_full_eval = False - elif not bf16_full_eval and not fp16_full_eval: - args.bf16_full_eval = args.bf16 - args.fp16_full_eval = args.fp16 - if tuple(getattr(args, k, False) for k in ('fp16', 'bf16', 'fp16_full_eval', 'bf16_full_eval')) != _unsloth_precision_requested: - from unsloth.models.loader_utils import sync_deepspeed_mixed_precision - sync_deepspeed_mixed_precision(args) - if model is not None: - from unsloth.models.loader_utils import apply_deepspeed_zero3_model_setup - apply_deepspeed_zero3_model_setup(model) - if model is not None: - _warnings_issued = getattr(model, 'warnings_issued', None) - if _warnings_issued is None: - model.warnings_issued = {} - elif not isinstance(_warnings_issued, dict): - try: - model.warnings_issued = dict(_warnings_issued) - except Exception: - model.warnings_issued = {} - if 'max_seq_length' not in locals() and not hasattr(args, 'max_seq_length'): - pass - else: - model_max_seq_length = getattr(model, 'max_seq_length', None) - args_max_seq_length = getattr(args, 'max_seq_length', None) - if args_max_seq_length is None and model_max_seq_length is not None: - max_seq_length = model.max_seq_length - if hasattr(args, 'max_seq_length'): args.max_seq_length = max_seq_length - elif args_max_seq_length is not None and model_max_seq_length is not None: - if args_max_seq_length > model_max_seq_length: - print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but ' - 'the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.') - args.max_seq_length = model_max_seq_length - if model is not None and hasattr(model, 'for_training'): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right' - if 'processing_class' in locals(): - if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right' - if hasattr(processing_class, 'tokenizer') and hasattr(processing_class.tokenizer, 'padding_side'): processing_class.tokenizer.padding_side = 'right' - other_metrics = [] - - from unsloth_zoo.logging_utils import PatchRLStatistics - PatchRLStatistics('server_distillation_trainer', other_metrics) - - # [TODO] Fix up DataParallel multiplying batch sizes - # [TODO] DDP works, but DP seems to not work? [TODO] - if getattr(args, "parallel_mode", None) == ParallelMode.NOT_DISTRIBUTED and args.n_gpu > 1: - if getattr(args, "_n_gpu", 1) != 1: - args._n_gpu = 1 - if "model" in locals(): - try: - from unsloth.models._gc_config_default import honour_gradient_checkpointing_off - honour_gradient_checkpointing_off(args, model) - except Exception: - pass - if "model" in locals() and hasattr(model, "for_training"): - _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) - model.for_training(use_gradient_checkpointing=_use_gc) - if "model" in locals(): - try: - from unsloth.models._ddp_defaults import unsloth_default_ddp_find_unused_parameters - unsloth_default_ddp_find_unused_parameters(args, model) - except Exception: - pass - super().__init__( - model = model, - args = args, - train_dataset = train_dataset, - eval_dataset = eval_dataset, - processing_class = processing_class, - callbacks = callbacks, - quantization_config = quantization_config, - peft_config = peft_config,**kwargs) - # A tensor parallel model (tp_native.tp_trainer_setup) stays on its training path: for_inference would - # route the decoder layers onto decode paths outside the tensor parallel regions. - if "model" in locals() and hasattr(model, "for_inference") and not getattr(self, "_unsloth_tensor_parallel", False): - model.for_inference() - if hasattr(self, 'neftune_hook_handle'): - self.neftune_hook_handle.remove() - if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle - if getattr(args, 'neftune_noise_alpha', None) is not None: - model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha - pass - if hasattr(self, 'accelerator'): - scaler = self.accelerator.scaler - current_model = model - while hasattr(current_model, 'model'): - current_model.accelerator_scaler = scaler - current_model = current_model.model - current_model.accelerator_scaler = scaler - pass - if hasattr(self, 'train'): - self.train = MethodType(prepare_for_training_mode(self.__class__.train), self) - pass - if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'): - _vllm_tok = self.llm.get_tokenizer() - _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None) - if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None: - _vllm_tok.chat_template = _pc.chat_template - pass - if getattr(self, 'aux_loss_enabled', False) and hasattr(getattr(self, 'model', None), 'config'): - _text_config = self.model.config - if hasattr(_text_config, 'get_text_config'): _text_config = _text_config.get_text_config() - _n_experts = [getattr(_text_config, _k) for _k in ('num_local_experts', 'num_experts', 'n_routed_experts', 'moe_num_experts') if isinstance(getattr(_text_config, _k, None), int)] - if _n_experts and all(_n <= 1 for _n in _n_experts): - self.aux_loss_enabled = False - _text_config.output_router_logits = False - pass - -pass diff --git a/unsloth_compiled_cache/moe_utils.py b/unsloth_compiled_cache/moe_utils.py deleted file mode 100644 index afd40afa7d0c41f862ce8ac5315cf08ce874af06..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/moe_utils.py +++ /dev/null @@ -1,6221 +0,0 @@ -# Unsloth Zoo - Utilities for Unsloth -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . -import torch -import torch.nn.functional as F -import contextlib -import inspect -import json -import math -import os -import shutil -import stat -import tempfile -import sys -import threading -import time -import warnings -import importlib -import importlib.util -from typing import Optional, Tuple -from torch.autograd import Function -from unsloth_zoo.mlx import is_mlx_available -from unsloth_zoo.mxfp4_dequant import is_mxfp4_expert_param - -UNSLOTH_COMPILE_LOCATION = os.environ.get( - "UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache" -) - -try: - import bitsandbytes as bnb - from bitsandbytes.nn import Params4bit - # unsloth_zoo installs a permissive bitsandbytes stub wherever the real package is - # absent, macOS arm64 among others, and every attribute of that stub resolves to a - # placeholder object rather than a class. `isinstance(x, Params4bit)` against one - # raises TypeError: isinstance() arg 2 must be a type, so a non-class Params4bit has - # to count as no bitsandbytes at all. On a real install this is just True. - HAS_BNB = isinstance(Params4bit, type) - if not HAS_BNB: - Params4bit = None -except Exception: - # Not just ImportError: a bitsandbytes mismatched with torch fails its own import with AttributeError. - HAS_BNB = False - Params4bit = None - -if not isinstance(Params4bit, type): - # A bitsandbytes that imports but does not expose Params4bit as a class, which is what - # the macOS build does, makes every `isinstance(param, Params4bit)` below raise - # TypeError instead of answering False. Probe the object, not the platform or the - # version, and treat that install as no bitsandbytes at all: there is no 4-bit expert - # path without the class. - HAS_BNB = False - Params4bit = None - - -def _get_compile_location() -> str: - return os.path.abspath( - os.environ.get("UNSLOTH_COMPILE_LOCATION", UNSLOTH_COMPILE_LOCATION) - ) - - -def _log_info(message: str): - if os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1": - print(message) - - -def _warn_without_raising(message): - """warnings.warn, except that it cannot be the thing that fails the import. - - install_to_cache() runs at import and simplefilter("error") would turn this - into a raise, for a condition already handled by not using the file. - """ - try: - warnings.warn(message) - except Exception: - print(message) - - -def _read_file_bytes(path): - """The file's bytes, or None when it cannot be read.""" - try: - with open(path, "rb") as handle: - return handle.read() - except OSError: - return None - - -def _replace_with_copy(current_file, destination): - """Copy over `destination` by replacing the path, not by writing into it. - - shutil.copy() opens the destination, so a read-only one refuses it, and it - is not atomic while nothing locks this file and every rank installs it at - import. os.replace() needs only the directory and is atomic. - """ - directory = os.path.dirname(destination) or "." - descriptor, temporary = tempfile.mkstemp( - prefix = f".{os.path.basename(destination)}.", suffix = ".tmp", dir = directory, - ) - try: - # Through the DESCRIPTOR, not by reopening the name, or the copy and the - # mode follow that name somewhere else; this directory is writable by - # whoever would plant a copy here. Same shape as - # compiler._replace_compiled_cache_file. - with os.fdopen(descriptor, "wb") as handle: - descriptor = None - with open(current_file, "rb") as source: - shutil.copyfileobj(source, handle) - handle.flush() - # mkstemp is owner-only; keep the mode shutil.copy() would have left. - mode = stat.S_IMODE(os.stat(current_file).st_mode) - if hasattr(os, "fchmod"): - os.fchmod(handle.fileno(), mode) - else: - os.chmod(temporary, mode) - os.replace(temporary, destination) - except BaseException: - if descriptor is not None: - try: - os.close(descriptor) - except OSError: - pass - try: - os.remove(temporary) - except OSError: - pass - raise - - -def install_to_cache(source_path, destination_filename=None): - """Copy a file into unsloth_compiled_cache so compiled modules can use it. - - Returns whether the copy is known to match `source_path`. Runs at import - time, so an unwritable directory is not fatal; a destination that exists and - does NOT match is reported, since _load_cached_moe_utils_module() execs it. - """ - compile_location = _get_compile_location() - if not os.path.exists(compile_location): - try: - os.makedirs(compile_location) - except: - pass - - current_file = os.path.abspath(source_path) - if destination_filename is None: - destination_filename = os.path.basename(current_file) - - destination = os.path.abspath(os.path.join(compile_location, destination_filename)) - - if current_file == destination: - return True - - current = _read_file_bytes(current_file) - if current is not None and _read_file_bytes(destination) == current: - # Already installed. This runs on every `import unsloth_zoo`, on every - # rank, and copies this whole module, so without the compare the steady - # state was replacing a file with an identical copy of itself. - return True - - try: - _replace_with_copy(current_file, destination) - except Exception as replace_error: - # A directory we cannot add a temp file to can still hold a destination - # we can write through, so the plain copy is worth one attempt. It gives - # up the atomicity above, so the readback below is what keeps a partial - # copy from being used. - try: - # Writing through a planted symlink would clobber whatever it targets. - if os.path.islink(destination) or os.path.isdir(destination): - raise OSError("destination is a symlink or directory") - shutil.copy(current_file, destination) - except Exception as copy_error: - _log_info( - f"Unsloth: Could not install {destination}: " - f"{replace_error}; {copy_error}" - ) - - current = _read_file_bytes(current_file) - installed = _read_file_bytes(destination) - if installed is None or current is None: - # Nothing landed, or there is nothing to compare it against. Either way - # no compiled module gets to use it. - return False - if installed == current: - return True - _warn_without_raising( - f"Unsloth: {destination} does not match {current_file} and could not be " - "replaced, so it will not be used. Delete it to restore the compiled " - "cache copy." - ) - return False - - -install_to_cache(__file__, "moe_utils.py") - -_CACHED_FORWARD_MOE_BACKEND = None -_CACHED_MOE_UTILS_MODULE = None - - -_WARNED_STALE_CACHE = set() - - -def _cached_copy_is_current(cache_file, current_file) -> bool: - """Whether the compiled cache holds byte for byte what this module is. - - They diverge when the copy could not be written, leaving an OLDER - unsloth_zoo that every caller below would prefer, silently installing that - release's patches over this one's. Prefer this module, and say so once. - """ - cached = _read_file_bytes(cache_file) - current = _read_file_bytes(current_file) - if cached is None or current is None: - return False - if cached == current: - return True - - if cache_file not in _WARNED_STALE_CACHE: - _WARNED_STALE_CACHE.add(cache_file) - logger = _moe_utils_logger() - if logger is not None: - logger.warning( - f"Unsloth: {cache_file} is from a different version of unsloth_zoo and " - f"could not be refreshed, so it is being ignored. Delete that directory " - f"if MoE behaviour looks stale." - ) - return False - - -def _remove_cached_bytecode(source_file): - """Drop the pyc beside a cache copy we are about to execute. - - The comparison covers the .py only, and CPython runs an unchecked-hash pyc - without consulting it, so a pyc we cannot remove means using our own defs. - """ - try: - bytecode_location = importlib.util.cache_from_source(source_file) - except NotImplementedError: - return True - # A symlinked `__pycache__` sends this unlink outside the cache. Refusing - # costs the cache copy and falls back to our own definitions, which beats - # deleting somebody else's file. sys.pycache_prefix is the user's own. - if not getattr(sys, "pycache_prefix", None): - try: - if os.path.islink(os.path.dirname(bytecode_location)): - return False - except OSError: - return False - try: - os.remove(bytecode_location) - except FileNotFoundError: - pass - except OSError: - return not os.path.isfile(bytecode_location) - return True - - -def cached_copy_is_importable(directory) -> bool: - """Whether a generated module may be allowed to import moe_utils from here. - - A path that exists and is not a real directory is never trusted: zipimport - is a default sys.path hook, so a ZIP named `unsloth_compiled_cache` passed - every isfile check and served moe_utils out of the archive. - - True when there is no copy at all; the caller runs - compiler._reject_shadowing_import_candidates first, which refuses a - moe_utils package, extension or sourceless pyc. Matching bytes are necessary - and NOT sufficient, hence the name: a bare import prefers the pyc, so it is - dropped here and one that cannot be dropped means no. - """ - try: - # Exists but is not a directory, not merely "is not a directory": a path - # that does not exist yet imports nothing and stays trusted, which is - # what a cache folder looks like before it is created. - if os.path.exists(directory) and not os.path.isdir(directory): - return False - cache_file = os.path.abspath(os.path.join(directory, "moe_utils.py")) - except Exception: - return False - current_file = os.path.abspath(__file__) - if cache_file == current_file or not os.path.isfile(cache_file): - return True - if not _cached_copy_is_current(cache_file, current_file): - return False - return _remove_cached_bytecode(cache_file) - - -def _load_cached_moe_utils_module(): - global _CACHED_MOE_UTILS_MODULE - - cache_file = os.path.abspath(os.path.join(_get_compile_location(), "moe_utils.py")) - current_file = os.path.abspath(__file__) - if not os.path.isfile(cache_file) or cache_file == current_file: - return None - if not _cached_copy_is_current(cache_file, current_file): - _CACHED_MOE_UTILS_MODULE = None - return None - - # The cache copy is only ever a copy of this file, so bytes that differ are - # bytes install_to_cache() did not put there, and exec_module() below runs - # whatever is in the file. Use this module's own definitions instead, which - # is what every caller falls back to anyway. - cached_bytes = _read_file_bytes(cache_file) - if cached_bytes is None or cached_bytes != _read_file_bytes(current_file): - return None - - try: - module_name = "unsloth_cached_moe_utils" - module = sys.modules.get(module_name, None) - if module is not None and os.path.abspath(getattr(module, "__file__", "")) == cache_file: - _CACHED_MOE_UTILS_MODULE = module - return module - - if not _remove_cached_bytecode(cache_file): - return None - - spec = importlib.util.spec_from_file_location(module_name, cache_file) - if spec is None: - return None - module = importlib.util.module_from_spec(spec) - sys.modules[module_name] = module - try: - # The bytes compared above, not a fresh read of the path. exec_module - # reopens the file, and the comparison and that open are two moments: - # a writer with access to the shared cache replaces the file between - # them and its bytes run as unsloth_cached_moe_utils having matched - # nothing. compiler.py's loader closed the same window; this one is - # the copy of it that lives here. - exec(compile(cached_bytes, cache_file, "exec"), module.__dict__) - except Exception: - sys.modules.pop(module_name, None) - raise - _CACHED_MOE_UTILS_MODULE = module - return module - except Exception: - return None - - -def _select_moe_backend_eagerly(module): - # Selecting inside a trace sets a guarded global and recompiles the graph once. - if torch.compiler.is_compiling(): - return - try: - module.select_moe_backend() - except Exception: - pass - - -def get_forward_moe_backend(): - """Resolve forward_moe_backend from the compiled cache copy, else the local def.""" - global _CACHED_FORWARD_MOE_BACKEND - module = _load_cached_moe_utils_module() - if module is not None and hasattr(module, "forward_moe_backend"): - _select_moe_backend_eagerly(module) - _CACHED_FORWARD_MOE_BACKEND = module.forward_moe_backend - return _CACHED_FORWARD_MOE_BACKEND - - _select_moe_backend_eagerly(sys.modules[__name__]) - _CACHED_FORWARD_MOE_BACKEND = forward_moe_backend - return _CACHED_FORWARD_MOE_BACKEND - -# Grouped MM wrapper around torch._grouped_mm; native backward works correctly. - -# Dynamo calls a function decorated with this eagerly at trace time and bakes in its result. -_assume_constant_result = getattr(torch.compiler, "assume_constant_result", None) or \ - getattr(torch._dynamo, "assume_constant_result", lambda fn: fn) - - -def _match_grouped_mm_dtypes(inputs, weight): - """torch._grouped_mm is not on autocast's cast list and needs equal operand dtypes. - - fp16 full finetuning keeps float32 master expert weights; under fp16 autocast the - activations reaching the experts are fp16 and the grouped GEMM raised "expected mat1 and - mat2 to have the same dtype ... Half != float". Do what autocast does for F.linear: run in - the autocast dtype (the weight cast is autograd tracked, so the fp32 master still gets its - grad). Without autocast, promote to the wider dtype. Identity when the dtypes agree. - """ - if inputs.dtype == weight.dtype or not (inputs.is_floating_point() and weight.is_floating_point()): - return inputs, weight - device_type = inputs.device.type - try: - autocast_dtype = ( - torch.get_autocast_dtype(device_type) if torch.is_autocast_enabled(device_type) else None - ) - except (AttributeError, RuntimeError, TypeError): - autocast_dtype = None - if autocast_dtype in (torch.float16, torch.bfloat16): - return inputs.to(autocast_dtype), weight.to(autocast_dtype) - dtype = torch.promote_types(inputs.dtype, weight.dtype) - return inputs.to(dtype), weight.to(dtype) - - -def _grouped_mm_with_backward_fix( - inputs: torch.Tensor, weight: torch.Tensor, offsets: torch.Tensor -) -> torch.Tensor: - """Grouped matmul; passes the weight as a transposed view (no copy) when safe. - - Forcing weight.contiguous() copies the frozen base stack (~805 MB for gate_up on Qwen3-30B, - ~57% of MoE GPU time) every step. torch._grouped_mm takes the non-contiguous view directly, - and a one-time probe confirms the view matches the contiguous copy on this device before we - skip the copy; else we keep the copy. Falls back to a per-group matmul when the device has no - torch._grouped_mm, and on the 16-byte stride error. Bit-exact vs the always-contiguous - path in forward and backward. Small experts may first take the Triton grouped GEMM. - """ - if _triton_grouped_mm_wanted(inputs, weight): - return _triton_grouped_mm(inputs, weight, offsets) - if ( - inputs.dtype == torch.float16 - and weight.dtype == torch.float16 - and _GROUPED_MM_FP16_OP is not None - and torch.compiler.is_compiling() - ): - return _GROUPED_MM_FP16_OP(inputs, weight, offsets) - return _grouped_mm_eager(inputs, weight, offsets) - - -def _grouped_mm_eager(inputs, weight, offsets): - # Mixed dtypes always land here: the Triton and compiled fp16 routes need equal dtypes. - if inputs.dtype != weight.dtype: - inputs, weight = _match_grouped_mm_dtypes(inputs, weight) - inputs = inputs.contiguous() - # The Triton backend is picked precisely when the probe says no, yet its separated - # LoRA delta still routed here. torch 2.8 hard-raises unless `dprops->major == 9` - # (Blas.cpp, sm90_only) and 2.6/2.7 have no `_grouped_mm`, so a LoRA MoE died on - # every card but an H100; 2.9 falls back internally. Probe is cached: one read. - if not _grouped_mm_supported_on(inputs.device): - return _manual_grouped_mm(inputs, weight, offsets) - if not _transposed_view_safe_on(inputs.device): - weight = weight.contiguous() # view path unproven on this device -> safe copy - try: - return torch._grouped_mm(inputs, weight, offs=offsets) - except RuntimeError as exc: - if "strides should be multiple of 16 bytes" not in str(exc): - raise - weight = weight.contiguous() - try: - return torch._grouped_mm(inputs, weight, offs=offsets) - except RuntimeError as exc: - if "strides should be multiple of 16 bytes" not in str(exc): - raise - return _manual_grouped_mm(inputs, weight, offsets) - - -def _grouped_mm_fp16_wgrad_eager(inputs, grad_output, offsets): - """dW[g] = inputs[g]^T @ grad_output[g] per group, as [G, K, N]; empty groups are zero.""" - # torch._grouped_mm rejects stride-0 (the grad of y.sum()) or unaligned operands, which sent eager dW to the - # per-expert loop below: a host sync, and a different reduction than the compiled path. No-op when contiguous. - inputs = inputs.contiguous() - grad_output = grad_output.contiguous() - try: - if _grouped_mm_supported_on(inputs.device): - return torch._grouped_mm(inputs.transpose(-2, -1), grad_output, offs = offsets).contiguous() - except RuntimeError: - pass - out = inputs.new_zeros((offsets.shape[0], inputs.shape[1], grad_output.shape[1])) - start = 0 - for g, end in enumerate(offsets.tolist()): - if end > start: - out[g] = inputs[start:end].transpose(0, 1) @ grad_output[start:end] - start = end - return out - - -def _is_transposed_stack(weight): - # [E, K, N] view of a contiguous [E, N, K] parameter (HF stores gate_up_proj / down_proj as [E, out, in]). - return weight.dim() == 3 and not weight.is_contiguous() and weight.transpose(-2, -1).is_contiguous() - - -def _wgrad_in_weight_layout(wgrad, inputs, grad, offsets, transposed): - # dW^T built contiguous and returned as its view: the leaf's layout, so AccumulateGrad does not copy it. - if transposed: - return wgrad(grad, inputs, offsets).transpose(-2, -1) - return wgrad(inputs, grad, offsets) - - -def _register_grouped_mm_fp16_op(): - """torch._grouped_mm on float16 as an opaque op for compiled callers. The eager kernel takes float16 - (torch 2.10 - 2.14), but its meta/fake impl rejects it (2.14 accepts only with cuBLASLt grouped GEMM, - CUDA 13.3+), so a compiled float16 MoE frame failed to trace and, under suppress_errors, silently ran - eagerly. Eager calls and bf16 never reach this op. Registered once per process: the compiled-cache copy of - this file reuses the first registration.""" - if not hasattr(torch, "library") or not hasattr(torch.library, "custom_op"): - return None - try: - return torch.ops.unsloth_zoo.grouped_mm_fp16.default - except (AttributeError, RuntimeError): - pass - try: - @torch.library.custom_op("unsloth_zoo::grouped_mm_fp16", mutates_args = ()) - def grouped_mm_fp16(inputs: torch.Tensor, weight: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor: - # contiguous: the real kernel pads the row stride to 16 bytes, the fake below does not. - return _grouped_mm_eager(inputs, weight, offsets).contiguous() - - @grouped_mm_fp16.register_fake - def _(inputs, weight, offsets): - return inputs.new_empty((inputs.shape[0], weight.shape[-1])) - - @torch.library.custom_op("unsloth_zoo::grouped_mm_fp16_wgrad", mutates_args = ()) - def grouped_mm_fp16_wgrad(inputs: torch.Tensor, grad_output: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor: - return _grouped_mm_fp16_wgrad_eager(inputs, grad_output, offsets) - - @grouped_mm_fp16_wgrad.register_fake - def _(inputs, grad_output, offsets): - return inputs.new_empty((offsets.shape[0], inputs.shape[1], grad_output.shape[1])) - - def _setup_context(ctx, inputs, output): - x, w, offs = inputs - ctx.save_for_backward(x, w, offs) - ctx.w_transposed = _is_transposed_stack(w) - - def _backward(ctx, grad): - x, w, offs = ctx.saved_tensors - grad = grad.contiguous() - gx = gw = None - if ctx.needs_input_grad[0]: - gx = torch.ops.unsloth_zoo.grouped_mm_fp16(grad, w.transpose(-2, -1), offs) - if ctx.needs_input_grad[1]: - gw = _wgrad_in_weight_layout( - torch.ops.unsloth_zoo.grouped_mm_fp16_wgrad, x.contiguous(), grad, offs, ctx.w_transposed, - ) - return gx, gw, None - - grouped_mm_fp16.register_autograd(_backward, setup_context = _setup_context) - return torch.ops.unsloth_zoo.grouped_mm_fp16.default - except Exception: - return None - - -_GROUPED_MM_FP16_OP = _register_grouped_mm_fp16_op() - - -# Off sm90 / sm100 torch._grouped_mm is a per-expert cuBLAS loop; one Triton launch beats it for small experts. -# Measured max average rows per expert: "lora" (a dim <= 32), "many" / "few" and "dw" / "dw_few" (base dW) at >= / < 64 -# experts. sm75 has no Triton MMA, sm86 is unmeasured, sm90 / sm100 have native grouped GEMMs. sm121 (GB10) takes sm120's -# limits (unsloth-zoo#1661) except "dw" 256, measured on RTX PRO 6000 only. sm80 "few" 512: Triton 4-40% faster at 257-512 -# rows on A100. sm120 "dw" 256: 13-26% faster at 65-256 rows; "dw_few" stays 64 (Triton loses ~10% on Mixtral gate_up). -_TRITON_GROUPED_MM_AUTO_ROWS = { - (8, 0): {"lora": -1, "many": 2048, "few": 512, "dw": 128, "dw_few": 128}, - (12, 0): {"lora": -1, "many": 2048, "few": 256, "dw": 256, "dw_few": 64}, - (12, 1): {"lora": -1, "many": 2048, "few": 256, "dw": 64, "dw_few": 64}, - (8, 9): {"lora": -1, "many": 32, "few": 32, "dw": 0, "dw_few": 0}, -} -_TRITON_GROUPED_MM_CAPABILITY = {} -_TRITON_GROUPED_MM_POLICY = {} - - -def _triton_grouped_mm_capability(index): - cap = _TRITON_GROUPED_MM_CAPABILITY.get(index) - if cap is None: - cap = _TRITON_GROUPED_MM_CAPABILITY[index] = tuple(torch.cuda.get_device_capability(index)) - return cap - - -@_assume_constant_result -def _triton_grouped_mm_max_rows(index, kind = "lora"): - """Row limit for class `kind` on device `index`: 0 = off, -1 = no limit. - UNSLOTH_MOE_GROUPED_TRITON=auto|1|0, UNSLOTH_DISABLE_MOE_TRITON=1, UNSLOTH_MOE_GROUPED_TRITON_MAX_ROWS. - Baked in under torch.compile: an env change after compiling needs torch._dynamo.reset().""" - environ = os.environ - key = ( - index, kind, environ.get("UNSLOTH_MOE_GROUPED_TRITON"), environ.get("UNSLOTH_DISABLE_MOE_TRITON"), - environ.get("UNSLOTH_MOE_GROUPED_TRITON_MAX_ROWS"), - ) - limit = _TRITON_GROUPED_MM_POLICY.get(key) - if limit is None: - limit = _TRITON_GROUPED_MM_POLICY[key] = _triton_grouped_mm_policy(index, kind, *key[2:]) - return limit - - -def _triton_grouped_mm_policy(index, kind, mode, disabled, override): - mode = (mode or "auto").strip().lower() - if mode in ("0", "false", "off") or disabled == "1": - return 0 - if torch.version.hip is not None or _GROUPED_MM_TRITON_OP is None or index is None: - return 0 - try: - from unsloth_zoo.temporary_patches.moe_grouped_fp16 import triton as _triton - if _triton is None: - return 0 - cap = _triton_grouped_mm_capability(index) - except Exception: - return 0 - if mode in ("1", "true", "on"): - limit = -1 if cap >= (8, 0) else 0 - else: - limit = _TRITON_GROUPED_MM_AUTO_ROWS.get(cap, {}).get(kind, 0) - if limit != 0 and override: - try: - limit = int(override) - except ValueError: - pass - return limit - - -def _triton_grouped_mm_kind(K, N, E): - return "lora" if min(K, N) <= 32 else ("many" if E >= 64 else "few") - - -def _rows_within(limit, M, E): - return limit != 0 and (limit < 0 or M <= limit * E) - - -_PLAIN_TENSOR_TYPES = (torch.Tensor, torch.nn.Parameter) - - -def _triton_grouped_mm_wanted(inputs, weight) -> bool: - """Static gate (never reads offsets' values): use Triton when M <= limit * E for the GEMM's class.""" - device = inputs.device - if device.type != "cuda": - return False - if _triton_grouped_mm_max_rows(device.index, "lora") == 0 or inputs.dim() != 2 or weight.dim() != 3 \ - or weight.device != device: - return False - dtype = inputs.dtype - if weight.dtype != dtype or (dtype != torch.bfloat16 and dtype != torch.float16): - return False - # Tensor subclasses (DTensor, float8 / quantized wrappers) keep torch._grouped_mm's own dispatch. - if type(inputs) not in _PLAIN_TENSOR_TYPES or type(weight) not in _PLAIN_TENSOR_TYPES: - return False - E, K, N = weight.shape - if E == 0 or inputs.shape[1] != K or N == 0: - return False - return _triton_grouped_mm_rows_ok(device, inputs.shape[0], E, K, N) - - -def _triton_grouped_mm_rows_ok(device, M, E, K, N): - limit = _triton_grouped_mm_max_rows(device.index, _triton_grouped_mm_kind(K, N, E)) - return _rows_within(limit, M, E) - - -def _grouped_mm_triton_impl(inputs, weight, offsets): - from unsloth_zoo.temporary_patches import moe_grouped_fp16 as mg - if not mg.triton_grouped_available(inputs.device): - mg.GENERIC_CALLS["fallback"] += 1 - return _grouped_mm_eager(inputs, weight, offsets).contiguous() - return mg.generic_grouped_mm(inputs, weight, offsets) - - -def _grouped_mm_triton_wgrad_impl(inputs, grad, offsets): - from unsloth_zoo.temporary_patches import moe_grouped_fp16 as mg - if not mg.triton_grouped_available(inputs.device): - mg.GENERIC_CALLS["fallback"] += 1 - return _grouped_mm_fp16_wgrad_eager(inputs, grad, offsets).contiguous() - E, K, N = offsets.shape[0], inputs.shape[1], grad.shape[1] - dw_kind = "dw" if E >= 64 else "dw_few" - if min(K, N) > 32 and not _rows_within(_triton_grouped_mm_max_rows(inputs.device.index, dw_kind), inputs.shape[0], E): - # A trainable base stack's dW (full finetuning): cuBLAS beats the Triton reduction at these sizes. - return _grouped_mm_fp16_wgrad_eager(inputs, grad, offsets).contiguous() - return mg.generic_grouped_wgrad(inputs, grad, offsets, E) - - -class _GroupedMMTriton(torch.autograd.Function): - """Eager twin of unsloth_zoo::grouped_mm_triton, skipping the custom op's ~30 us dispatch cost.""" - - @staticmethod - def forward(ctx, inputs, weight, offsets): - ctx.save_for_backward( - inputs if ctx.needs_input_grad[1] else None, weight if ctx.needs_input_grad[0] else None, offsets, - ) - ctx.w_transposed = _is_transposed_stack(weight) - return _grouped_mm_triton_impl(inputs, weight, offsets) - - @staticmethod - def backward(ctx, grad): - x, w, offs = ctx.saved_tensors - if torch.is_grad_enabled(): - return _grouped_mm_differentiable_backward(ctx, x, w, offs, grad) - gx = gw = None - if ctx.needs_input_grad[0]: - gx = _grouped_mm_triton_impl(grad, w.transpose(-2, -1), offs) - if ctx.needs_input_grad[1]: - gw = _wgrad_in_weight_layout(_grouped_mm_triton_wgrad_impl, x, grad, offs, ctx.w_transposed) - return gx, gw, None - - -def _grouped_mm_differentiable_backward(ctx, x, w, offs, grad): - # create_graph=True: the raw kernels would detach the grads, so take torch._grouped_mm (double backward). - gx = _grouped_mm_eager(grad, w.transpose(-2, -1), offs) if ctx.needs_input_grad[0] else None - gw = _grouped_mm_fp16_wgrad_eager(x, grad, offs) if ctx.needs_input_grad[1] else None - return gx, gw, None - - -def _triton_grouped_mm(inputs, weight, offsets): - """torch._grouped_mm on the Triton kernels: custom op when traced, autograd.Function in eager.""" - if torch.compiler.is_compiling(): - return _GROUPED_MM_TRITON_OP(inputs, weight, offsets) - return _GroupedMMTriton.apply(inputs, weight, offsets) - - -def _register_grouped_mm_triton_op(): - """torch._grouped_mm(inputs [M, K], weight [E, K, N], offs = cumulative int32 ends) on moe_grouped_fp16's - Triton kernels, as an opaque op so compiled MoE frames trace it (its autograd.Function does not). - Counts come from the offsets on device: no host sync. Rows past offs[-1] are left unwritten, as - torch._grouped_mm leaves them. Registered once per process (the compiled-cache copy reuses it).""" - if not hasattr(torch, "library") or not hasattr(torch.library, "custom_op"): - return None - try: - return torch.ops.unsloth_zoo.grouped_mm_triton.default - except (AttributeError, RuntimeError): - pass - try: - @torch.library.custom_op("unsloth_zoo::grouped_mm_triton", mutates_args = ()) - def grouped_mm_triton(inputs: torch.Tensor, weight: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor: - return _grouped_mm_triton_impl(inputs, weight, offsets) - - @grouped_mm_triton.register_fake - def _(inputs, weight, offsets): - return inputs.new_empty((inputs.shape[0], weight.shape[-1])) - - @torch.library.custom_op("unsloth_zoo::grouped_mm_triton_wgrad", mutates_args = ()) - def grouped_mm_triton_wgrad(inputs: torch.Tensor, grad: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor: - return _grouped_mm_triton_wgrad_impl(inputs, grad, offsets) - - @grouped_mm_triton_wgrad.register_fake - def _(inputs, grad, offsets): - return inputs.new_empty((offsets.shape[0], inputs.shape[1], grad.shape[1])) - - def _setup_context(ctx, inputs, output): - x, w, offs = inputs - # Only what each gradient reads: a frozen base keeps no activation alive for a dW it never takes. - ctx.save_for_backward(x if w.requires_grad else None, w if x.requires_grad else None, offs) - ctx.w_transposed = _is_transposed_stack(w) - - def _backward(ctx, grad): - x, w, offs = ctx.saved_tensors - if torch.is_grad_enabled(): - return _grouped_mm_differentiable_backward(ctx, x, w, offs, grad) - gx = gw = None - if ctx.needs_input_grad[0]: - gx = torch.ops.unsloth_zoo.grouped_mm_triton(grad, w.transpose(-2, -1), offs) - if ctx.needs_input_grad[1]: - gw = _wgrad_in_weight_layout( - torch.ops.unsloth_zoo.grouped_mm_triton_wgrad, x, grad, offs, ctx.w_transposed, - ) - return gx, gw, None - - grouped_mm_triton.register_autograd(_backward, setup_context = _setup_context) - return torch.ops.unsloth_zoo.grouped_mm_triton.default - except Exception: - return None - - -_GROUPED_MM_TRITON_OP = _register_grouped_mm_triton_op() - - -def _grouped_matmul_loop(inputs, weight, offsets, bounds = None): - """out[s:e] = inputs[s:e] @ weight[g], group by group. No autograd. - - `bounds` is the already-decoded group ends: the signature packs the offsets into - its own transfer, so re-reading them here was a second stream synchronization per - grouped matmul, over every layer's base and LoRA projections. - """ - outputs = [] - start = 0 - if bounds is None: bounds = offsets.detach().cpu().tolist() - for expert_idx, end in enumerate(bounds): - if start < end: - outputs.append(torch.matmul(inputs[start:end], weight[expert_idx])) - start = end - if outputs: - return torch.cat(outputs, dim=0) - return inputs.new_empty((0, weight.shape[-1])) - - -# Projection strides for the routing signature. A CHECKSUM, not an identity: four -# projections leave a `hidden - 4` null space, so rows differing within it project -# alike; an exact identity means copying the whole `[T, hidden]` input, the 512MB -# transient this file avoids. A swap escapes only by matching ALL of these AND the -# norm below, each term weighted by the row index. -_SIGNATURE_STRIDES = (12.9898, 78.233, 43.7585, 96.4271) - -# Row 2-norm, carried beside the projections. Not linear in the row, so that null -# space does not reach it: rows differing within it must also match norms to swap unseen. -_SIGNATURE_EXTRA = 1 - -# Trailing checksum entries of a packed signature, not group boundaries. Named because -# offsets are read back by slicing it off, and a bare `[:-1]` read checksums as experts. -_SIGNATURE_WIDTH = len(_SIGNATURE_STRIDES) + _SIGNATURE_EXTRA - - -def _routing_signature(inputs, offsets): - """Offsets plus an order-sensitive checksum of the expert-sorted rows. - - Offsets are only the routing histogram, so a replay swapping two tokens between - equal-sized experts leaves them identical while `inputs` holds a different - sequence; the checksum moves with the rows, so the swap is visible. Bit-cast into - the offsets vector so the whole thing is one device sync, the one the group loop - already pays. - - Rows reduce through a fixed pseudo-random projection, not a sum: a sum ignores - WHERE a row's values sit, so `[1, 0]` and `[0, 1]` reduced alike and a swap across - an expert boundary went unseen, letting the guard accept gradients from a different - routing. The projection derives from the hidden size alone, so forward and - recompute build the same one without carrying state. - - The matvec runs in the input's own dtype, not an FP32 copy: upcasting on the way in - materialised a whole `[routed_tokens, hidden]` transient, 512MB at 32K by 4096, to - produce one number per row -- an OOM risk on exactly the memory-constrained runs - this fallback exists for. cuBLAS accumulates a half-precision gemv in FP32 anyway. - """ - inputs = inputs.detach() - hidden = inputs.shape[-1] - # Autocast off, explicitly: `mv`/`mm` are on the autocast lower-precision list and - # only ONE of the two calls sees it (forward with the caller's autocast live, - # backward with it disabled), so FP32 inputs hashed bf16 one side and fp32 the - # other and every backward raised the routing error below on unchanged routing. - with torch.autocast(device_type = inputs.device.type, enabled = False): - weights = torch.linspace( - 1.0, 2.0, hidden, device = inputs.device, dtype = torch.float32) - # Irrational stride: a linear ramp alone sums the same under a reversal. - strides = torch.tensor( - _SIGNATURE_STRIDES, device = inputs.device, dtype = torch.float32) - weights = torch.sin( - weights.unsqueeze(1) * strides).to(inputs.dtype) - # Several projections, not one: a single dot maps each row to one scalar, so - # rows orthogonal to it hash like the zero row (`[p[1], -p[0], 0, ...]` and - # zeros both give exactly 0) and swapping that pair across an expert boundary - # survived the check. - rows = (inputs @ weights).float() - # Plus the norm, not linear in the row, so the projections' shared null space - # does not reach it. `vector_norm` accumulates in fp32 from the input's own - # dtype; `(x.float() * x.float()).sum(...)` held three concurrent 512MB - # temporaries at the 32K-by-4096 shape this fallback targets. Promoted, not - # pinned to fp32: `vector_norm` refuses a dtype that narrows its input. - norm = torch.linalg.vector_norm( - inputs, dim = -1, keepdim = True, - dtype = torch.promote_types(inputs.dtype, torch.float32)) - rows = torch.cat((rows, norm.to(rows.dtype)), -1) - ramp = torch.arange( - 1, rows.shape[0] + 1, device = rows.device, dtype = rows.dtype) - checksum = (rows * ramp.unsqueeze(1)).sum(0) - packed = torch.cat(( - offsets.detach().reshape(-1).to(torch.int64), - checksum.view(torch.int32).to(torch.int64), - )) - return packed.cpu().tolist() - - -class _ManualGroupedMM(torch.autograd.Function): - """The loop above, saving what torch._grouped_mm saves and nothing else. - - A Function rather than plain autograd because the naive loop tapes every per-group - SLICE, whose shape is the group size the router decides. Non-reentrant checkpointing - replays the forward and compares saved metadata, so any routing difference between - the two passes surfaces as - - CheckpointError: Recomputed values ... have different metadata - saved: torch.Size([38, 8]) recomputed: torch.Size([39, 8]) - - one row apart, in a hundred groups at once. `torch._grouped_mm` never shows that: - it is one op saving the whole `[T, K]` input, so a drop-in has to save the same - shape-stable set (inputs, weight, offsets) and rebuild the slices in backward. - - This does not make routing deterministic, and does not pretend to: it restores the - numerics the fused path already has on an H100. - """ - @staticmethod - def forward(ctx, inputs, weight, offsets): - # A plain list, NOT a tensor: non-reentrant checkpointing swaps the saved - # TENSORS for the replay's, so a saved `offsets` reports the replay's routing, - # never the forward's. This copy survives, and backward compares the two. - ctx.forward_routing = routing = _routing_signature(inputs, offsets) - ctx.save_for_backward(inputs, weight, offsets) - with torch.no_grad(): - return _grouped_matmul_loop( - inputs, weight, offsets, routing[:-_SIGNATURE_WIDTH]) - - @staticmethod - def backward(ctx, grad_output): - inputs, weight, offsets = ctx.saved_tensors - routing = _routing_signature(inputs, offsets) - if routing != ctx.forward_routing: - # Shape-stable saves would let this through silently, and pairing the - # original `grad_output` with the replay's partition is a gradient for a - # routing that never produced the loss. Louder than CheckpointError. - were = ctx.forward_routing[:-_SIGNATURE_WIDTH] - now = routing[:-_SIGNATURE_WIDTH] - how = ("the same experts in a different order" - if were == now else f"expert ends {were} then {now}") - raise RuntimeError( - "Unsloth: the MoE router assigned tokens differently in the " - f"activation-checkpoint replay than in the forward ({how}), so " - "the gradients would belong to a routing that never produced " - "the loss. Turn gradient checkpointing off for this run, or " - "make the router deterministic." - ) - bounds = routing[:-_SIGNATURE_WIDTH] - need_x, need_w, _ = ctx.needs_input_grad - grad_output = grad_output.contiguous() - # `backward` runs OUTSIDE the forward's autocast, so `grad_output` can be bf16 - # while the saved tensors are still fp32 and the first matmul raises. Aligned - # by hand, not with `torch.amp.custom_fwd/custom_bwd`, whose `device_type` is - # fixed at class definition: this fallback also runs on CPU and XPU. - compute_dtype = grad_output.dtype - grad_inputs = torch.zeros_like(inputs) if need_x else None - grad_weight = torch.zeros_like(weight) if need_w else None - # Write each group straight into its slice when no cast is due. The assignment - # below allocates a temporary and copies it in; `out=` does neither, and there - # are two per group -- 256 on a 128-expert layer, 37% of this loop measured. - # A cast needs the temporary anyway, and a non-contiguous destination would put - # the copy back, so both keep the plain path. - # Grad mode first: it is off for an ordinary backward and ON only under - # `create_graph = True`, where `out=` raises "functions with out=... arguments - # don't support automatic differentiation". The fused path double-backwards - # fine, so the fallback has to as well; the plain branch is differentiable. - direct = ( - not torch.is_grad_enabled() - and inputs.dtype == weight.dtype == compute_dtype - and (grad_inputs is None or grad_inputs.is_contiguous()) - and (grad_weight is None or grad_weight.is_contiguous()) - ) - start = 0 - for expert_idx, end in enumerate(bounds): - if start < end: - g = grad_output[start:end] - if need_x: - w = weight[expert_idx] - if direct: - torch.matmul(g, w.transpose(-2, -1), - out = grad_inputs[start:end]) - else: - grad_inputs[start:end] = ( - g @ w.to(compute_dtype).transpose(-2, -1) - ).to(inputs.dtype) - if need_w: - if direct: - torch.matmul(inputs[start:end].transpose(-2, -1), g, - out = grad_weight[expert_idx]) - else: - x = inputs[start:end].to(compute_dtype) - grad_weight[expert_idx] = ( - x.transpose(-2, -1) @ g).to(weight.dtype) - start = end - return grad_inputs, grad_weight, None - - -def _manual_grouped_mm( - inputs: torch.Tensor, weight: torch.Tensor, offsets: torch.Tensor -) -> torch.Tensor: - """Differentiable grouped matmul fallback for torch._grouped_mm alignment gaps. - - Not compilable, and marked so: group boundaries come off a tensor, so `start < end` - is a data-dependent branch Dynamo cannot guard, and a tensorized rewrite is worse - (gathering `weight[expert_id]` per row materializes `[T, K, N]`; masking runs every - expert over every row). `torch.compiler.disable` makes a compiled caller break - cleanly here rather than abort mid-trace; under `fullgraph = True` a break is still - fatal, which is what the eager fallback in `temporary_patches/utils.py` is for. - """ - # Grad off means `apply` builds no autograd node, so nothing will ever read the - # signature: skip the Function rather than pay a `[T, hidden]` projection per call. - # Covers `inference_mode` and `no_grad` alike. - # - # KNOWN GAP, deliberately not closed here. A REENTRANT checkpoint also runs its - # loss-producing forward with grad off and its replay with grad ON, so the replay - # compares its own routing against itself and a reroute passes silently. Not - # fixable from inside this function: `ctx` cannot cross a reentrant boundary, and a - # side channel keyed on the weight fails too because - # `_canonical_lora_weights_for_grouped_mm` rebuilds a contiguous tensor per call. - # It needs `unsloth_checkpoint` to announce its region and call order. The - # non-reentrant path is unaffected: its forward runs with grad on, so the forward - # ctx survives and only its SAVED TENSORS are swapped. - if not torch.is_grad_enabled(): - return _grouped_matmul_loop(inputs, weight, offsets) - return _ManualGroupedMM.apply(inputs, weight, offsets) - - -if hasattr(torch, "compiler") and hasattr(torch.compiler, "disable"): - _manual_grouped_mm = torch.compiler.disable(_manual_grouped_mm) - - -# Recompute-in-backward for the frozen base expert GEMM: the dequantized bf16 stack -# is rebuilt from the 4-bit Params4bit in backward (dX only; the base is frozen and -# LoRA is a separate additive grouped_mm) instead of being pinned on the tape. Output -# is unchanged. See _moe_recompute_enabled for the pin-vs-recompute policy. - - -def _base_is_recomputable(source) -> bool: - """True iff the base expert weight can be rebuilt in backward (frozen and - grouped-mm capable). A trainable or unsupported base must use the pinned path.""" - try: - if not _should_use_separated_lora(): # merged LoRA folds the delta into base - return False - if not _check_torch_grouped_mm_supported(): - return False - param = source - while hasattr(param, "base_layer"): - param = param.base_layer - if HAS_BNB and Params4bit is not None and isinstance(param, Params4bit): - if getattr(param, "quant_state", None) is None: - return False - return not param.requires_grad - if is_mxfp4_expert_param(param): - return not param.requires_grad - if isinstance(param, torch.Tensor): - return (not param.requires_grad) and param.dtype in ( - torch.bfloat16, torch.float16, torch.float32, - ) - except Exception: - return False - return False - - -def _moe_recompute_default(prefer_memory: bool = False) -> bool: - """Pin-vs-recompute decision independent of the source weight. - - Adaptive by default: pin (return False) inside a gradient-checkpoint recompute - pass, where the stack is rebuilt immediately before the layer's own backward so - the pin is momentary and cheap; otherwise recompute (return True), so a - non-checkpointed forward does not hold every layer's dense stack across the whole - backward. UNSLOTH_MOE_RECOMPUTE overrides it: "1" forces recompute (max memory - saving), "0" forces pinning (max speed for memory-rich runs). - - ``prefer_memory`` biases the no-override case toward recompute even inside a GC - recompute pass. It is set for bases whose pinned form is a large dense dequant of - a compressed weight (bnb 4-bit MoE experts): there the "momentary" pin still - materializes the full bf16 expert stack the 4-bit storage exists to avoid, which - can be several GiB per layer on large MoEs, so recompute is the better default - when the model is already memory-constrained (4-bit + gradient checkpointing). - - The GC branch reads a thread-local; under torch.compile of the MoE forward that - read can be traced away and frozen at the first trace, so the adaptive choice may - not re-evaluate per GC pass. That only trades memory for speed (pin and recompute - are dX-identical), never correctness, and the grouped GEMM path generally runs - eager anyway; set UNSLOTH_MOE_RECOMPUTE explicitly to pin the choice if needed.""" - override = os.environ.get("UNSLOTH_MOE_RECOMPUTE") - if override == "1": - return True - if override == "0": - return False - if prefer_memory: - return True - try: - from unsloth_zoo.gradient_checkpointing import in_gradient_checkpoint_recompute - return not in_gradient_checkpoint_recompute() - except Exception: - return True # safe default: recompute rather than pin across a full backward - - -def _source_pins_large_dequant(source) -> bool: - """True iff pinning this base means holding a large dense dequant of a compressed - weight (a frozen bnb 4-bit MoE expert). For these the pinned path materializes the - full bf16 expert stack, so recompute is the memory-preserving default under - gradient checkpointing; a plain bf16/fp16/fp32 base pins its own storage (no - extra dequant) and keeps the speed-oriented adaptive policy.""" - if not (HAS_BNB and Params4bit is not None): - return False - try: - param = source - while hasattr(param, "base_layer"): - param = param.base_layer - return isinstance(param, Params4bit) and getattr(param, "quant_state", None) is not None - except Exception: - return False - - -def _in_gc_recompute() -> bool: - try: - from unsloth_zoo.gradient_checkpointing import in_gradient_checkpoint_recompute - return bool(in_gradient_checkpoint_recompute()) - except Exception: - return False - - -_MOMENTARY_PIN_HEADROOM = 2.0 # free memory must cover this many copies of one layer's stack - - -def _momentary_pin_fits(source, dtype = None) -> bool: - """Whether one layer's dequantized stack fits in free + cached-unused memory; unmeasurable means False.""" - try: - param = source - while hasattr(param, "base_layer"): - param = param.base_layer - device = param.device - backend = getattr(torch, device.type, None) if device.type in ("cuda", "xpu") else None - if backend is None or not callable(getattr(backend, "mem_get_info", None)): - return False - shape = _logical_expert_shape(param) - if not shape: - return False - dtype = dtype or getattr(getattr(param, "quant_state", None), "dtype", None) or torch.bfloat16 - need = math.prod(int(d) for d in shape) * dtype.itemsize - # Cached briefly: mem_get_info per source per replay costs more than the dequant it saves. - key = (device.type, device.index, need, _MOMENTARY_PIN_HEADROOM) - now = time.monotonic() - hit = _MOMENTARY_PIN_FITS_CACHE.get(key) - if hit is not None and now - hit[1] < _MOMENTARY_PIN_FITS_TTL_S: - return hit[0] - free, _ = backend.mem_get_info(device) - free += backend.memory_reserved(device) - backend.memory_allocated(device) - fits = free >= _MOMENTARY_PIN_HEADROOM * need - _MOMENTARY_PIN_FITS_CACHE[key] = (fits, now) - return fits - except Exception: - return False - - -_MOMENTARY_PIN_FITS_CACHE = {} -_MOMENTARY_PIN_FITS_TTL_S = 0.5 - - -def _moe_recompute_enabled(source, dtype = None) -> bool: - """Whether to recompute the dequantized base stack in backward (True) or pin it - for reuse (False). Only a frozen, grouped-mm-capable base can be recomputed; for - everything else the pinned eager path is used. A bnb 4-bit base recomputes unless - UNSLOTH_MOE_GC_REPLAY_PIN=1 and, inside a gradient-checkpoint replay, the stack fits: - the pin then saves a dequant for the same layer's backward but raises peak memory.""" - if not _base_is_recomputable(source): - return False - override = os.environ.get("UNSLOTH_MOE_RECOMPUTE") - if override == "1": - return True - if override == "0": - return False - if _source_pins_large_dequant(source): - if ( - os.environ.get("UNSLOTH_MOE_GC_REPLAY_PIN") == "1" - and _in_gc_recompute() - and _momentary_pin_fits(source, dtype = dtype) - ): - return False - return True - return _moe_recompute_default() - - -class _GroupedMMRecompute(torch.autograd.Function): - """grouped_mm(inputs, W) for a frozen W from weight_provider(): saves only offsets and rebuilds - W in backward (dX only) instead of pinning the dense stack.""" - - @staticmethod - def forward(ctx, inputs, offsets, weight_provider): - ctx.weight_provider = weight_provider - ctx.save_for_backward(offsets) # inputs is unused in backward (frozen base -> dX only) - with torch.no_grad(): - out = _grouped_mm_with_backward_fix(inputs, weight_provider(), offsets) - return out - - @staticmethod - def backward(ctx, grad_output): - (offsets,) = ctx.saved_tensors - with torch.no_grad(): - transposed = getattr(ctx.weight_provider, "transposed", None) - # The view, not a copy: _grouped_mm_with_backward_fix copies only where the view path is unproven. - weight_t = transposed() if transposed is not None else ctx.weight_provider().transpose(-2, -1) - grad_input = _grouped_mm_with_backward_fix(grad_output.contiguous(), weight_t, offsets) - return grad_input, None, None - - -_EXPERT_BIAS_ONE_HOT_ELEMENTS = 1 << 24 - - -class _ExpertBiasGather(torch.autograd.Function): - """bias.index_select(0, expert_ids) whose backward is one-hot(expert_ids) @ grad: one fp32-accumulating - GEMM instead of a 16-bit index_add_, which sm80 emulates with CAS loops and which rounds every partial sum.""" - - @staticmethod - def forward(ctx, bias, expert_ids): - ctx.save_for_backward(expert_ids) - ctx.num_experts = bias.shape[0] - return bias.index_select(0, expert_ids) - - @staticmethod - def backward(ctx, grad): - (expert_ids,) = ctx.saved_tensors - experts = torch.arange(ctx.num_experts, device = expert_ids.device, dtype = expert_ids.dtype).unsqueeze(1) - # The one-hot is at most _EXPERT_BIAS_ONE_HOT_ELEMENTS: long contexts go in row chunks, summed in fp32. - chunk = max(1, _EXPERT_BIAS_ONE_HOT_ELEMENTS // ctx.num_experts) - rows = grad.shape[0] - # Autocast off: under fp16 autocast (incl. compiled backward's "same_as_forward") `mm` rounds the - # per-expert sum of fp32 grads to fp16 and overflows past 65504. - with torch.autocast(device_type = grad.device.type, enabled = False): - if rows <= chunk: - return (expert_ids.unsqueeze(0) == experts).to(grad.dtype) @ grad, None - # GEMMs only (no index_add_ atomics) for reproducibility; fp32 partials, since 16-bit chunk sums can overflow to inf. - if torch.compiler.is_compiling(): - # A loop over symbolic rows would specialize them: one uncapped GEMM instead. - return _one_hot_mm_fp32(expert_ids.unsqueeze(0) == experts, grad).to(grad.dtype), None - if not _mm_fp32_out_ok(grad): - # The fallback upcasts each grad chunk to fp32: cap that copy too. - chunk = max(1, _EXPERT_BIAS_ONE_HOT_ELEMENTS // max(ctx.num_experts, grad.shape[1])) - out = torch.zeros(ctx.num_experts, grad.shape[1], dtype = torch.float32, device = grad.device) - for start in range(0, rows, chunk): - out += _one_hot_mm_fp32(expert_ids[start : start + chunk].unsqueeze(0) == experts, grad[start : start + chunk]) - return out.to(grad.dtype), None - - -# torch.mm(..., out_dtype = torch.float32): 16-bit inputs, fp32 output (aten::mm.dtype, CUDA only, torch 2.9+). -_MM_OUT_DTYPE = "dtype" in torch.ops.aten.mm.overloads() - - -def _mm_fp32_out_ok(grad): - return _MM_OUT_DTYPE and grad.is_cuda and torch.version.hip is None and grad.dtype in (torch.float16, torch.bfloat16) - - -def _one_hot_mm_fp32(one_hot, grad): - """one_hot (bool [E, rows]) @ grad as fp32 [E, D], with no 16-bit rounding of the sum.""" - if grad.dtype == torch.float32: - return one_hot.to(grad.dtype) @ grad - if _mm_fp32_out_ok(grad): - return torch.mm(one_hot.to(grad.dtype), grad, out_dtype = torch.float32) - return one_hot.to(torch.float32) @ grad.float() - - -def _gather_expert_bias(bias, expert_ids): - expert_ids = expert_ids.to(bias.device) - if bias.requires_grad and torch.is_grad_enabled() and bias.dim() == 2: - return _ExpertBiasGather.apply(bias, expert_ids) - return bias.index_select(0, expert_ids) - - -def _base_grouped_mm(inputs, offsets, weight_provider, recompute, plan = None): - """recompute -> rebuild W in backward; else the prior eager grouped_mm. - plan (recompute only): rebuild just the planned experts, chunk by chunk.""" - if recompute: - if plan is not None: - return _GroupedMMSelectiveRecompute.apply(inputs, weight_provider, plan) - return _GroupedMMRecompute.apply(inputs, offsets, weight_provider) - return _grouped_mm_with_backward_fix(inputs, weight_provider(), offsets) - - -def _nf4_fused_base_mm(inputs, offsets, source, proj_type, hidden_dim, model_type, experts_module): - """UNSLOTH_MOE_NF4_FUSED grouped GEMM on a frozen NF4 stack; None keeps the dequant path.""" - if os.environ.get("UNSLOTH_MOE_NF4_FUSED", "0") in ("", "0") or not _base_is_recomputable(source): - return None - try: - from unsloth_zoo.temporary_patches.moe_nf4_grouped import maybe_nf4_grouped_mm - except ImportError: - return None - preprocess = lambda w: preprocess_weight(w, proj_type, hidden_dim, model_type, experts_module = experts_module) - return maybe_nf4_grouped_mm(inputs, offsets, source, preprocess) - - -# UNSLOTH_MOE_SELECTIVE_DEQUANT: "auto" = num_rows experts on device (sync-free) if rows <= 7/8 E; -# "1" = exact routed set (one host sync); "0" = full stack. -SELECTIVE_DEQUANT_STATS = {"plans": 0, "dequant_calls": 0, "experts_dequantized": 0, "experts_total": 0} - - -class _SelectiveDequantPlan: - __slots__ = ("num_experts", "num_rows", "chunks") - - def __init__(self, num_experts, num_rows, chunks): - self.num_experts = num_experts - self.num_rows = num_rows - self.chunks = chunks # [(expert slice or int64 ids, row_lo, row_hi, offsets relative to row_lo)] - - -def _selective_dequant_chunk_experts() -> int: - try: - return max(int(os.environ.get("UNSLOTH_MOE_DEQUANT_CHUNK_EXPERTS", "0")), 0) - except ValueError: - return 0 - - -def _bnb4bit_expert_split(param): - """(E, per-expert shape) when every expert of this Params4bit owns whole packed elements - and whole absmax blocks, so a gather of experts is exact.""" - while hasattr(param, "base_layer"): - param = param.base_layer - if not (HAS_BNB and isinstance(param, Params4bit)): - return None - qs = getattr(param, "quant_state", None) - if qs is None or param.requires_grad: - return None - shape = getattr(param, "_original_shape", None) or qs.shape - if shape is None or len(shape) != 3: - return None - E, per_expert = int(shape[0]), int(shape[1]) * int(shape[2]) - weights_per_element = 2 * param.data.element_size() - if E == 0 or per_expert == 0 or per_expert % qs.blocksize or per_expert % weights_per_element: - return None - if param.data.numel() * weights_per_element != E * per_expert or qs.absmax.numel() != E * (per_expert // qs.blocksize): - return None - return E, tuple(int(x) for x in shape[1:]) - - -# Gathered ids are decoded in pieces of at least _SELECTIVE_GATHER_MIN experts and at most 1/8 of -# the request, so the packed copy stays near 1/32 of the fp16 output without a launch per expert. -_SELECTIVE_GATHER_PIECES = 8 -_SELECTIVE_GATHER_MIN = 32 - - -def _dequantize_bnb4bit_experts(param, experts, target_dtype = None): - """Dequantize only `experts` of a packed 4-bit stack: a slice reads views of the packed bytes - and absmax blocks, an int64 id tensor gathers them in a few pieces. Each piece is a smaller - QuantState decoded by the same kernels, so every element matches the full-stack read.""" - from bitsandbytes.functional import QuantState - - while hasattr(param, "base_layer"): - param = param.base_layer - E, inner = _bnb4bit_expert_split(param) - qs = param.quant_state - blocks = inner[0] * inner[1] // qs.blocksize - packed_all = param.data.reshape(E, -1) - absmax_all, offset, absmax2_all = qs.absmax, qs.offset, None - if qs.nested: - s2 = qs.state2 - if blocks % s2.blocksize == 0: - absmax2_all = s2.absmax.reshape(E, -1) - else: - # Second-level blocks straddle experts: denest as bitsandbytes does (multiply, then - # add the offset in fp32), which is what both 4-bit decoders compute per block. - absmax_all = (bnb.functional.dequantize_blockwise(absmax_all, s2) + offset).float() - offset = None - absmax_all = absmax_all.reshape(E, blocks) - - if isinstance(experts, slice): - n = len(range(*experts.indices(E))) - pieces = [experts] - else: - n = int(experts.numel()) - pieces = list(experts.split(max(_SELECTIVE_GATHER_MIN, -(-n // _SELECTIVE_GATHER_PIECES)))) - - def take(t, piece): - return t[piece] if isinstance(piece, slice) else t.index_select(0, piece) - - out = None - if len(pieces) > 1: - out = torch.empty((n,) + inner, dtype = qs.dtype, device = packed_all.device) - row = 0 - for piece in pieces: - packed = take(packed_all, piece) - k = packed.shape[0] - dest = out[row : row + k] if out is not None else None - state2 = None - if absmax2_all is not None: - state2 = QuantState( - absmax = take(absmax2_all, piece).reshape(-1), code = qs.state2.code, - blocksize = qs.state2.blocksize, quant_type = qs.state2.quant_type, dtype = qs.state2.dtype, - ) - piece_shape = (k,) + inner - sub = QuantState( - absmax = take(absmax_all, piece).reshape(-1), shape = torch.Size(piece_shape), code = qs.code, - blocksize = qs.blocksize, quant_type = qs.quant_type, dtype = qs.dtype, offset = offset, state2 = state2, - ) - weight = None - try: - from unsloth_zoo.temporary_patches.moe_triton_kernels import nf4_dequant_triton - weight = nf4_dequant_triton(packed, sub, piece_shape, out = dest) - except ImportError: - pass - if weight is None: - sub_param = torch.Tensor._make_subclass(Params4bit, packed.reshape(-1, 1)) - sub_param.quant_state = sub - sub_param._original_shape = torch.Size(piece_shape) - slicer = _get_dequantize_4bit_in_slices() - weight = slicer(sub_param) if slicer is not None else None - if weight is None: - weight = bnb.functional.dequantize_4bit(packed.reshape(-1, 1), sub) - del packed, sub, state2 - if dest is None: - out = weight.reshape(piece_shape) - elif weight.data_ptr() != dest.data_ptr(): - dest.copy_(weight.reshape(piece_shape)) - row += k - del weight, dest - if not torch.compiler.is_compiling(): # a traced global write re-traces every call - SELECTIVE_DEQUANT_STATS["dequant_calls"] += 1 - SELECTIVE_DEQUANT_STATS["experts_dequantized"] += n - SELECTIVE_DEQUANT_STATS["experts_total"] += E - return out.to(target_dtype) if target_dtype is not None else out - - -def _selective_dequant_plan(sources, num_tokens_per_expert, num_rows): - """Chunks of experts to dequantize, or None to keep the full-stack path (off, traced, - capturing, a source that cannot be split exactly, or nothing to save).""" - mode = os.environ.get("UNSLOTH_MOE_SELECTIVE_DEQUANT", "auto") - selective = mode == "1" - chunk = _selective_dequant_chunk_experts() - if not selective and chunk == 0: - if mode == "0": - return None - return _capped_dequant_plan(sources, num_tokens_per_expert, num_rows) - if torch.compiler.is_compiling(): - return None - device = num_tokens_per_expert.device - if device.type == "cuda" and torch.cuda.is_current_stream_capturing(): - return None - E = int(num_tokens_per_expert.numel()) - for source in sources: - split = _bnb4bit_expert_split(source) - if split is None or split[0] != E: - return None - counts = num_tokens_per_expert.tolist() - experts = [e for e, c in enumerate(counts) if c > 0] if selective else list(range(E)) - if not experts: - return None - if not chunk and 8 * len(experts) > 7 * E: - return None # under 1/8 of the stack to save, not worth the gathers - chunk = chunk or len(experts) - ends = [] - total = 0 - for c in counts: - total += c - ends.append(total) - chunks = [] - for i in range(0, len(experts), chunk): - ids = experts[i : i + chunk] - row_lo = ends[ids[0]] - counts[ids[0]] - row_hi = ends[ids[-1]] - if row_hi == row_lo: - continue # chunked without selective: no token routed to any of these experts - contiguous = ids[-1] - ids[0] + 1 == len(ids) - chunks.append(( - slice(ids[0], ids[-1] + 1) if contiguous else torch.tensor(ids, dtype = torch.int64, device = device), - row_lo, - row_hi, - torch.tensor([ends[e] - row_lo for e in ids], dtype = torch.int32, device = device), - )) - if not chunks: - return None - SELECTIVE_DEQUANT_STATS["plans"] += 1 - return _SelectiveDequantPlan(E, num_rows, chunks) - - -def _capped_dequant_plan(sources, num_tokens_per_expert, num_rows): - """Sync-free plan from shapes only: num_rows experts, routed ids ascending, then unrouted padding.""" - E = int(num_tokens_per_expert.numel()) - if num_rows <= 0 or 8 * num_rows > 7 * E: - return None - for source in sources: - split = _bnb4bit_expert_split(source) - if split is None or split[0] != E: - return None - counts = num_tokens_per_expert - unrouted_last = (counts == 0).to(torch.int64) * E + torch.arange(E, device = counts.device) - experts = torch.argsort(unrouted_last)[:num_rows] - offsets = torch.cumsum(counts.index_select(0, experts), dim = 0, dtype = torch.int32) - if not torch.compiler.is_compiling(): - SELECTIVE_DEQUANT_STATS["plans"] += 1 - return _SelectiveDequantPlan(E, num_rows, [(experts, 0, num_rows, offsets)]) - - -def _slice_takes_triton(inputs, weight, plan): - E, K, N = weight.shape - if not ( - _GROUPED_MM_TRITON_OP is not None - and inputs.device.type == "cuda" - and inputs.dtype in (torch.bfloat16, torch.float16) - and weight.dtype == inputs.dtype - and _triton_grouped_mm_rows_ok(inputs.device, plan.num_rows, plan.num_experts, K, N) - ): - return False - from unsloth_zoo.temporary_patches import moe_grouped_fp16 as mg - return mg.triton_grouped_available(inputs.device) - - -def _grouped_mm_slice(inputs, weight, offsets, plan): - """_grouped_mm_with_backward_fix for a slice of a planned problem: the Triton choice and - its tiles follow the whole (rows, experts) problem, so each group's result matches it.""" - if _slice_takes_triton(inputs, weight, plan): - from unsloth_zoo.temporary_patches import moe_grouped_fp16 as mg - return mg.generic_grouped_mm( - inputs.contiguous(), weight, offsets, config_rows = (plan.num_rows, plan.num_experts), - ) - return _grouped_mm_eager(inputs, weight, offsets) - - -def _slice_reads_views(inputs, weight, plan): - # Triton and the sm90 / sm100 grouped GEMMs (behind the view probe) give the copy's bits. A per-expert cuBLAS - # loop (torch._grouped_mm elsewhere, our fallback) does not: a transposed operand takes another kernel (T4 fp16, K 5760: 3x max error). - if _slice_takes_triton(inputs, weight, plan): - return True - device = inputs.device - if device.type != "cuda" or torch.version.hip is not None: - return False - index = device.index if device.index is not None else torch.cuda.current_device() - return _triton_grouped_mm_capability(index)[0] in (9, 10) and _grouped_mm_supported_on(device) - - -def _transposed_for_slice(weight, inputs, plan): - weight = weight.transpose(-2, -1) - return weight if _slice_reads_views(inputs, weight, plan) else weight.contiguous() - - -def _selective_grouped_mm(inputs, weight_provider, plan, transpose): - if len(plan.chunks) == 1: - experts, _, _, offsets = plan.chunks[0] - weight = weight_provider(experts) - if transpose: - weight = _transposed_for_slice(weight, inputs, plan) - return _grouped_mm_slice(inputs, weight, offsets, plan) - out = None - for experts, row_lo, row_hi, offsets in plan.chunks: - weight = weight_provider(experts) - if transpose: - weight = _transposed_for_slice(weight, inputs, plan) - rows = inputs[row_lo : row_hi] - if rows.data_ptr() % 16: - rows = rows.clone() # grouped GEMMs want 16-byte aligned operands - part = _grouped_mm_slice(rows, weight, offsets, plan) - del weight - if out is None: - out = inputs.new_empty((inputs.shape[0], part.shape[-1])) - out[row_lo : row_hi] = part - del part - return out - - -class _GroupedMMSelectiveRecompute(torch.autograd.Function): - """_GroupedMMRecompute over a _SelectiveDequantPlan: forward and dX read only the planned experts.""" - - @staticmethod - def forward(ctx, inputs, weight_provider, plan): - ctx.weight_provider = weight_provider - ctx.plan = plan - with torch.no_grad(): - return _selective_grouped_mm(inputs, weight_provider, plan, False) - - @staticmethod - def backward(ctx, grad_output): - with torch.no_grad(): - grad_input = _selective_grouped_mm(grad_output.contiguous(), ctx.weight_provider, ctx.plan, True) - return grad_input, None, None - - -def _mxfp4_expert_layout(source, proj_type, hidden_dim, model_type, experts_module): - """(packed param, transpose_b: weight is P^T of the packed rows) for a frozen packed MXFP4 stack, else None.""" - param = source - while hasattr(param, "base_layer"): - param = param.base_layer - if not is_mxfp4_expert_param(param) or param.requires_grad: - return None - cache = experts_module.__dict__.setdefault("_unsloth_mxfp4_layout", {}) - key = (proj_type, hidden_dim, model_type, bool(param.mxfp4_transposed), tuple(param._original_shape)) - if key not in cache: - # preprocess_weight on a meta stand-in: identity or a transpose view, anything else is unsupported. - meta = torch.empty(tuple(param._original_shape), device = "meta") - try: - weight = preprocess_weight(meta, proj_type, hidden_dim, model_type, experts_module = experts_module) - except Exception: - weight = None - flipped = meta.transpose(-2, -1) - if weight is meta or (weight is not None and weight.shape == meta.shape and weight.stride() == meta.stride()): - cache[key] = bool(param.mxfp4_transposed) - elif weight is not None and weight.shape == flipped.shape and weight.stride() == flipped.stride(): - cache[key] = not param.mxfp4_transposed - else: - cache[key] = None - transpose_b = cache[key] - return None if transpose_b is None else (param, transpose_b) - - -def _mxfp4_fused_enabled(param, dtype, rows = None) -> bool: - """UNSLOTH_MXFP4_FUSED_GEMM: "auto" fused below UNSLOTH_MXFP4_FUSED_MAX_ROWS (192) rows/expert (above, dequant + - cuBLAS wins on B200); "1" always, "dequant" never, "0" never plus the legacy decode-slot path.""" - if dtype != torch.bfloat16 or param.device.type != "cuda": - return False - mode = os.environ.get("UNSLOTH_MXFP4_FUSED_GEMM", "auto") - if mode in ("0", "dequant"): - return False - if mode != "1" and rows is not None: - limit = int(os.environ.get("UNSLOTH_MXFP4_FUSED_MAX_ROWS", "192")) - if rows >= limit * param.shape[0]: - return False - from unsloth_zoo.mxfp4_gemm import mxfp4_grouped_mm_available - return mxfp4_grouped_mm_available(param.device) - - -def _mxfp4_base_grouped_mm(inputs, offsets, counts, weight_provider, recompute, layout): - """Fused decode-in-GEMM, else dequantize with a transposed decode for backward.""" - param, transpose_b = layout - scales = param.mxfp4_scales - if scales.device != param.device: - scales = param.mxfp4_scales = scales.to(param.device) - if ( - inputs.dtype == torch.float32 and inputs.device.type == "cuda" - and torch.is_autocast_enabled("cuda") and torch.get_autocast_dtype("cuda") == torch.bfloat16 - ): - inputs = inputs.to(torch.bfloat16) # what autocast does to torch._grouped_mm's operands - if _mxfp4_fused_enabled(param, inputs.dtype, inputs.shape[0]): - from unsloth_zoo.mxfp4_gemm import mxfp4_expert_grouped_mm - return mxfp4_expert_grouped_mm(inputs, param.data, scales, counts, transpose_b = transpose_b) - if recompute and os.environ.get("UNSLOTH_MXFP4_FUSED_GEMM") != "0": - weight_provider.transposed = lambda: param.dequantize( - inputs.dtype, token_counts = counts, transpose = not transpose_b, - ) - return _base_grouped_mm(inputs, offsets, weight_provider, recompute) - - -_GROUPED_GEMM_AVAILABLE = None -_TORCH_GROUPED_MM_AVAILABLE = hasattr(torch, "_grouped_mm") - -# GPU support for torch._grouped_mm, verified via runtime probe. -_TORCH_GROUPED_MM_SUPPORTED = None - - -@torch.compiler.disable -def _run_probe_eagerly(probe): - """Graph-break so lazy probes never run on FakeTensors, whose meta check can reject what the real - kernel accepts and latch the cached flag False (torch 2.14 `_grouped_mm_fp16_cublaslt_supported` needs CUDA 13.3+).""" - return probe() - - -def _grouped_mm_probe_device(): - # Typed device, not a bare index: an int resolves to the default accelerator, losing this branch. - if torch.cuda.is_available(): - return torch.device("cuda", torch.cuda.current_device()) - if hasattr(torch, "xpu") and torch.xpu.is_available(): - return torch.device("xpu", torch.xpu.current_device()) - return None - - -def _grouped_mm_device_kind(backend, index): - try: - p = backend.get_device_properties(index) - return (p.name, getattr(p, "major", None), getattr(p, "minor", None), getattr(p, "gcnArchName", None)) - except Exception: - return index - - -_GROUPED_MM_FOREIGN = {} # (type, index) -> None if same GPU model as the probe device, else its model -# (type, index) the process-wide flags were probed on; a later current-device switch must not move it. -_GROUPED_MM_HOME = None -_GROUPED_MM_HOME_LOCK = threading.Lock() -# Every visible accelerator is the home card's model: per-call gates read the process-wide flags. -_GROUPED_MM_HOMOGENEOUS = False - - -def _grouped_mm_home(): - """Card the process-wide flags describe; locked so concurrent first callers on different cards agree.""" - global _GROUPED_MM_HOME, _GROUPED_MM_HOMOGENEOUS - if _GROUPED_MM_HOME is None: - device = _grouped_mm_probe_device() - anchored = False - with _GROUPED_MM_HOME_LOCK: - if _GROUPED_MM_HOME is None and device is not None and device.index is not None: - _GROUPED_MM_HOME = (device.type, device.index) - _GROUPED_MM_FOREIGN.setdefault(_GROUPED_MM_HOME, None) - anchored = True - # Classify the other visible cards now (metadata only, no kernel probe): a traced first use of - # another card must not reach the compiler-disabled resolver (graph break / fullgraph error). - if anchored: - try: - count = getattr(torch, device.type).device_count() - except Exception: - count = 0 - for i in range(count): - if (device.type, i) not in _GROUPED_MM_FOREIGN: - _resolve_grouped_mm_foreign_kind(torch.device(device.type, i), (device.type, i)) - _GROUPED_MM_HOMOGENEOUS = _grouped_mm_all_cards_alike(device.type, count) - return torch.device(*_GROUPED_MM_HOME) if _GROUPED_MM_HOME is not None else _grouped_mm_probe_device() - - -def _grouped_mm_all_cards_alike(home_type, count): - """Every visible card of the home's type is its model, and no other accelerator type has cards.""" - if count < 1 or any(_GROUPED_MM_FOREIGN.get((home_type, i), 0) is not None for i in range(count)): - return False - other = "xpu" if home_type == "cuda" else "cuda" - try: - backend = getattr(torch, other, None) - return not (backend is not None and backend.is_available() and backend.device_count() > 0) - except Exception: - return False - - -def _grouped_mm_traced_kind(device): - # Trace-time lookup, never resolves: None = the home model (or not a GPU), False = never seen eagerly. - if device is None or device.type not in ("cuda", "xpu") or device.index is None: return None - return _GROUPED_MM_FOREIGN.get((device.type, device.index), False) - - -def _grouped_mm_foreign_kind(device): - if device is None or device.type not in ("cuda", "xpu") or device.index is None: return None - key = (device.type, device.index) - if key in _GROUPED_MM_FOREIGN: return _GROUPED_MM_FOREIGN[key] - return _run_probe_eagerly(lambda: _resolve_grouped_mm_foreign_kind(device, key)) - - -def _resolve_grouped_mm_foreign_kind(device, key): - home = _grouped_mm_home() - kind = None - if home is not None and (home.type, home.index) != (device.type, device.index): - # Typed: a CUDA home must not answer for an XPU card (or the reverse) on a host with both. - kind = (device.type, _grouped_mm_device_kind(getattr(torch, device.type), device.index)) - if home.type == device.type and kind == (home.type, _grouped_mm_device_kind(getattr(torch, home.type), home.index)): - kind = None - _GROUPED_MM_FOREIGN[key] = kind - return kind - - -# Per GPU model other than the home card's. -_TORCH_GROUPED_MM_SUPPORTED_BY_KIND = {} - - -def _check_torch_grouped_mm_supported(device = None): - """Check torch._grouped_mm support on `device` (default: current); a runtime probe is the only reliable check.""" - global _TORCH_GROUPED_MM_SUPPORTED - if device is not None: - kind = _grouped_mm_foreign_kind(device) - if kind is not None: - supported = _TORCH_GROUPED_MM_SUPPORTED_BY_KIND.get(kind) - if supported is not None: return supported - if not _TORCH_GROUPED_MM_AVAILABLE: return False - return _run_probe_eagerly(lambda: _probe_torch_grouped_mm_supported(device)) - if _TORCH_GROUPED_MM_SUPPORTED is not None: return _TORCH_GROUPED_MM_SUPPORTED - # Definitive negatives stay traceable so fullgraph traces don't break. - if not _TORCH_GROUPED_MM_AVAILABLE or _grouped_mm_probe_device() is None: - _TORCH_GROUPED_MM_SUPPORTED = False - return False - return _run_probe_eagerly(_probe_torch_grouped_mm_supported) - - -def _torch_version_tuple(): - try: - return tuple(int(x) for x in torch.__version__.split("+")[0].split(".")[:2]) - except Exception: - return (0, 0) - - -# CUDA (not ROCm) torch 2.8 hard-raises unless `dprops->major == 9` (Blas.cpp, sm90_only). 2.9+ sends -# any other arch to _grouped_mm_fallback instead of raising (GroupedBlas.cpp, `_grouped_mm_cuda`). -_TORCH_IS_CUDA_BUILD = getattr(torch.version, "hip", None) is None -_TORCH_GROUPED_MM_SM90_ONLY = _TORCH_IS_CUDA_BUILD and _torch_version_tuple() == (2, 8) -_TORCH_GROUPED_MM_FALLS_BACK = _TORCH_IS_CUDA_BUILD and _torch_version_tuple() >= (2, 9) - - -def _grouped_mm_falls_back_internally(device_type, kind): - """torch itself serves grouped_mm on a card no probe reached: CUDA 2.9+ (natively or through its - fallback), 2.8 only on sm90. ROCm (some cards crash, unsloth#12391) and XPU make no such promise.""" - if device_type != "cuda": return False - if _TORCH_GROUPED_MM_FALLS_BACK: return True - info = kind[1] if isinstance(kind, tuple) and len(kind) == 2 else None - return _TORCH_GROUPED_MM_SM90_ONLY and isinstance(info, tuple) and info[1] == 9 - - -def _grouped_mm_supported_on(device): - # Absent symbol (torch 2.6/2.7): the traceable definitive negative, no classification. - if _GROUPED_MM_HOMOGENEOUS or not _TORCH_GROUPED_MM_AVAILABLE: - return _check_torch_grouped_mm_supported() - if torch.compiler.is_compiling(): - return _grouped_mm_supported_while_tracing(device) - kind = _grouped_mm_foreign_kind(device) - if kind is None: - return _check_torch_grouped_mm_supported() - return _check_torch_grouped_mm_supported(device) - - -def _grouped_mm_supported_while_tracing(device): - # Never probes (graph break; fullgraph error). An unwarmed card takes the home answer only where torch - # falls back internally, else the manual loop (safe, but breaks a fullgraph trace). - home = _check_torch_grouped_mm_supported() - kind = _grouped_mm_traced_kind(device) - if kind is None or _GROUPED_MM_HOMOGENEOUS: return home - known = None if kind is False else _TORCH_GROUPED_MM_SUPPORTED_BY_KIND.get(kind) - if known is not None: return known - return home and _grouped_mm_falls_back_internally(device.type, kind) - - -# Some ROCm GPUs segfault in grouped_mm instead of raising (unsloth#12391): run the kernels in a throwaway interpreter first. -_GROUPED_MM_CRASH_PROBE = r""" -import sys -if sys.platform == "win32": # no "python.exe has stopped working" dialog when the kernel faults - import ctypes; ctypes.windll.kernel32.SetErrorMode(0x0003) -import torch -d = torch.device("cuda", int(sys.argv[1])) -def probe(x, w, offs): - offs = torch.tensor(offs, device = d, dtype = torch.int32) - try: torch._grouped_mm(x, w, offs = offs); torch.cuda.synchronize(d) - except Exception: pass # a Python error is caught in-process too; only a crash or hang matters here -half = dict(device = d, dtype = torch.float16) -bf16 = dict(device = d, dtype = torch.bfloat16) -w_t = torch.ones(4, 64, 32, **bf16).transpose(-2, -1) -probe(torch.ones(1, 8, **half), torch.ones(1, 8, 8, **half), [1]) -probe(torch.ones(32, 32, **bf16), w_t, [8, 16, 24, 32]) -probe(torch.ones(32, 32, **bf16), w_t.contiguous(), [8, 16, 24, 32]) -""" -_GROUPED_MM_SURVIVES = None -_GROUPED_MM_SURVIVES_BY_KIND = {} -_GROUPED_MM_SURVIVES_LOCK = threading.Lock() - - -def _grouped_mm_survives_out_of_process(device): - global _GROUPED_MM_SURVIVES - kind = _grouped_mm_foreign_kind(device) - if kind is not None: - with _GROUPED_MM_SURVIVES_LOCK: - if kind not in _GROUPED_MM_SURVIVES_BY_KIND: - _GROUPED_MM_SURVIVES_BY_KIND[kind] = _run_grouped_mm_crash_probe(device) - return _GROUPED_MM_SURVIVES_BY_KIND[kind] - if _GROUPED_MM_SURVIVES is not None: return _GROUPED_MM_SURVIVES - # Publish once known: concurrent first callers wait for the child, never run the kernel. - with _GROUPED_MM_SURVIVES_LOCK: - if _GROUPED_MM_SURVIVES is None: _GROUPED_MM_SURVIVES = _run_grouped_mm_crash_probe(device) - return _GROUPED_MM_SURVIVES - - -def _run_grouped_mm_crash_probe(device): - # AMD only: CUDA and XPU keep the in-process probe. - if getattr(torch.version, "hip", None) is None or device.type != "cuda": return True - # Studio's Python _grouped_mm fallback (gfx120X) is active: this process never runs the native kernel. - try: - dump = torch._C._dispatch_dump("aten::_grouped_mm").splitlines() - if any(l.startswith("CUDA:") and ".cpp:" not in l for l in dump): return True - except Exception: - pass - import subprocess - index, env = device.index, None - # ROCr / ordinal selectors compose with the HIP mask unmirrorably: keep the parent's env and index. - layered = any((os.environ.get(var) or "").strip() for var in ("ROCR_VISIBLE_DEVICES", "GPU_DEVICE_ORDINAL")) - if torch.cuda.device_count() > 1 and not layered: - # Child sees only this card; HIP indexes into the inherited mask. - mask = os.environ.get("HIP_VISIBLE_DEVICES") or os.environ.get("CUDA_VISIBLE_DEVICES") - mask = (mask or "").strip() - if len(mask) >= 2 and mask[0] == mask[-1] == '"': mask = mask[1:-1] - # clr keeps the first of a repeated entry. - mask = list(dict.fromkeys(m.strip() for m in mask.split(","))) if mask else None - # Only plain ordinals map token -> logical index (UUIDs alias them); PAL (Windows) enumerates in physical order. - plain = mask is None or ( - all(m.isdigit() for m in mask) and (sys.platform != "win32" or mask == sorted(mask, key = int)) - ) - if plain and (mask is None or index < len(mask)): - visible = mask[index] if mask else str(index) - index, env = 0, dict(os.environ, HIP_VISIBLE_DEVICES = visible, CUDA_VISIBLE_DEVICES = visible) - # Not multiprocessing: spawn re-imports the caller's __main__. - try: - result = subprocess.run( - [sys.executable, "-c", _GROUPED_MM_CRASH_PROBE, str(index)], env = env, - stdout = subprocess.DEVNULL, stderr = subprocess.DEVNULL, timeout = 180, - creationflags = getattr(subprocess, "CREATE_NO_WINDOW", 0), - ).returncode - except Exception as e: - result = type(e).__name__ - if result == 0: return True - _log_info(f"Unsloth: torch._grouped_mm probe crashed on ROCm GPU {device.index} ({result}); disabling grouped_mm.") - return False - - -def _grouped_mm_runs_on(device): - try: - # Dummy call verifies real support (symbol may exist but hardware unsupported, e.g. < H100). - dtype = torch.float16 - - # 1 expert, 1 token, dim 8 (safe alignment). - x = torch.ones((1, 8), device=device, dtype=dtype) - w = torch.ones((1, 8, 8), device=device, dtype=dtype) - offs = torch.tensor([1], device=device, dtype=torch.int32) - - torch._grouped_mm(x, w, offs=offs) - del x, w, offs - return True - except Exception: - return False - - -def _probe_torch_grouped_mm_supported(device = None): - global _TORCH_GROUPED_MM_SUPPORTED - kind = _grouped_mm_foreign_kind(device) - if kind is not None: - supported = _TORCH_GROUPED_MM_SUPPORTED_BY_KIND.get(kind) - if supported is None: - supported = _TORCH_GROUPED_MM_AVAILABLE and _grouped_mm_survives_out_of_process(device) and _grouped_mm_runs_on(device) - _TORCH_GROUPED_MM_SUPPORTED_BY_KIND[kind] = supported - # Same eager region: a traced first use of this model must not break again on the view check. - if kind not in _TRANSPOSED_VIEW_SAFE_BY_KIND: - _TRANSPOSED_VIEW_SAFE_BY_KIND[kind] = supported and _transposed_view_probe(device) - return supported - if _TORCH_GROUPED_MM_SUPPORTED is not None: return _TORCH_GROUPED_MM_SUPPORTED - - # The anchored card, not whichever is current now: the flags must describe one card. - device = _grouped_mm_home() if _TORCH_GROUPED_MM_AVAILABLE else None - if device is None or not _grouped_mm_survives_out_of_process(device): - _TORCH_GROUPED_MM_SUPPORTED = False - else: - _TORCH_GROUPED_MM_SUPPORTED = _grouped_mm_runs_on(device) - return _TORCH_GROUPED_MM_SUPPORTED - - -def _warm_grouped_mm_devices(devices): - """Mixed host only, eagerly: probe one card of each other GPU model among `devices` (the cards - holding this process's expert weights), so a later trace there reads its own answers. Never a - card the process does not use: torchrun ranks see every GPU, and probing foreign ones would put - a CUDA context and a grouped_mm call from every rank on one representative card.""" - if _GROUPED_MM_HOMOGENEOUS or torch.compiler.is_compiling(): return - seen = set() - for device in devices: - try: - if device is None or device.type not in ("cuda", "xpu") or device.index is None: continue - kind = _grouped_mm_foreign_kind(device) - if kind is None or kind in seen: continue - seen.add(kind) - _check_torch_grouped_mm_supported(device) - except Exception: - pass - - -# One-time check that torch._grouped_mm on a transposed bf16 view matches the contiguous copy. Insurance -# only: pytorch/pytorch#186365 was a usage error (offs must be cumulative group ends), no view miscompute is known. -_TRANSPOSED_VIEW_GROUPED_MM_SAFE = None - - -def _transposed_view_grouped_mm_is_safe(): - global _TRANSPOSED_VIEW_GROUPED_MM_SAFE - if _TRANSPOSED_VIEW_GROUPED_MM_SAFE is not None: - return _TRANSPOSED_VIEW_GROUPED_MM_SAFE - if not _TORCH_GROUPED_MM_AVAILABLE or _grouped_mm_probe_device() is None: - _TRANSPOSED_VIEW_GROUPED_MM_SAFE = False - return False - return _run_probe_eagerly(_probe_transposed_view_grouped_mm_is_safe) - - -_TRANSPOSED_VIEW_SAFE_BY_KIND = {} - - -def _transposed_view_safe_on(device): - """View probe on `device`'s GPU model; the home model keeps the process-wide flag.""" - if _GROUPED_MM_HOMOGENEOUS or not _TORCH_GROUPED_MM_AVAILABLE: return _transposed_view_grouped_mm_is_safe() - if torch.compiler.is_compiling(): - kind = _grouped_mm_traced_kind(device) - if kind is None: return _transposed_view_grouped_mm_is_safe() - # Unprobed card: the contiguous copy (correct, traceable), never the home card's verdict. - return False if kind is False else bool(_TRANSPOSED_VIEW_SAFE_BY_KIND.get(kind, False)) - kind = _grouped_mm_foreign_kind(device) - if kind is None: return _transposed_view_grouped_mm_is_safe() - safe = _TRANSPOSED_VIEW_SAFE_BY_KIND.get(kind) - if safe is not None: return safe - return _run_probe_eagerly(lambda: _probe_transposed_view_grouped_mm_is_safe(device)) - - -def _probe_transposed_view_grouped_mm_is_safe(device = None): - global _TRANSPOSED_VIEW_GROUPED_MM_SAFE - kind = _grouped_mm_foreign_kind(device) - if kind is not None: - if kind not in _TRANSPOSED_VIEW_SAFE_BY_KIND: - _TRANSPOSED_VIEW_SAFE_BY_KIND[kind] = _transposed_view_probe(device) - return _TRANSPOSED_VIEW_SAFE_BY_KIND[kind] - if _TRANSPOSED_VIEW_GROUPED_MM_SAFE is not None: - return _TRANSPOSED_VIEW_GROUPED_MM_SAFE - _TRANSPOSED_VIEW_GROUPED_MM_SAFE = _transposed_view_probe(_grouped_mm_home()) - return _TRANSPOSED_VIEW_GROUPED_MM_SAFE - - -def _transposed_view_probe(device): - safe = False - try: - if _TORCH_GROUPED_MM_AVAILABLE and device is not None and _grouped_mm_survives_out_of_process(device): - E, N, K, M = 4, 64, 32, 32 - # local generator: never touch the process-wide RNG (manual_seed would shift training) - gen = torch.Generator(device=device).manual_seed(0) - A = torch.randn(M, K, dtype=torch.bfloat16, device=device, generator=gen) - w = torch.randn(E, N, K, dtype=torch.bfloat16, device=device, generator=gen) - w_t = w.transpose(-2, -1) - w_tc = w_t.contiguous() - per = M // E - offs = torch.arange(per, M + 1, per, dtype=torch.int32, device=device)[:E] - offs[-1] = M - ok, ref = True, None - for _ in range(6): - row_wise_max = A.abs().amax(dim=-1, keepdim=True) - _ = A / (row_wise_max / 448.0) # allocator churn from the #186365 report - r_view = torch._grouped_mm(A, w_t, offs=offs) - r_contig = torch._grouped_mm(A, w_tc, offs=offs) - if (r_view - r_contig).abs().max().item() > 1e-2: # view disagrees with contiguous - ok = False; break - if ref is None: - ref = r_view - elif (r_view - ref).abs().max().item() > 1e-2: # view not stable across calls - ok = False; break - safe = ok - except Exception: - safe = False # anything unexpected -> keep the safe contiguous copy - return safe - - -_TRITON_ALLOCATOR_INITIALIZED = False -_PERSISTENT_BUFFER = {} # per device index: a card must never be handed another card's scratch -_original_peft_get_peft_model = None - - -def _init_triton_allocator(): - """Initialize a persistent Triton allocator to avoid per-call allocation overhead.""" - global _TRITON_ALLOCATOR_INITIALIZED - if _TRITON_ALLOCATOR_INITIALIZED: return - - try: - import triton - - # Persistent buffer that grows as needed, avoiding per-kernel allocations. - def persistent_alloc_fn(size: int, alignment: int, stream): - # Triton launches on, and asks for scratch on, the current device. - device = torch.cuda.current_device() - buffer = _PERSISTENT_BUFFER.get(device) - # Round up to nearest 128 bytes for alignment / fewer reallocations. - rounded_size = ((size + 128 - 1) // 128) * 128 - - if ( - buffer is None - or buffer.numel() * buffer.element_size() - < rounded_size - ): - # 10% headroom; uint8 for raw byte storage. - buffer = _PERSISTENT_BUFFER[device] = torch.empty( - int(rounded_size * 1.1), device=torch.device("cuda", device), dtype=torch.uint8 - ) - buffer.__hibernate__ = {"type": "ignore"} - return buffer - - triton.set_allocator(persistent_alloc_fn) - triton._unsloth_allocator_set = True - _TRITON_ALLOCATOR_INITIALIZED = True - except Exception: - pass - - -def _check_grouped_gemm_available(): - """Check if Unsloth grouped GEMM kernels are available.""" - if os.environ.get("UNSLOTH_DISABLE_MOE_TRITON", "0") == "1": return False - if is_mlx_available(): return False - - global _GROUPED_GEMM_AVAILABLE - if _GROUPED_GEMM_AVAILABLE is not None: return _GROUPED_GEMM_AVAILABLE - - # The kernel asserts `X.device.type == "cuda"` on entry, so importable is not the - # same as usable: on a CPU box with unsloth installed this used to answer yes, - # select_moe_backend picked "unsloth_triton", and the first expert forward died on - # "X and W must be on CUDA" instead of taking the native loop. Asked here rather - # than at the call site, the way _check_torch_grouped_mm_supported already refuses - # without an accelerator. - if not torch.cuda.is_available(): - _GROUPED_GEMM_AVAILABLE = False - return False - - # Triton 3.3 (torch 2.7) lacks tl.make_tensor_descriptor, which the JIT resolves while hashing: use the native loop. - try: - import triton.language as tl - if not hasattr(tl, "make_tensor_descriptor"): - _GROUPED_GEMM_AVAILABLE = False - return False - except Exception: - _GROUPED_GEMM_AVAILABLE = False - return False - - try: - from unsloth.kernels.moe.grouped_gemm.interface import grouped_gemm, supports_tma - _GROUPED_GEMM_AVAILABLE = True - _init_triton_allocator() - except (ImportError, ModuleNotFoundError): - # A temporary patch asking while unsloth itself imports: circular, not absent. Decide later. - if _unsloth_importing(): - return False - _GROUPED_GEMM_AVAILABLE = False - return _GROUPED_GEMM_AVAILABLE - - -from functools import lru_cache, wraps - - -def _unsloth_importing(): - return bool(getattr(getattr(sys.modules.get("unsloth"), "__spec__", None), "_initializing", False)) - - -_SELECTED_MOE_BACKEND = None - - -def select_moe_backend(): - """Select MoE backend from UNSLOTH_MOE_BACKEND + availability. - - Choices: "grouped_mm", "unsloth_triton", "native_torch" (default "grouped_mm"). - """ - # This Unsloth Zoo code section is licensed under AGPL3 - global _SELECTED_MOE_BACKEND - if _SELECTED_MOE_BACKEND is not None: - return _SELECTED_MOE_BACKEND - backend = _select_moe_backend() - # Chosen while unsloth imports (its Triton kernels cannot load yet): choose again on the next call. - # Not inside a trace either: the write fails the `is None` guard and recompiles the caller once. - if not _unsloth_importing() and not torch.compiler.is_compiling(): - _SELECTED_MOE_BACKEND = backend - return backend - - -def _clear_selected_moe_backend(): - global _SELECTED_MOE_BACKEND - _SELECTED_MOE_BACKEND = None - - -select_moe_backend.cache_clear = _clear_selected_moe_backend - - -def select_moe_backend_for_interleaved_gate_up(): - """select_moe_backend() for interleaved gate_up (GPT-OSS), which Triton cannot run: grouped_mm, not the loop.""" - # This Unsloth Zoo code section is licensed under AGPL3 - backend = select_moe_backend() - if backend in ("unsloth_triton", "sonic_moe") and _check_torch_grouped_mm_supported(): - return "grouped_mm" - return backend - - -def select_moe_backend_without_sonic(device = None): - """select_moe_backend() for dispatchers with no sonic route (bnb4bit / FP8 experts): sonic_moe runs as grouped_mm.""" - # This Unsloth Zoo code section is licensed under AGPL3 - backend = select_moe_backend() - if backend == "sonic_moe": - # sonic_moe was selected on a card with torch._grouped_mm; on a mixed host the expert's card may lack it. - if getattr(device, "type", None) == "cuda" and not _grouped_mm_supported_on(device): - return "native_torch" - return "grouped_mm" - return backend - - -def _check_sonic_moe_available(): - """Opt-in kernels-community/sonic-moe backend; guarded import.""" - if os.environ.get("UNSLOTH_DISABLE_MOE_SONIC", "0") == "1": - return False - try: - from unsloth_zoo.temporary_patches.moe_utils_sonic import sonic_moe_available - # Probes must not run under Dynamo tracing. - return _run_probe_eagerly(sonic_moe_available) - except Exception: - return False - - -def _select_moe_backend(): - requested = os.environ.get("UNSLOTH_MOE_BACKEND") - if requested: - if requested == "grouped_mm" and _check_torch_grouped_mm_supported(): - return "grouped_mm" - if requested == "unsloth_triton" and _check_grouped_gemm_available(): - return "unsloth_triton" - if requested == "native_torch": - return "native_torch" - # Opt-in only; still needs torch._grouped_mm for LoRA GEMMs and fallbacks. - if ( - requested == "sonic_moe" - and _check_torch_grouped_mm_supported() - and _check_sonic_moe_available() - ): - _log_info("Unsloth: Using MoE backend 'sonic_moe'") - return "sonic_moe" - _log_info(f"Unsloth: '{requested}' backend requested but is not available. Falling back to next available.") - - if _check_torch_grouped_mm_supported(): - _log_info("Unsloth: Using MoE backend 'grouped_mm'") - return "grouped_mm" - if _check_grouped_gemm_available(): - _log_info("Unsloth: Using MoE backend 'unsloth_triton'") - return "unsloth_triton" - return "native_torch" - - -def moe_compute_dtype(hidden_states): - """W8A16 / W4A16 compute dtype: a half activation's own, else autocast's, else bf16 (or fp16). Never float32.""" - dtype = hidden_states.dtype - if dtype in (torch.bfloat16, torch.float16): - return dtype - device_type = hidden_states.device.type - try: - if device_type != "meta" and torch.amp.is_autocast_available(device_type) and torch.is_autocast_enabled(device_type): - autocast_dtype = torch.get_autocast_dtype(device_type) - if autocast_dtype in (torch.bfloat16, torch.float16): - return autocast_dtype - except (AttributeError, RuntimeError, TypeError): - pass - # The activation's own device; is_bf16_supported() reads the current one and counts sm < 80 emulation. - if ( - device_type == "cuda" - and torch.version.hip is None - and torch.cuda.get_device_capability(hidden_states.device)[0] < 8 - ): - return torch.float16 - return torch.bfloat16 - - -def swap_moe_weights_for_call(experts_module, gate_up_proj, down_proj, forward_fn, *args): - """Temporarily install dequantized weights for one forward call, then restore. - - Registered parameters are swapped in place in `_parameters` (no nn.Module - __setattr__, so no (de)registration or hooks); any other attribute goes through - object.__setattr__. Used by the FP8 and bnb4bit MoE dispatchers. - """ - # Swap `_parameters` entries only: a __dict__ write shadows later FSDP2 swaps (freed buffer), and - # `__dict__.pop` / `del` / an `object()` sentinel graph-break; dict setitem traces on every torch. - names = ("gate_up_proj", "down_proj") - params = experts_module._parameters - instance_dict = experts_module.__dict__ - registered = {name: params[name] for name in names if name in params} - saved = {name: instance_dict[name] for name in names if name not in registered and name in instance_dict} - for name, tensor in zip(names, (gate_up_proj, down_proj)): - if name in registered: - params[name] = tensor - else: - object.__setattr__(experts_module, name, tensor) - try: - return forward_fn(experts_module, *args) - finally: - for name in names: - if name in registered: - params[name] = registered[name] - elif name in saved: - instance_dict[name] = saved[name] - else: - instance_dict.pop(name, None) - - -def forward_moe_backend( - self, - hidden_states: torch.Tensor, - top_k_index: torch.Tensor, - top_k_weights: torch.Tensor, -) -> torch.Tensor: - """Dispatch MoE forward to the selected backend (keeps model-specific patches minimal).""" - # This Unsloth Zoo code section is licensed under AGPL3 - - # Absolute imports: this function is also copied into - # unsloth_compiled_cache/moe_utils.py where relative imports of sibling - # helpers don't resolve (only the dispatcher is copied). - # Keep `except ImportError` around ONLY the import; runtime errors in the - # bnb4bit/fp8 path must propagate, not fall through to a crashing backend. - _moe_uses_bnb4bit_expert_weights = forward_moe_backend_bnb4bit = None - try: - from unsloth_zoo.temporary_patches.moe_utils_bnb4bit import ( - _moe_uses_bnb4bit_expert_weights, - forward_moe_backend_bnb4bit, - ) - except ImportError: - pass - if _moe_uses_bnb4bit_expert_weights is not None and _moe_uses_bnb4bit_expert_weights(self): - result = forward_moe_backend_bnb4bit(self, hidden_states, top_k_index, top_k_weights) - if result is not None: - return result - - _moe_uses_fp8_expert_weights = forward_moe_backend_fp8 = None - try: - from unsloth_zoo.temporary_patches.moe_utils_fp8 import ( - _moe_uses_fp8_expert_weights, - forward_moe_backend_fp8, - ) - except ImportError: - pass - if _moe_uses_fp8_expert_weights is not None and _moe_uses_fp8_expert_weights(self): - return forward_moe_backend_fp8(self, hidden_states, top_k_index, top_k_weights) - - # BF16 / FP16 decode-sized no-grad calls: only the routed experts (moe_routed.py). - routed_moe_forward = None - try: - from unsloth_zoo.temporary_patches.moe_routed import routed_moe_forward - except ImportError: - pass - if routed_moe_forward is not None: - result = routed_moe_forward(self, hidden_states, top_k_index, top_k_weights) - if result is not None: - return result - - backend = select_moe_backend() - if backend == "sonic_moe": - # Absolute import: this function is copied into unsloth_compiled_cache. - forward_sonic_moe = None - try: - from unsloth_zoo.temporary_patches.moe_utils_sonic import forward_sonic_moe - except ImportError: - pass - if forward_sonic_moe is not None: - return forward_sonic_moe(self, hidden_states, top_k_index, top_k_weights) - backend = "grouped_mm" - if backend == "grouped_mm": - return forward_native_grouped_mm(self, hidden_states, top_k_index, top_k_weights) - # Triton kernels assume split gate_up, SiLU, no bias: interleaved GPT-OSS gate_up takes the eager loop. - if backend == "unsloth_triton" and not _gate_up_is_interleaved(self): - return forward_triton_grouped_gemm(self, hidden_states, top_k_index, top_k_weights) - return forward_native_moe_loop(self, hidden_states, top_k_index, top_k_weights) - - -# Test-only call counter, wired at import time and OFF by default: incrementing a -# module global is a Dynamo side effect, and leaving it always-on re-traced the -# compiled MoE block every step (3.6 ms -> 572 ms). -_EXPERT_COUNT_CALLS = [0] -_COUNT_DEBUG = os.environ.get("UNSLOTH_MOE_COUNT_DEBUG", "0") == "1" - - -def _count_tokens_per_expert( - flat_experts: torch.Tensor, - num_experts: int, - dtype: torch.dtype = torch.int32, -) -> torch.Tensor: - """Per-expert token counts, WITHOUT synchronising the device. - - Drop-in replacement for ``torch.bincount(flat_experts, minlength=num_experts)`` - returning shape ``[num_experts]`` and the requested ``dtype``. - - Why not bincount: on CUDA the output size of ``bincount`` depends on the data - (it must learn ``max(input)`` even when ``minlength`` is given), so the kernel - D2H-copies and blocks the host. Measured on B200, E=256 / 16384 routed rows: - 2 sync warnings under ``torch.cuda.set_sync_debug_mode("warn")`` and ~58 us of - host time per call, versus 0 warnings and ~19 us here. With one MoE layer per - block that drain happens tens of times per step. - - ``scatter_add_`` accumulates with atomics, so the ORDER of accumulation is - nondeterministic; integer addition is associative and commutative and every - addend is exactly 1, so the SUM is bit-exact and run-to-run deterministic - regardless of order. Verified bit-identical to ``bincount`` over random, - skewed, empty and single-token routings, and it does not trip - ``torch.use_deterministic_algorithms(True)``. - - Caveat, same precondition as the caller already relies on: ``flat_experts`` - must hold values in ``[0, num_experts)``. That holds for router ``topk`` - output by construction. Out-of-range values make ``bincount`` silently return - a LONGER tensor (breaking the downstream ``offs=`` shape) whereas this writes - out of bounds, so neither is safe; the difference is only in how it fails. - """ - # This Unsloth Zoo code section is licensed under AGPL3 - if flat_experts.dim() != 1: - flat_experts = flat_experts.reshape(-1) - # scatter_add_ needs int64; router topk already gives one, so no copy. - index = flat_experts if flat_experts.dtype == torch.int64 else flat_experts.long() - counts = torch.zeros(num_experts, dtype=dtype, device=flat_experts.device) - counts.scatter_add_(0, index, torch.ones_like(index, dtype=dtype)) - return counts - - -if _COUNT_DEBUG: - def count_tokens_per_expert(flat_experts, num_experts, dtype=torch.int32): - _EXPERT_COUNT_CALLS[0] += 1 - return _count_tokens_per_expert(flat_experts, num_experts, dtype) - - count_tokens_per_expert.__doc__ = _count_tokens_per_expert.__doc__ -else: - count_tokens_per_expert = _count_tokens_per_expert - - -def moe_routing_sort(top_k_index, num_experts): - """(counts int32 [E], stable argsort int64 [N], cumsum(counts) int32 [E]), sync-free; Triton on CUDA.""" - # This Unsloth Zoo code section is licensed under AGPL3 - if _TRITON_ROUTING is not None: - return _TRITON_ROUTING(top_k_index, num_experts) - flat_experts = top_k_index.reshape(-1) - counts = count_tokens_per_expert(flat_experts, num_experts, torch.int32) - sorted_indices = flat_experts.argsort(stable=True) - return counts, sorted_indices, torch.cumsum(counts, dim=0, dtype=torch.int32) - - -try: - from unsloth_zoo.temporary_patches.moe_routing_triton import moe_routing_sort as _TRITON_ROUTING -except Exception: - _TRITON_ROUTING = None - - -@torch.no_grad() -def _get_routing_indices(selected_experts, num_experts): - """Compute token->expert mapping for grouped GEMM. - - Returns (token_counts_by_expert (num_experts,), gather_indices (total_tokens,)). - """ - # This Unsloth Zoo code section is licensed under AGPL3 - - token_counts_by_expert, gather_indices, _ = moe_routing_sort(selected_experts, num_experts) - return token_counts_by_expert, gather_indices - - -def combine_permuted_moe_outputs( - permuted_output: torch.Tensor, - sorted_indices: torch.Tensor, - num_tokens: int, - top_k: int, - out_dtype = None, -) -> torch.Tensor: - """Sum the top_k expert outputs belonging to each token, in a fixed order. - - The obvious spelling of this reduction is - ``zeros(num_tokens, hidden).index_add_(0, token_indices, permuted_output)`` - with ``token_indices = sorted_indices // top_k``. Because every token index - appears ``top_k`` times in one call, that lands in ``index_add_``'s CUDA - atomicAdd path, and atomicAdd fixes no accumulation order. Float addition is - not associative, so the result moves from run to run: repeating one identical - forward 40 times on a single already-loaded Gemma-4 MoE, with no optimizer - step and no data change, gave 40 distinct losses spread over 0.0284 nats, - while the same 40 repeats through the reduction below returned one value. - - ``sorted_indices`` is ``argsort`` of the flat expert assignment, so it is a - permutation of ``range(num_tokens * top_k)`` - every slot exactly once. That - means the permutation can simply be undone (a gather through the inverse - permutation: unique indices, no accumulation and so no atomics) and the - ``top_k`` axis reduced with an ordinary ``sum``, which has a fixed reduction - order. Same arithmetic, same values up to that ordering, reproducible. - - ``out_dtype`` is applied AFTER the reduction, not before it. When the routed - slots arrive wider than ``out_dtype`` - the fp32-router case, where the - routing-weight multiply promotes the expert output to fp32 - casting first - would round every one of the ``top_k`` summands to bf16 and only then add - them. Summing first and rounding once is what ``transformers`` - (``integrations/moe.py``) does, and it is measurably closer to an fp64 - reference. With a bf16 router the two orders are bit-identical. - """ - # This Unsloth Zoo code section is licensed under AGPL3 - - # The gather below is safe only because sorted_indices is an argsort, so the indices - # are unique and, given the count check here matches, cover every row: nothing is left - # unread and no slot is counted twice. A bare `assert` would vanish under `python -O`, - # which is exactly the build where a silently wrong reduction is hardest to notice. - if sorted_indices.numel() != permuted_output.shape[0]: - raise ValueError( - f"Unsloth: expected a full permutation of {permuted_output.shape[0]} routed slots, " - f"got {sorted_indices.numel()} indices" - ) - # inverse_indices[sorted_indices[i]] = i, i.e. "which routed slot holds original row j". - # Gathering with it is one [num_tokens * top_k, hidden] buffer, in the incoming dtype; - # the narrowing cast then only ever touches the small [num_tokens, hidden] result. - inverse_indices = torch.empty_like(sorted_indices) - inverse_indices.scatter_( - 0, - sorted_indices, - torch.arange( - sorted_indices.numel(), - device = sorted_indices.device, - dtype = sorted_indices.dtype, - ), - ) - combined = ( - permuted_output - .index_select(0, inverse_indices) - .view(num_tokens, top_k, permuted_output.shape[-1]) - .sum(dim = 1) - ) - if out_dtype is not None: - combined = combined.to(out_dtype) - return combined - - -class _RoutedRowGather(torch.autograd.Function): - """hidden_states[sorted_indices // top_k] with a fixed-order backward: Inductor lowers the - default index_put backward to tl.atomic_add (non-reproducible LoRA grads). An unrolled top_k - sum was miscompiled under dynamic shapes (torch 2.13), so do not use it.""" - # This Unsloth Zoo code section is licensed under AGPL3 - - @staticmethod - def forward(ctx, hidden_states, sorted_indices, top_k): - ctx.save_for_backward(sorted_indices) - ctx.num_tokens = hidden_states.shape[0] - ctx.top_k = top_k - return hidden_states.index_select(0, sorted_indices // top_k) - - @staticmethod - def backward(ctx, grad_rows): - (sorted_indices,) = ctx.saved_tensors - grad = combine_permuted_moe_outputs(grad_rows, sorted_indices, ctx.num_tokens, ctx.top_k) - return grad, None, None - - -# Default from top_k 8 (L4: -55% at k 8, +22% at k 4). Env =1: every top_k > 2, =0: never. -_DETERMINISTIC_GATHER = os.environ.get("UNSLOTH_MOE_DETERMINISTIC_GATHER", "") -_DETERMINISTIC_GATHER_MIN_TOP_K = {"1": 3, "0": 1 << 30}.get(_DETERMINISTIC_GATHER, 8) - - -def gather_routed_rows(hidden_states, sorted_indices, top_k): - """sorted_indices must be the full argsort from moe_routing_sort. Eager backward is already - reproducible, so it keeps plain indexing.""" - # This Unsloth Zoo code section is licensed under AGPL3 - if ( - top_k >= _DETERMINISTIC_GATHER_MIN_TOP_K - and torch.compiler.is_compiling() - and hidden_states.requires_grad - and torch.is_grad_enabled() - and sorted_indices.numel() == hidden_states.shape[0] * top_k - ): - return _RoutedRowGather.apply(hidden_states, sorted_indices, top_k) - return hidden_states[sorted_indices // top_k] - - -def _silu_and_mul(x): - """Fused SiLU + element-wise multiply for gate/up projections.""" - gate, up = x.chunk(2, dim=-1) - return F.silu(gate) * up - - -try: - from unsloth_zoo.temporary_patches import moe_fused_swiglu as _MOE_FUSED_SWIGLU -except Exception: - _MOE_FUSED_SWIGLU = None - - -def _act_fn_is_silu(act_fn) -> bool: - if act_fn is F.silu: - return True - # Exact types only: a subclass may override forward. - return type(act_fn) is torch.nn.SiLU or type(act_fn).__name__ == "SiLUActivation" - - -def _fused_moe_activation(module, gate_up, interleaved, allow_gptoss = True): - """Fused (clamped) SwiGLU over gate_up, or None so the caller runs its bit-identical eager path.""" - _sw = _MOE_FUSED_SWIGLU - if _sw is None or not _sw.fused_swiglu_supported(gate_up): - return None - # Compiled frames round once, matching Inductor's fusion of the eager expression. - eager_rounding = not torch.compiler.is_compiling() - alpha, limit = 1.702, 7.0 - if "GptOssExperts" in module.__class__.__name__: - if not allow_gptoss: - return None - act = _sw.ACT_GPTOSS - alpha, limit = float(getattr(module, "alpha", 1.702)), float(getattr(module, "limit", 7.0)) - else: - act_fn = getattr(module, "act_fn", None) - if act_fn is not None and callable(act_fn) and not _act_fn_is_silu(act_fn): - return None - act = _sw.ACT_SILU - if eager_rounding and not _sw.fused_swiglu_self_check(gate_up.device, gate_up.dtype, act, interleaved, alpha, limit): - return None - return _sw.moe_fused_swiglu(gate_up, act, interleaved, alpha, limit, eager_rounding) - - -# Separated LoRA helpers. - - -def _has_lora_adapters(param) -> bool: - """Check for active LoRA adapters (PEFT ParamWrapper).""" - if not hasattr(param, "lora_A") or not hasattr(param, "lora_B"): - return False - if hasattr(param, "disable_adapters") and param.disable_adapters: - return False - if hasattr(param, "merged") and param.merged: - return False - return len(param.lora_A) > 0 - - -# How the flat `num_experts * rank` axis of a fused expert `lora_B` is laid out. -# -# PEFT stores a fused MoE expert LoRA as two ordinary Linear weights, lora_A of shape -# (num_experts * rank, in) and lora_B of shape (out, num_experts * rank), and the two are -# NOT flattened the same way. lora_A is expert-slowest, `reshape(E, r, in)`. lora_B is -# expert-FASTEST: `reshape(out, r, E)`, so column j belongs to expert `j % E`. That -# asymmetry is easy to miss, and it is the whole of this constant's reason to exist. -# -# Expert-fastest is what PEFT's own forward (`ParamWrapper.get_delta_factors`) and its -# merge (`ParamWrapper.get_delta_weight`) both do, unchanged across PEFT 0.18 to 0.21; -# what vLLM's `_stack_moe_lora_weights` does when it serves the adapter; and what prime-rl -# writes. It is the only reading any consumer of a saved adapter implements, so it is what -# the separated forward has to train. -LORA_B_LAYOUT_RANK_MAJOR = "rank_major" - -# What Unsloth's separated forward read before this fix: expert-SLOWEST, a rank-wide -# contiguous block of columns per expert. Self-consistent, and understood by nothing else. -# Kept only so an adapter trained that way can still be read back, see -# `moe_lora_b_layout()`. -LORA_B_LAYOUT_GROUPED_BY_EXPERT = "grouped_by_expert" - -_LORA_B_LAYOUTS = (LORA_B_LAYOUT_RANK_MAJOR, LORA_B_LAYOUT_GROUPED_BY_EXPERT) - - -class MoELoRABLayoutError(ValueError): - """`UNSLOTH_MOE_LORA_B_LAYOUT` is set to something that is not a layout. - - Its own class because `_extract_lora_from_wrapper` turns every exception into "this - wrapper has no LoRA", which would silently drop the adapter and train the base model. - A configuration mistake has to be louder than that, so it is re-raised there by type. - """ - - -def _read_moe_lora_b_layout_from_env() -> str: - """`UNSLOTH_MOE_LORA_B_LAYOUT`, validated. Split out so the value can be refreshed into - a module global on every eager call, which is what makes the switch visible to a - compiled graph.""" - # This Unsloth Zoo code section is licensed under AGPL3 - - layout = os.environ.get("UNSLOTH_MOE_LORA_B_LAYOUT", LORA_B_LAYOUT_RANK_MAJOR) - if layout not in _LORA_B_LAYOUTS: - raise MoELoRABLayoutError( - f"Unsloth: UNSLOTH_MOE_LORA_B_LAYOUT must be one of {_LORA_B_LAYOUTS}, got " - f"{layout!r}." - ) - return layout - - -_MOE_LORA_B_LAYOUT = _read_moe_lora_b_layout_from_env() - - -def moe_lora_b_layout() -> str: - """Which packing to read a fused expert `lora_B` with. - - `rank_major` (PEFT's, and therefore everyone's) unless - `UNSLOTH_MOE_LORA_B_LAYOUT=grouped_by_expert` asks for the pre-fix reading, which is - what an adapter trained by an older Unsloth needs. - - It is a process-wide mode, not a load-time option. Nothing is recorded on the adapter, - and every forward and every merge resolves it again, so setting it only around - `load_adapter` leaves the rest of the run reading a legacy adapter as `rank_major` and - scrambling it. Set it before the process does any MoE work and leave it set. Reading - the layout back off the adapter itself needs the `adapter_config.json` marker, which is - the migration path, not this switch. - """ - global _MOE_LORA_B_LAYOUT - # Refreshed on every call, including while Dynamo traces. Reading the environment at - # trace time bakes the then-current value into the graph, which is what the supported - # usage needs: `import unsloth` runs this module's import long before the application - # sets the variable, so an import-time value alone would be stale for exactly the - # normal case of setting it before training starts. Gating the refresh on - # `not is_compiling()` got this wrong, because a run whose first MoE call is already - # compiled never executes the eager branch at all. - # - # What this still does not do is notice a change made AFTER the first compiled call: - # Dynamo installs a guard on an `os.environ.get` only when the key is set at trace - # time, so there is nothing to invalidate the graph. That is the documented contract, - # a process-wide mode set before the process does any MoE work. - _MOE_LORA_B_LAYOUT = _read_moe_lora_b_layout_from_env() - return _MOE_LORA_B_LAYOUT - - -def _resolve_moe_lora_b_layout(layout) -> str: - """The layout to use, whether the caller named one or left it to the environment. - - A caller that names one still gets it validated. Both readers branch on "is this - rank_major" and fall through to grouped_by_expert otherwise, so an unvalidated typo - is not a no-op: it permutes the columns of a standard adapter, silently, which is the - exact damage `moe_lora_b_layout()` validates the environment variable to prevent. A - converter passing the layout in from a marker or a command line is the likeliest - source of one.""" - # This Unsloth Zoo code section is licensed under AGPL3 - - if layout is None: - return moe_lora_b_layout() - if layout not in _LORA_B_LAYOUTS: - raise MoELoRABLayoutError( - f"Unsloth: lora_B layout must be one of {_LORA_B_LAYOUTS}, got {layout!r}." - ) - return layout - - -def _moe_lora_b_column_permutation(num_experts, rank_per_expert, device, invert = False): - """The column permutation between PEFT's rank-major packing and expert-major order. - - Forward: `out[:, k] = weight_B[:, perm[k]]`. The inverse is the same construction with - the two axes swapped, so neither direction needs the other's indices kept around.""" - total = num_experts * rank_per_expert - rows, columns = (num_experts, rank_per_expert) if invert else (rank_per_expert, num_experts) - return torch.arange(total, device = device).view(rows, columns).t().reshape(-1) - - -class _ExpertMajorGroupedOperand(torch.autograd.Function): - """`(out, E*rank)` -> `(E, rank, out)` for the grouped GEMM, in ONE materialization. - - Gathering into expert-major column order and then permuting to the operand's shape - copies the whole tensor twice per projection, on every expert layer of every step. - Transposing to `(E*rank, out)` first makes the gather itself produce the operand: the - rows land in expert-major order, `index_select` on dim 0 writes a contiguous result, - and the final reshape is free. - - Still a real gather, which is the property the grouped GEMM needs: under rank-major - packing the expert axis has stride 1, so a pure view chain lets Inductor elide the - copy and hand `aten._grouped_mm` an operand that is neither row nor column major. - Still saves nothing for the backward either, since the permutation is recomputed from - two integers, which is what non-reentrant checkpointing requires.""" - # This Unsloth Zoo code section is licensed under AGPL3 - - @staticmethod - def forward(ctx, weight_B, num_experts, rank_per_expert, layout): - ctx.num_experts = num_experts - ctx.rank_per_expert = rank_per_expert - ctx.layout = layout - ctx.dim_B = weight_B.shape[0] - rows = weight_B.t() - if layout == LORA_B_LAYOUT_RANK_MAJOR: - columns = _moe_lora_b_column_permutation( - num_experts, rank_per_expert, weight_B.device, - ) - rows = rows.index_select(0, columns) - else: - rows = rows.contiguous() - return rows.reshape(num_experts, rank_per_expert, ctx.dim_B) - - @staticmethod - def backward(ctx, grad_output): - rows = grad_output.reshape(ctx.num_experts * ctx.rank_per_expert, ctx.dim_B) - if ctx.layout == LORA_B_LAYOUT_RANK_MAJOR: - columns = _moe_lora_b_column_permutation( - ctx.num_experts, ctx.rank_per_expert, grad_output.device, invert = True, - ) - rows = rows.index_select(0, columns) - return rows.t(), None, None, None - - -class _ExpertMajorColumns(torch.autograd.Function): - """Reorder `lora_B`'s columns expert-major, saving nothing for the backward. - - `index_select` would do this in one line, but it saves its index tensor, and - `torch.utils.checkpoint(use_reentrant = False)` counts every saved tensor and refuses - when the recompute saves a different number than the forward. The measuring call in - `_patched_param_wrapper_forward` already runs the experts forward a different number - of times than the recompute does, so one more saved tensor on the default path is - enough to turn that into `CheckpointError`, which is transformers 5's default - configuration and Unsloth's own distributed vision path. - - The permutation is fully determined by two integers, so the backward rebuilds the - inverse from `ctx` rather than from a saved tensor: nothing is saved, the count is - unchanged from the pure-view spelling, and the gradient still reaches `lora_B`. - A real gather is also what Inductor cannot fold back into a strided view of the - source, which is the other half of why this is not a plain `permute`.""" - # This Unsloth Zoo code section is licensed under AGPL3 - - @staticmethod - def forward(ctx, weight_B, num_experts, rank_per_expert): - ctx.num_experts = num_experts - ctx.rank_per_expert = rank_per_expert - columns = _moe_lora_b_column_permutation( - num_experts, rank_per_expert, weight_B.device, - ) - return weight_B.index_select(1, columns) - - @staticmethod - def backward(ctx, grad_output): - columns = _moe_lora_b_column_permutation( - ctx.num_experts, ctx.rank_per_expert, grad_output.device, invert = True, - ) - return grad_output.index_select(1, columns), None, None - - -def _moe_lora_b_columns_expert_major(weight_B, num_experts, rank_per_expert, layout): - """`lora_B`'s columns reordered so the expert index is slowest, without copying. - - Shared by `unflatten_moe_lora_b` and the grouped-GEMM path so the two cannot disagree - about the column order, while each finishes with its OWN single permute into the shape - it wants. Doing the reorder here and the permute there is what keeps the copy count at - one: composing the two public shapes instead (unflatten, then transpose + contiguous) - materialises the tensor twice, which is an extra saved tensor and breaks non-reentrant - gradient checkpointing, whose recompute must save exactly what the forward did. - - Returns `weight_B` untouched under grouped-by-expert, where the columns already are - expert-slowest.""" - # This Unsloth Zoo code section is licensed under AGPL3 - - layout = _resolve_moe_lora_b_layout(layout) - if layout != LORA_B_LAYOUT_RANK_MAJOR: - return weight_B - # PEFT packs expert index fastest: column j holds expert `j % num_experts`. - return _ExpertMajorColumns.apply(weight_B, num_experts, rank_per_expert) - - -def unflatten_moe_lora_b( - weight_B: torch.Tensor, - num_experts: int, - rank_per_expert: int, - dim_B: int, - layout: str = None, -) -> torch.Tensor: - """`(out, num_experts * rank)` -> `(num_experts, out, rank)`. - - The single place the expert axis of a fused `lora_B` is resolved, so the forward, the - merge and any converter cannot drift apart. Both layouts end in `.contiguous()` on a - permuted view, so neither is cheaper than the other and the choice is purely one of - which convention the stored tensor was written in. - - The rank-major branch regroups the columns with `index_select` before it views them, - rather than permuting the expert axis out of a three-way view. The two spell the same - tensor, but only the first survives `torch.compile`. Under rank-major packing the - expert axis has stride 1 in `weight_B`, so every view that puts experts first leaves - BOTH remaining axes with a stride above 1, and `aten._grouped_mm` rejects a `mat_b` - that is neither row nor column major. Eager is fine because `.contiguous()` really - copies; Inductor folds the whole view chain into one strided read of `weight_B` and - picks the source layout, so the copy disappears and the consumer sees - `(1, num_experts, num_experts * rank)` and raises `Invalid strides/sizes`. Measured on - `gemma-4-26B-A4B-it` (E=128, rank=8, out=1408) and reproduced standalone; the - pre-existing grouped-by-expert spelling escaped only because eliding ITS copy happens - to leave a legal column-major operand. `index_select` is a real gather that Inductor - cannot fold into a view, it keeps the gradient flowing to `lora_B`, and its indices are - a permutation, so its `index_add` backward has no duplicate targets to reduce - nondeterministically.""" - weight_B = _moe_lora_b_columns_expert_major( - weight_B, num_experts, rank_per_expert, layout, - ) - return weight_B.reshape(dim_B, num_experts, rank_per_expert).permute(1, 0, 2).contiguous() - - -def moe_lora_b_expert_columns( - expert_idx: int, - num_experts: int, - rank_per_expert: int, - layout: str = None, -) -> slice: - """Which columns of a flat `(out, num_experts * rank)` `lora_B` belong to one expert. - - A `slice`, so indexing with it stays a view. The companion of - `unflatten_moe_lora_b` for the merge paths that walk one expert at a time; the two - must agree, which is why both live here. `lora_A`'s rows for the same expert are - always `expert_idx * rank : (expert_idx + 1) * rank`, in both layouts, and the - resulting column order matches that rank order.""" - layout = _resolve_moe_lora_b_layout(layout) - if layout == LORA_B_LAYOUT_RANK_MAJOR: - return slice(expert_idx, num_experts * rank_per_expert, num_experts) - return slice(expert_idx * rank_per_expert, (expert_idx + 1) * rank_per_expert) - - -def _canonical_lora_weights_for_grouped_mm( - weight_A: torch.Tensor, - weight_B: torch.Tensor, - num_experts: int, - rank_per_expert: int, - dim_A: int, - dim_B: int, -) -> Tuple[torch.Tensor, torch.Tensor]: - first_weight = weight_A.view(num_experts, rank_per_expert, dim_A) - first_weight = first_weight.permute(0, 2, 1).contiguous() - # (num_experts, rank, out) for X @ first @ second, in ONE copy. Going through - # unflatten_moe_lora_b and transposing its result materialises the tensor twice, and - # the extra saved tensor makes non-reentrant gradient checkpointing recompute a - # different number of tensors than the forward saved. - second_weight = _ExpertMajorGroupedOperand.apply( - weight_B, num_experts, rank_per_expert, _resolve_moe_lora_b_layout(None), - ) - return first_weight, second_weight - - -def _reversed_lora_weights_for_grouped_mm( - weight_A: torch.Tensor, - weight_B: torch.Tensor, - num_experts: int, - rank_per_expert: int, - dim_A: int, - dim_B: int, -) -> Tuple[torch.Tensor, torch.Tensor]: - first_weight = unflatten_moe_lora_b(weight_B, num_experts, rank_per_expert, dim_B) - second_weight = weight_A.view(num_experts, rank_per_expert, dim_A).contiguous() - return first_weight, second_weight - - -def _get_param_shape_from_module(module, parameter_name): - if module is None or parameter_name is None or not hasattr(module, parameter_name): - return None - param = getattr(module, parameter_name) - if hasattr(param, "get_param"): - param = param.get_param() - elif hasattr(param, "weight"): - param = param.weight - if is_mxfp4_expert_param(param): - return tuple(param._original_shape) - return tuple(param.shape) - - -def _get_moe_lora_io_dims(wrapper, experts_module=None): - base = None - if wrapper is not None and hasattr(wrapper, "get_base_layer"): - base = wrapper.get_base_layer() - if experts_module is None: - experts_module = base - if experts_module is None: - experts_module = getattr(wrapper, "base_layer", None) - - parameter_name = getattr(wrapper, "parameter_name", None) - source = experts_module if experts_module is not None else base - if source is None: - return None, None - _set_gpt_oss_grouped_mm_format_on_experts(source) - - shape = _get_param_shape_from_module(source, parameter_name) - if shape is not None and len(shape) >= 3: - grouped_mm_format = bool(getattr(source, "_unsloth_grouped_mm_format", False)) - if grouped_mm_format: - return shape[-2], shape[-1] - return shape[-1], shape[-2] - - hidden_dim = getattr(source, "hidden_dim", None) - intermediate_dim = getattr(source, "intermediate_dim", None) - if hidden_dim is None or intermediate_dim is None: - return None, None - if parameter_name == "gate_up_proj": - return hidden_dim, 2 * intermediate_dim - if parameter_name == "down_proj": - return intermediate_dim, hidden_dim - return None, None - - -def extract_moe_lora_weights_for_grouped_mm( - wrapper, - weight_A: torch.Tensor, - weight_B: torch.Tensor, - scaling, - num_experts: int, - *, - experts_module=None, - input_dim=None, - output_dim=None, - model_name: str = "MoE", - enable_logging: bool = None, - logger_obj=None, -) -> Tuple[torch.Tensor, torch.Tensor, float, int]: - total_rank = weight_A.shape[0] - rank_per_expert = total_rank // num_experts - dim_A = weight_A.shape[1] - dim_B = weight_B.shape[0] - - if num_experts <= 1: - return weight_A.T, weight_B.T, scaling, num_experts - - if input_dim is None or output_dim is None: - inferred_input_dim, inferred_output_dim = _get_moe_lora_io_dims( - wrapper, experts_module=experts_module, - ) - if input_dim is None: - input_dim = inferred_input_dim - if output_dim is None: - output_dim = inferred_output_dim - - canonical_match = ( - input_dim is not None - and output_dim is not None - and dim_A == input_dim - and dim_B == output_dim - ) - reversed_match = ( - input_dim is not None - and output_dim is not None - and dim_A == output_dim - and dim_B == input_dim - ) - - if canonical_match and reversed_match: - # Square stacks match both readings by shape, and which layouts PEFT swaps changed in 0.19.0 / 0.19.1: neither flag decides alone. - swapped = bool(getattr(wrapper, "_did_swap_in_out_features", False)) - base = wrapper.get_base_layer() if hasattr(wrapper, "get_base_layer") else None - stored_in_out = ( - getattr(base, "is_transposed", None) is True - or bool(getattr(base, "_unsloth_grouped_mm_format", False)) - ) - if swapped != stored_in_out: - first_weight, second_weight = _canonical_lora_weights_for_grouped_mm( - weight_A, weight_B, num_experts, rank_per_expert, dim_A, dim_B, - ) - else: - first_weight, second_weight = _reversed_lora_weights_for_grouped_mm( - weight_A, weight_B, num_experts, rank_per_expert, dim_A, dim_B, - ) - return first_weight, second_weight, scaling, num_experts - - if canonical_match: - first_weight, second_weight = _canonical_lora_weights_for_grouped_mm( - weight_A, weight_B, num_experts, rank_per_expert, dim_A, dim_B, - ) - return first_weight, second_weight, scaling, num_experts - - if reversed_match: - first_weight, second_weight = _reversed_lora_weights_for_grouped_mm( - weight_A, weight_B, num_experts, rank_per_expert, dim_A, dim_B, - ) - return first_weight, second_weight, scaling, num_experts - - if logger_obj is not None: - if enable_logging is None: - enable_logging = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" - if enable_logging and (input_dim is not None or output_dim is not None): - logger_obj.warning( - f"Unsloth: {model_name} LoRA extractor could not match either layout " - f"(weight_A={tuple(weight_A.shape)}, weight_B={tuple(weight_B.shape)}, " - f"expected input_dim={input_dim}, output_dim={output_dim}, " - f"num_experts={num_experts}). Falling back to canonical layout. " - "If this is a new PEFT version, the LoRA delta may be wrong." - ) - - first_weight, second_weight = _canonical_lora_weights_for_grouped_mm( - weight_A, weight_B, num_experts, rank_per_expert, dim_A, dim_B, - ) - return first_weight, second_weight, scaling, num_experts - - -def _extract_lora_from_wrapper( - wrapper, adapter_name: str = "default", experts_module=None -) -> Optional[Tuple[torch.Tensor, torch.Tensor, float, int]]: - """Extract LoRA weights from a PEFT ParamWrapper for MoE separated grouped_mm. - - PEFT 3D ParamWrapper gives lora_A: (E*R, in_dim), lora_B: (out_dim, E*R); - reshaped to first_weight (E, in_dim, R), second_weight (E, R, out_dim) so - delta = X @ first @ second. Handles both standard (E, out, in) Qwen3-MoE and - transposed (E, in, out) Qwen3-VL-MoE base weight layouts. - - Returns (first_weight, second_weight, scaling, num_experts) or None. - """ - # This Unsloth Zoo code section is licensed under AGPL3 - - try: - if not hasattr(wrapper, "lora_A") or not hasattr(wrapper, "lora_B"): - return None - - if hasattr(wrapper, "disable_adapters") and wrapper.disable_adapters: - return None - if hasattr(wrapper, "merged") and wrapper.merged: - return None - - if not wrapper.lora_A: - return None - - if adapter_name not in wrapper.lora_A: - adapter_name = list(wrapper.lora_A.keys())[0] - - lora_A_module = wrapper.lora_A[adapter_name] - lora_B_module = wrapper.lora_B[adapter_name] - - weight_A = lora_A_module.weight # (E*R, dim1) - weight_B = lora_B_module.weight # (dim2, E*R) - scaling = wrapper.scaling[adapter_name] - num_experts = getattr(wrapper, "num_experts", 1) - - if experts_module is None: - experts_module = wrapper.get_base_layer() if hasattr(wrapper, "get_base_layer") else None - - # Model-specific LoRA extractor attached to the experts module, if any. - extractor_fn = getattr(experts_module, "_unsloth_lora_extractor_fn", None) - - if extractor_fn is not None: - return extractor_fn(wrapper, weight_A, weight_B, scaling, num_experts) - - return extract_moe_lora_weights_for_grouped_mm( - wrapper, - weight_A, - weight_B, - scaling, - num_experts, - experts_module=experts_module, - model_name="MoE", - ) - except MoELoRABLayoutError: - # A misspelled UNSLOTH_MOE_LORA_B_LAYOUT must not read as "no adapter here". - raise - except Exception: - return None - - -def _extract_lora_weights( - param, adapter_name: str = "default", num_experts: int = None, experts_module=None -) -> Optional[Tuple[torch.Tensor, torch.Tensor, float]]: - """Compat wrapper around _extract_lora_from_wrapper; returns (first, second, scaling).""" - # This Unsloth Zoo code section is licensed under AGPL3 - - # Pass num_experts through so _extract_lora_from_wrapper can use it. - if num_experts is not None and not hasattr(param, "num_experts"): - param.num_experts = num_experts - - result = _extract_lora_from_wrapper(param, adapter_name, experts_module=experts_module) - if result is None: - return None - return result[0], result[1], result[2] - - -_DEQUANTIZE_4BIT_IN_SLICES = None # unresolved -_DEQUANTIZE_4BIT_IN_SLICES_MISSING = False # resolved, and there is none - - -def _get_dequantize_4bit_in_slices(): - """The sliced 4-bit read, resolved once. - - Absolute, like the dispatcher above and for the same reason: this file is - also copied to unsloth_compiled_cache/moe_utils.py and imported as a - top-level module, where a relative import of a sibling raises and the - caller would quietly put an oversized stack back on the call that aborts. - Only the import is guarded; a real failure inside the helper must propagate. - """ - global _DEQUANTIZE_4BIT_IN_SLICES, _DEQUANTIZE_4BIT_IN_SLICES_MISSING - if _DEQUANTIZE_4BIT_IN_SLICES is None and not _DEQUANTIZE_4BIT_IN_SLICES_MISSING: - try: - from unsloth_zoo.temporary_patches.moe_utils_bnb4bit import ( - _dequantize_4bit_in_slices, - ) - _DEQUANTIZE_4BIT_IN_SLICES = _dequantize_4bit_in_slices - except ImportError: - _DEQUANTIZE_4BIT_IN_SLICES_MISSING = True - return _DEQUANTIZE_4BIT_IN_SLICES - - -def _get_base_weight(param, target_dtype=None, token_counts=None, experts=None): - """Get base weight from a potentially wrapped parameter or module. target_dtype (recompute - providers) restores the packed Params4bit to its logical shape and casts. experts (a - _SelectiveDequantPlan chunk) returns only those experts of a packed 4-bit stack.""" - # This Unsloth Zoo code section is licensed under AGPL3 - - while hasattr(param, "base_layer"): - param = param.base_layer - - if experts is not None: - return _dequantize_bnb4bit_experts(param, experts, target_dtype) - - if is_mxfp4_expert_param(param): - return param.dequantize(dtype = target_dtype, token_counts = token_counts) - - if HAS_BNB and isinstance(param, Params4bit): - if getattr(param, "quant_state", None) is None: - raise RuntimeError( - "unsloth: _get_base_weight saw a Params4bit with quant_state=None. " - "This usually means the model was used in forward before loading " - "completed quantization (meta placeholder still in place), or the " - "MoE quantizer patch did not fire for this expert. " - f"data.shape={tuple(param.data.shape)}, device={param.device}." - ) - # Absolute import: this file is also copied into unsloth_compiled_cache and imported top-level. - weight = None - try: - from unsloth_zoo.temporary_patches.moe_triton_kernels import nf4_dequant_triton - except ImportError: - nf4_dequant_triton = None - if nf4_dequant_triton is not None: - weight = nf4_dequant_triton( - param.data, param.quant_state, getattr(param, "_original_shape", None), - ) - if weight is None: - # >= 2**31-element stacks abort in bitsandbytes dequantize (csrc/ops.cu:93); slice like the load does. - slicer = _get_dequantize_4bit_in_slices() - weight = slicer(param) if slicer is not None else None - if weight is None: - weight = bnb.functional.dequantize_4bit(param.data, param.quant_state) - original_shape = getattr(param, "_original_shape", None) - if original_shape is not None and weight.shape != original_shape: - weight = weight.reshape(original_shape) - if target_dtype is not None: - weight = weight.to(target_dtype) - return weight - - if hasattr(param, "get_param"): - return param.get_param() - - if hasattr(param, "weight"): - return param.weight - - return param - - -def _get_lora_wrapper_for_param(experts_module, param_name): - """Get the PEFT ParamWrapper for gate_up_proj or down_proj; does not lazily set up wrappers. - - A forward that asks for the wrapper applies that parameter's LoRA itself rather than - reading the stash, so finding one counts as a read (see `mark_moe_lora_stash_read`). - """ - # This Unsloth Zoo code section is licensed under AGPL3 - - wrapper = None - if hasattr(experts_module, f"{param_name}_lora_wrapper"): - wrapper = getattr(experts_module, f"{param_name}_lora_wrapper") - else: - if isinstance(inspect.getattr_static(type(experts_module), param_name, None), property): - # Don't evaluate a property (packed MXFP4 decodes on read); its setter stores the wrapper. - attr = experts_module.__dict__.get("_" + param_name, None) - else: - attr = getattr(experts_module, param_name, None) - if attr is not None and hasattr(attr, "lora_A"): # ParamWrapper - wrapper = attr - - if wrapper is not None: - mark_moe_lora_stash_read(experts_module, param_name) - return wrapper - - -def native_moe_grouped_mm( - inputs: torch.Tensor, weight: torch.Tensor, offsets: torch.Tensor -) -> torch.Tensor: - """Grouped_mm with backward fix for PyTorch's grouped_mm backward stride bug.""" - return _grouped_mm_with_backward_fix(inputs, weight, offsets) - - -def _lora_rank_pad_needed(inputs, first_weight) -> bool: - """False when both LoRA GEMMs take Triton (masks any K / N, no 16-byte stride rule). - rank <= 32 keeps both GEMMs (min(K, N) = rank) in the same "lora" row class.""" - if inputs is None or first_weight.shape[-1] > 32: - return True - if os.environ.get("UNSLOTH_MOE_LORA_PAD_ALWAYS", "0") == "1": - return True - if not _triton_grouped_mm_wanted(inputs, first_weight): - return True - # A failed self-check falls back to torch._grouped_mm, which needs the pad. - from unsloth_zoo.temporary_patches import moe_grouped_fp16 as mg - if torch.compiler.is_compiling(): - # Never run the self-check while tracing; skip the pad only if it already passed. - return not (mg._GENERIC_DISABLED_REASON is None and mg._GENERIC_SELF_CHECKED.get(inputs.device.index) is True) - return not mg.triton_grouped_available(inputs.device) - - -def _pad_lora_rank_for_grouped_mm(first_weight, second_weight, inputs = None): - """Zero-pad the LoRA rank of (E, in, R) / (E, R, out) so torch._grouped_mm gets 16-byte - aligned rows (bf16 rank 4 or 6 fell back to the per-expert loop, and aborted compile). - Zero columns meet zero rows, so the product is exact. Minimum 8 even for fp32: autocast - lowers grouped_mm operands to bf16.""" - align = max(8, 16 // first_weight.element_size()) - rank = first_weight.shape[-1] - pad = (-rank) % align - if pad == 0 or not _lora_rank_pad_needed(inputs, first_weight): - return first_weight, second_weight - first_weight = F.pad(first_weight, (0, pad)) - second_weight = F.pad(second_weight, (0, 0, 0, pad)) - return first_weight, second_weight - - -def _apply_lora_grouped_mm( - inputs: torch.Tensor, - lora_B: torch.Tensor, - lora_A: torch.Tensor, - offsets: torch.Tensor, - scaling: float, - grouped_mm_func=native_moe_grouped_mm, -) -> torch.Tensor: - """Apply LoRA via grouped GEMM: result = ((X @ B) @ A) * scaling. - - inputs (total_tokens, in_dim); lora_B (E, in_dim, R); lora_A (E, R, out_dim). - """ - # This Unsloth Zoo code section is licensed under AGPL3 - - lora_B, lora_A = _pad_lora_rank_for_grouped_mm(lora_B, lora_A, inputs) - # X @ B then result @ A; both already in native (E, ...) layout, no transpose. - lora_intermediate = grouped_mm_func(inputs, lora_B.contiguous(), offsets) - lora_delta = grouped_mm_func(lora_intermediate, lora_A.contiguous(), offsets) - - return lora_delta * scaling - - -def _should_use_separated_lora() -> bool: - """Use separated LoRA (default True); UNSLOTH_MOE_LORA_MERGED=1 forces the merged path.""" - return os.environ.get("UNSLOTH_MOE_LORA_MERGED", "0") != "1" - - -# Model-specific weight preprocessing hooks: each model registers a transposition -# function so the generic backend works across weight layouts. -# -# unsloth_compiled_cache holds copies of this file, loaded as their own module objects -# (unsloth_cached_moe_utils, whose forward get_forward_moe_backend() prefers, and the -# bare `moe_utils` compiled modules import), each with its own empty dict. Resolving -# through the package keeps registration and lookup on one dict whichever copy runs, -# so a registration is not dropped and left to layout inference (#849). - -_WEIGHT_PREPROCESSORS = {} - - -def _weight_preprocessor_registry(): - """The registry shared by every loaded copy of this module (see above).""" - package_module = sys.modules.get("unsloth_zoo.temporary_patches.moe_utils") - if package_module is None: - return _WEIGHT_PREPROCESSORS # package copy not loaded, so nothing registered - return getattr(package_module, "_WEIGHT_PREPROCESSORS", _WEIGHT_PREPROCESSORS) - - -def register_weight_preprocessor(model_type: str, preprocessor_fn): - """Register a weight preprocessor (weight, proj_type, hidden_dim) -> weight for a model type.""" - _weight_preprocessor_registry()[model_type] = preprocessor_fn - - -def get_weight_preprocessor(model_type: str): - """Get registered weight preprocessor for model type.""" - return _weight_preprocessor_registry().get(model_type) - - -_LOGICAL_SHAPE_MAX_DEPTH = 16 # PEFT / quant wrapper chains are a few levels deep - - -def _logical_expert_shape(param): - """Return the logical expert shape without dequantizing or materializing. - - Resolve PEFT parameters before module wrappers. A recorded logical shape (bnb - Params4bit _original_shape, ParameterModule shape_3d) is read directly so a - get_param provider is never materialized just to read its shape. - """ - # This Unsloth Zoo code section is licensed under AGPL3 - - # Bounded walk, no id(): id() of a traced parameter guards on that object's identity, so a - # compiled decoder layer recompiled once per layer (each layer owns its expert parameters). - for _ in range(_LOGICAL_SHAPE_MAX_DEPTH): - if param is None: - break - - # Recorded logical shape avoids materializing a get_param provider. - recorded = getattr(param, "_original_shape", None) - if recorded is None: - recorded = getattr(param, "shape_3d", None) - if recorded is not None: - return tuple(int(x) for x in recorded) - - get_param = getattr(param, "get_param", None) - if callable(get_param): - try: - inner = get_param() - except Exception: - inner = None - if inner is not None and inner is not param: - param = inner - continue - - base_layer = getattr(param, "base_layer", None) - if base_layer is not None and base_layer is not param: - param = base_layer - continue - - # Linear modules store parameters on .weight. - weight = getattr(param, "weight", None) - if weight is not None and weight is not param: - param = weight - continue - - break - - shape = getattr(param, "shape", None) - if shape is None: - return None - try: - return tuple(int(x) for x in shape) - except TypeError: - return None - - -def _orientation_needs_transpose(shape, proj_type, hidden_dim): - """Return whether a weight needs transposition, or None when ambiguous.""" - # This Unsloth Zoo code section is licensed under AGPL3 - - if shape is None or len(shape) < 3: - return None - d1, d2 = int(shape[1]), int(shape[2]) - if d1 == d2: - return None - if proj_type == "gate_up": - # grouped_mm uses (E, hidden, 2I). - return d1 != hidden_dim - # grouped_mm down uses (E, I, hidden). - return d2 != hidden_dim - - -_WARNED_AMBIGUOUS_LAYOUTS = set() - - -def _warn_ambiguous_layout_once(proj_type, shape, hidden_dim): - # This Unsloth Zoo code section is licensed under AGPL3 - - key = (proj_type, tuple(shape), hidden_dim) - if key in _WARNED_AMBIGUOUS_LAYOUTS: - return - _WARNED_AMBIGUOUS_LAYOUTS.add(key) - # Always warn because a wrong guess corrupts training (#849). - print( - f"Unsloth: MoE '{proj_type}' expert weight of shape {tuple(shape)} has equal matmul " - f"dims (hidden_dim={hidden_dim}), so its layout is ambiguous, and no unambiguous " - f"sibling projection was reachable to disambiguate it. Assuming the transformers " - f"F.linear layout and transposing to grouped_mm layout. If these experts were " - f"already in grouped_mm layout this transpose is WRONG and trains on transposed " - f"weights (see unslothai/unsloth-zoo#849). Register an explicit preprocessor via " - f"register_weight_preprocessor('', ...) to remove the guess.", - file=sys.stderr, - ) - - -def preprocess_weight( - weight: torch.Tensor, proj_type: str, hidden_dim: int, model_type=None, - experts_module=None, -): - """Convert an expert weight to grouped_mm layout. - - Registered preprocessors take precedence. Square weights use the non-square sibling. - """ - # This Unsloth Zoo code section is licensed under AGPL3 - - # Guarded: this runs per projection per MoE forward and again in recompute - # backward, and the shipped models pass model_type=None. - if model_type: - registry = _weight_preprocessor_registry() - if model_type in registry: - return registry[model_type](weight, proj_type, hidden_dim) - - # Non-square shapes reveal layout directly. - needs_transpose = _orientation_needs_transpose(tuple(weight.shape), proj_type, hidden_dim) - if needs_transpose is not None: - return weight.transpose(-2, -1) if needs_transpose else weight - - if getattr(experts_module, "_unsloth_grouped_mm_format", False): - return weight - is_transposed = getattr(experts_module, "is_transposed", None) - if isinstance(is_transposed, bool): - return weight if is_transposed else weight.transpose(-2, -1) - - # One sibling is always non-square and reveals the shared layout. - if experts_module is not None: - sibling_name = "down_proj" if proj_type == "gate_up" else "gate_up_proj" - sibling_type = "down" if proj_type == "gate_up" else "gate_up" - sibling = getattr(experts_module, sibling_name, None) - if sibling is not None: - sibling_transpose = _orientation_needs_transpose( - _logical_expert_shape(sibling), sibling_type, hidden_dim, - ) - if sibling_transpose is not None: - return weight.transpose(-2, -1) if sibling_transpose else weight - - # Default to F.linear when no sibling is readable. - _warn_ambiguous_layout_once(proj_type, weight.shape, hidden_dim) - return weight.transpose(-2, -1) - - -# Generic MoE detection and ParamWrapper patching. - - -def _normalize_model_type(value) -> str: - if value is None: - return "" - return str(value).lower().replace("-", "_") - - -def _iter_model_configs(model): - seen = set() - queue = [model] - while queue and len(seen) < 8: - current = queue.pop(0) - if current is None: - continue - current_id = id(current) - if current_id in seen: - continue - seen.add(current_id) - - config = getattr(current, "config", None) - if config is not None: - yield config - - for attr in ("base_model", "model"): - nested = getattr(current, attr, None) - if nested is not None and nested is not current: - queue.append(nested) - - -def _is_gpt_oss_model(model) -> bool: - for config in _iter_model_configs(model): - model_type = _normalize_model_type(getattr(config, "model_type", None)) - if model_type == "gpt_oss": - return True - - for attr in ("_name_or_path", "name_or_path"): - name = getattr(config, attr, None) - if name is None: - continue - # Match only the final path component so parent directories like - # /data/gpt-oss-tests/qwen3-7b do not count as gpt-oss. - base = str(name).replace("\\", "/").rstrip("/").rsplit("/", 1)[-1] - if "gpt_oss" in _normalize_model_type(base): - return True - - return False - - -def _set_gpt_oss_grouped_mm_format_on_experts(module) -> bool: - if module is None: - return False - if bool(getattr(module, "_unsloth_grouped_mm_format", False)): - return False - if getattr(module, "is_transposed", None) is True: - module._unsloth_grouped_mm_format = True - return True - if module.__class__.__name__ != "GptOssExperts": - return False - # Require the gpt-oss (E, in, out) weight signature: gate_up's out dim is - # twice down's in dim. Same-named classes with other layouts stay unflagged. - gate_shape = _get_param_shape_from_module(module, "gate_up_proj") - down_shape = _get_param_shape_from_module(module, "down_proj") - if gate_shape is None or down_shape is None: - return False - if len(gate_shape) < 3 or len(down_shape) < 3: - return False - if gate_shape[0] != down_shape[0]: - return False - if gate_shape[-2] != down_shape[-1] or gate_shape[-1] != 2 * down_shape[-2]: - return False - module._unsloth_grouped_mm_format = True - return True - - -def patch_gpt_oss_grouped_mm_format(model) -> int: - """ - Mark GPT-OSS experts as storing weights in grouped_mm format. - - Stock transformers GPT-OSS experts use (E, in_dim, out_dim) tensors but do - not carry Unsloth's `_unsloth_grouped_mm_format` instance flag. Set it on - live expert modules so the shared MoE LoRA extractor chooses GPT-OSS - ordering instead of the Qwen-style fallback. - """ - # This Unsloth Zoo code section is licensed under AGPL3 - - if model is None or not _is_gpt_oss_model(model): - return 0 - - modules = getattr(model, "modules", None) - if not callable(modules): - return 0 - - updated = 0 - for module in modules(): - if _set_gpt_oss_grouped_mm_format_on_experts(module): - updated += 1 - return updated - - -def _patch_peft_get_peft_model_for_moe(): - # This Unsloth Zoo code section is licensed under AGPL3 - - global _original_peft_get_peft_model - if _original_peft_get_peft_model is not None: - return - - try: - import peft - except Exception: - return - - original_get_peft_model = getattr(peft, "get_peft_model", None) - if original_get_peft_model is None: - return - if getattr(original_get_peft_model, "_unsloth_moe_patched", False): - return - - _original_peft_get_peft_model = original_get_peft_model - - @wraps(original_get_peft_model) - def patched_get_peft_model(model, *args, **kwargs): - peft_model = original_get_peft_model(model, *args, **kwargs) - try: - patch_gpt_oss_grouped_mm_format(model) - if peft_model is not model: - patch_gpt_oss_grouped_mm_format(peft_model) - except Exception: - pass - return peft_model - - patched_get_peft_model._unsloth_moe_patched = True - peft.get_peft_model = patched_get_peft_model - - for module_name in ("peft.mapping_func", "peft.mapping"): - try: - module = importlib.import_module(module_name) - except Exception: - continue - if getattr(module, "get_peft_model", None) is original_get_peft_model: - module.get_peft_model = patched_get_peft_model - - -def _is_moe_experts_module(module) -> bool: - """Generic check for an MoE experts layer with stacked 3D expert weights. - - Matches gate_up_proj/down_proj (Qwen3-MoE etc.) or w1/w2/w3 (older models). - """ - # This Unsloth Zoo code section is licensed under AGPL3 - - import torch.nn as nn - - # After PEFT's parametrize wrapping, gate_up_proj is a Tensor (not Parameter), - # so accept both. - if hasattr(module, "gate_up_proj"): - param = module.gate_up_proj - # 4-bit params are packed into 2D tensors. - if HAS_BNB and isinstance(param, Params4bit) and param.ndim == 2: - return True - if is_mxfp4_expert_param(param): - return True - # Standard MoE weights are 3D (num_experts, in, out). - if isinstance(param, (nn.Parameter, torch.Tensor)) and param.ndim in (2, 3): - return True - - # w1/w2 pattern (separate gate/up projections). - if hasattr(module, "w1") and hasattr(module, "w2"): - w1 = module.w1 - if isinstance(w1, (nn.Parameter, torch.Tensor)) and w1.ndim in (2, 3): - return True - - return False - - -# Aliases for compatibility with gpt_oss.py -_get_moe_lora_weights = _extract_lora_from_wrapper - - -# --------------------------------------------------------------------------------------- -# Which convention packs lora_B's (num_experts * rank) axis. -# -# PEFT's ParamWrapper builds lora_A as (E*R, in) and lora_B as (out, E*R) and pairs them -# in get_delta_weight by reading A as (E, R, in) and B as (out, R, E): lora_A is grouped -# by expert, lora_B is NOT, its expert index is the fastest axis. That reading is -# unchanged in PEFT 0.18.0, 0.19.0, 0.19.1, 0.20.0 and 0.21.0, and it is what vLLM's -# `_stack_moe_lora_weights` and every other consumer of a saved adapter implements, so it -# is the only reading a stored fused expert lora_B can be expected to have. -# -# Unsloth's separated forward read the same lora_B as (out, E, R) instead -# (`_canonical_lora_weights_for_grouped_mm`), so a rank-contiguous block of columns -# belonged to one expert. Both reshapes always fit, since E*R == R*E, so nothing ever -# raises and the disagreement is always silent. They disagree whenever num_experts > 1 -# and rank > 1, which is every real MoE checkpoint: equal num_experts and rank does not -# make them agree, because the two readings send flat column r*E + e and flat column -# e*R + r to the same place and those coincide for all (e, r) only when E == 1 or R == 1. -# -# The separated forward is being moved onto PEFT's packing (the `lora_B` packing fix, -# unsloth_zoo#1269), with `UNSLOTH_MOE_LORA_B_LAYOUT=grouped_by_expert` left as the -# opt-in for reading an adapter an older Unsloth trained. What stays here is the -# *recording* side: which packing the adapter on this model was trained with, written -# into adapter_config.json on save (unsloth#6930), because a downstream converter has no -# wrapper to ask and Unsloth stamps no version into the file either, so a marker-less -# adapter cannot be dated from its bytes. -# --------------------------------------------------------------------------------------- - -# Also defined next to the packing helpers by unsloth_zoo#1269; same strings, so the two -# definitions are interchangeable and either branch can land first. -LORA_B_LAYOUT_GROUPED_BY_EXPERT = "grouped_by_expert" -LORA_B_LAYOUT_RANK_MAJOR = "rank_major" - - -def _process_lora_b_layout() -> str: - """Which packing the separated MoE forward uses in THIS process. - - `moe_lora_b_layout()` (unsloth_zoo#1269) is the authority once that fix is in: it - reads `UNSLOTH_MOE_LORA_B_LAYOUT` and defaults to PEFT's `rank_major`. It is looked - up rather than imported because it lives in this same module, so the two can land in - either order. Without it the separated forward is unconditionally grouped by expert, - which is what a build predating the packing fix does, so that is the honest answer - there.""" - resolver = globals().get("moe_lora_b_layout") - if callable(resolver): - return resolver() - return LORA_B_LAYOUT_GROUPED_BY_EXPERT - - -def _legacy_lora_b_layout_requested() -> bool: - """Whether the caller has explicitly declared that the adapters in this process were - packed the pre-fix way, by setting `UNSLOTH_MOE_LORA_B_LAYOUT=grouped_by_expert`. - - Deliberately the raw environment variable and not `_process_lora_b_layout()`: this - gates a patch over PEFT's own `get_delta_weight`, and the packing of a *stored* - adapter is a property of the adapter, not of how this process happens to be routing - its forwards. Only an explicit statement about provenance may redirect PEFT's merge; - an unset variable leaves PEFT's reconstruction, which is the correct one for every - adapter trained with the packing fix in place, completely untouched.""" - return os.environ.get("UNSLOTH_MOE_LORA_B_LAYOUT") == LORA_B_LAYOUT_GROUPED_BY_EXPERT - -# The only parameter names the separated forward claims. Anything else on an experts -# module (the unfused experts.gate_proj / experts.up_proj pair that NemotronH uses, for -# instance) stays on PEFT's own forward and keeps PEFT's own packing. -_SEPARATED_MOE_LORA_PARAMETER_NAMES = ("gate_up_proj", "down_proj") - - -def _wrapper_uses_separated_moe_lora(wrapper, experts_module = None) -> bool: - """True when `_patched_param_wrapper_forward` routes this wrapper to the separated - MoE forward instead of PEFT's `_activate_lora`. Kept in one place so the forward, the - delta reconstruction and the saved layout marker cannot drift apart.""" - if not _should_use_separated_lora(): - return False - if getattr(wrapper, "parameter_name", None) not in _SEPARATED_MOE_LORA_PARAMETER_NAMES: - return False - if experts_module is None: - get_base_layer = getattr(wrapper, "get_base_layer", None) - if get_base_layer is None: - return False - try: - experts_module = get_base_layer() - except Exception: - return False - return _is_moe_experts_module(experts_module) - - -def _wrapper_forward_applies_stash(wrapper, store_attr = None): - """The measured verdict for this wrapper's experts forward, or None if unmeasured. - - `_wrapper_uses_separated_moe_lora` answers a structural question, "does the separated - forward claim this parameter". `_patched_param_wrapper_forward` then asks a stronger - one at run time, "did the forward that actually ran read the stash", and falls back to - PEFT when the answer is False. The layout follows the forward that ran, so it has to - read the same verdict.""" - # This Unsloth Zoo code section is licensed under AGPL3 - - parameter_name = getattr(wrapper, "parameter_name", None) - if parameter_name is None: - return None - get_base_layer = getattr(wrapper, "get_base_layer", None) - if get_base_layer is None: - return None - try: - experts_module = get_base_layer() - except Exception: - return None - if experts_module is None: - return None - try: - return moe_lora_forward_applies_stash(experts_module, parameter_name, store_attr) - except Exception: - return None - - -def _wrapper_has_adapter(wrapper, adapter_name) -> bool: - """Whether `adapter_name` has a LoRA on this wrapper at all. A PeftModel can carry a - fused expert adapter and a dense adapter side by side, and the fused expert wrappers - then exist for the whole model while only one adapter has weights in them, so every - layout answer has to be scoped to one adapter or it describes the wrong one. - - `adapter_name = None` asks the weaker question "does this wrapper hold any adapter at - all", which is what the unnamed form of `moe_lora_b_layout_for_wrapper` needs. Still - a question: - `delete_adapter` empties lora_A and leaves the wrapper in place, and claiming a - packing for weights that are gone is the same mistake in a smaller form.""" - lora_A = getattr(wrapper, "lora_A", None) - if lora_A is None: - return False - try: - if adapter_name is None: - return len(lora_A) != 0 - return adapter_name in lora_A - except Exception: - return False - - -def moe_lora_b_layout_for_wrapper(wrapper, adapter_name = None) -> str: - """Which convention packs `wrapper`'s lora_B columns for `adapter_name`. - - A wrapper the separated MoE forward owns is packed the way that forward packs, which - is `_process_lora_b_layout()`: PEFT's `rank_major` with the packing fix in place, - `grouped_by_expert` under `UNSLOTH_MOE_LORA_B_LAYOUT=grouped_by_expert` or on a build - predating the fix. Every other wrapper is on PEFT's own forward and is therefore - `rank_major`. A single expert (or a wrapper PEFT collapses to a plain Linear) has no - packing to get wrong and is reported as rank_major, which is what PEFT's own - reconstruction does for it. - - Named for the wrapper on purpose: `moe_lora_b_layout()` is the process-wide setting - and takes no arguments, this one answers for one wrapper and one adapter. - - `adapter_name` defaults to "whichever adapter this wrapper holds", which is the old - behaviour and is right for a model with one adapter. Name an adapter that has no LoRA - on this wrapper and the answer is rank_major, so the caller falls back to PEFT rather - than claiming a packing for weights that are not there. - - `wrapper` must be the PEFT ParamWrapper. Hand it anything else, the experts module - itself for instance, and the answer is rank_major, because nothing on it says the - separated forward claimed a parameter. That reads as a confident "PEFT packed this" - when the truth is "cannot tell", so callers that may hold something else have to - check for `parameter_name` and `lora_A` first.""" - if int(getattr(wrapper, "num_experts", 1) or 1) <= 1: - return LORA_B_LAYOUT_RANK_MAJOR - if not _wrapper_has_adapter(wrapper, adapter_name): - return LORA_B_LAYOUT_RANK_MAJOR - if _wrapper_uses_separated_moe_lora(wrapper): - verdict = _wrapper_forward_applies_stash(wrapper) - if verdict is None: - # A compiled-only run never reaches the eager probe. - verdict = _wrapper_forward_applies_stash(wrapper, _MOE_LORA_COMPILED_VERDICT_ATTR) - if verdict is False: - # The structural test says the separated forward claims this wrapper, but the - # forward that actually ran did not read the stash, so `_patched_param_wrapper_forward` - # handed the wrapper back to PEFT and PEFT trained it rank-major. That happens - # for a stacked-expert family Unsloth does not patch (Olmoe, for one). Believing - # the structural answer here would write a grouped_by_expert marker for - # rank-major weights, and, under the legacy override, merge them with the wrong - # pairing. Only a measured False overrides it; None means "not measured yet", - # which is not evidence either way. - return LORA_B_LAYOUT_RANK_MAJOR - return _process_lora_b_layout() - return LORA_B_LAYOUT_RANK_MAJOR - - -# Did the experts forward actually read the LoRA that ParamWrapper.forward handed it? -# -# `_patched_param_wrapper_forward` deliberately does not let PEFT fold the expert LoRA into -# the stacked expert weight. It stashes the factors on the experts module as -# `_unsloth_lora_` and lets the experts forward apply them as a separate -# grouped GEMM, which is cheaper and keeps the base weight quantized. That is only correct -# when the forward that runs is one that reads the stash. Unsloth installs such a forward for -# the MoE families it patches. For a stacked-expert family it does not patch, transformers' -# own experts forward runs, the stash is written and then deleted unread, and the expert LoRA -# has no effect on the output at all while requires_grad and the optimizer still report a -# healthy adapter. `_forward_native_fp8_expert_loop` already refuses rather than train an -# adapter that nothing reads; this does the same job without having to enumerate the -# families, by recording the read and asking afterwards. -# -# Every stash read goes through `take_moe_lora_stash`, which records the read on the experts -# module. The wrapper resets the record, calls the base layer, and checks it: an unread stash -# means this forward does not apply the expert LoRA, so the wrapper redoes the call through -# PEFT's own `ParamWrapper.forward`, which folds the delta into the weight. The verdict is -# cached per parameter name against the forward it was measured with, so the double call -# happens at most once per experts module and never for a family whose forward does read it. - -_MOE_LORA_STASH_READ_ATTR = "_unsloth_moe_lora_stash_read" -_MOE_LORA_STASH_VERDICT_ATTR = "_unsloth_moe_lora_forward_applies" -_MOE_LORA_COMPILED_VERDICT_ATTR = "_unsloth_moe_lora_forward_applies_compiled" - - -def moe_lora_stash_name(parameter_name: str) -> str: - """Attribute name `_patched_param_wrapper_forward` stashes `parameter_name`'s LoRA under.""" - # This Unsloth Zoo code section is licensed under AGPL3 - - return f"_unsloth_lora_{parameter_name}" - - -def _moe_module_dict(module, attr): - """Per-module bookkeeping dict for `attr`, created on first use. - - Read through `__dict__` and written with `setattr`: these are plain dicts, so - `nn.Module.__setattr__` stores them in the instance `__dict__`, and going straight to - `__dict__` on the read path skips `nn.Module.__getattr__` for a value that is never a - parameter, buffer or submodule. - """ - # This Unsloth Zoo code section is licensed under AGPL3 - - try: - store = module.__dict__.get(attr) - except AttributeError: - return None - if store is None: - store = {} - try: - setattr(module, attr, store) - except Exception: - return None - return store - - -def mark_moe_lora_stash_read(experts_module, parameter_name: str) -> None: - """Record that this experts forward applies `parameter_name`'s expert LoRA itself. - - Called by `take_moe_lora_stash` for every stash read, and by - `_get_lora_wrapper_for_param` for the forwards that reach past the stash and pull the - LoRA straight off the PEFT wrapper (the MXFP4 GPT-OSS experts forward does that). - """ - # This Unsloth Zoo code section is licensed under AGPL3 - - store = _moe_module_dict(experts_module, _MOE_LORA_STASH_READ_ATTR) - if store is not None: - store[parameter_name] = True - - -def take_moe_lora_stash(experts_module, parameter_name: str): - """The stashed LoRA for `parameter_name`, or None, recording that this forward read it. - - Records the read whatever the value is: an attempted read is what proves the forward - knows about the stash, and the stash is legitimately absent when no adapter is attached. - """ - # This Unsloth Zoo code section is licensed under AGPL3 - - mark_moe_lora_stash_read(experts_module, parameter_name) - return getattr(experts_module, moe_lora_stash_name(parameter_name), None) - - -def _resolve_experts_forward(experts_module): - """The function object that `experts_module(...)` will run, or None.""" - # This Unsloth Zoo code section is licensed under AGPL3 - - # Not getattr().__func__: Dynamo yields the bound method, so a traced verdict never matches. - try: - forward = experts_module.__dict__.get("forward") - except AttributeError: - forward = None - if forward is None: - forward = getattr(type(experts_module), "forward", None) - return getattr(forward, "__func__", forward) - - -def moe_lora_forward_applies_stash(experts_module, parameter_name: str, store_attr = None): - """Cached verdict: does this experts forward apply the stashed expert LoRA itself? - - True or False once measured, None when it has not been measured for the forward that is - currently installed. Re-patching the forward invalidates the verdict. - """ - # This Unsloth Zoo code section is licensed under AGPL3 - - try: - cache = experts_module.__dict__.get(store_attr or _MOE_LORA_STASH_VERDICT_ATTR) - except AttributeError: - return None - if not cache: - return None - entry = cache.get(parameter_name) - if entry is None: - return None - forward, verdict = entry - if forward is not _resolve_experts_forward(experts_module): - return None - return verdict - - -def _record_moe_lora_forward_verdict( - experts_module, parameter_name: str, verdict, store_attr = None, -) -> None: - """Remember `verdict` for `parameter_name` against the forward currently installed.""" - # This Unsloth Zoo code section is licensed under AGPL3 - - store = _moe_module_dict(experts_module, store_attr or _MOE_LORA_STASH_VERDICT_ATTR) - if store is not None: - store[parameter_name] = (_resolve_experts_forward(experts_module), bool(verdict)) - - -def _reset_moe_lora_stash_read(experts_module, parameter_name: str) -> None: - """Clear the read record for `parameter_name` before calling the experts forward.""" - # This Unsloth Zoo code section is licensed under AGPL3 - - store = _moe_module_dict(experts_module, _MOE_LORA_STASH_READ_ATTR) - if store is not None: - store[parameter_name] = False - - -def _moe_lora_stash_was_read(experts_module, parameter_name: str) -> bool: - """Whether the experts forward read `parameter_name`'s stash since the last reset.""" - # This Unsloth Zoo code section is licensed under AGPL3 - - try: - store = experts_module.__dict__.get(_MOE_LORA_STASH_READ_ATTR) - except AttributeError: - return False - return bool(store and store.get(parameter_name)) - - -_MOE_LORA_STASH_UNREAD_LOGGED = set() - - -def _log_moe_lora_stash_unread_once(experts_module, parameter_name: str) -> None: - """One message per experts class and parameter, the first time the stash goes unread.""" - # This Unsloth Zoo code section is licensed under AGPL3 - - key = (type(experts_module).__name__, parameter_name) - if key in _MOE_LORA_STASH_UNREAD_LOGGED: - return - _MOE_LORA_STASH_UNREAD_LOGGED.add(key) - _log_info( - f"Unsloth: {key[0]}.forward does not apply Unsloth's separated expert LoRA for " - f"{parameter_name}, so the {parameter_name} adapter is applied through PEFT's " - "parameter path instead." - ) - - -_STASH_READ_MARKERS = frozenset(("take_moe_lora_stash", "moe_lora_stash_name")) - -# How far the scan follows a call before giving up. The installed forwards delegate at most -# twice (dispatcher -> backend -> helper), and a bound keeps a pathological import graph -# from turning a cold start into a walk of the whole module tree. -_STASH_SCAN_MAX_DEPTH = 4 - - -# Resolved at import, absolute: the traced wrapper cannot import, and the compiled-cache copy has no package. -try: - from unsloth_zoo.temporary_patches.moe_experts_interface import ( - interface_route_reads_stash as _interface_route_reads_stash_impl, - ) -except Exception: - _interface_route_reads_stash_impl = None - - -def _interface_route_reads_stash(experts_module) -> bool: - """transformers' experts dispatch hides Unsloth's forward from the bytecode scan; ask the interface.""" - if _interface_route_reads_stash_impl is None: - return False - try: - return bool(_interface_route_reads_stash_impl(experts_module)) - except Exception: - return False - - -def _forward_statically_reads_stash(experts_module): - """Does this experts forward reach the stash API at all, read from its bytecode? - - A static answer, which is the only kind available while Dynamo is tracing: the probe - cannot run inside a captured graph without being captured with it. Every forward that - applies Unsloth's separated expert LoRA reaches `take_moe_lora_stash`, so that name - appearing anywhere the forward can call is what distinguishes an Unsloth-installed - forward from transformers' own. Cheap and side-effect free: no call is made, only - names are read. - - The reach matters. The forward installed for the Qwen MoE families IS - `forward_moe_backend`, a dispatcher that names only `forward_native_grouped_mm`, - `forward_triton_grouped_gemm` and `forward_native_moe_loop`, each of which reads the - stash. Scanning the dispatcher's own code object alone calls those supported families - stash-ignorant, which routes the first compiled call through PEFT's dynamic - parametrization and brings back the `fullgraph=True` failure this patch exists to - avoid. So a global name that resolves to a Python function is followed too, bounded by - `_STASH_SCAN_MAX_DEPTH` and by the visited set. - - Only module globals are resolvable: a name imported inside the function body, or - reached through an instance attribute, is not bound anywhere this can read. That makes - the answer one-sided, True is certain and False is only "no route found", which is why - the caller treats False as a reason to log rather than a guarantee. - - Returns True, False, or None when there is no code object to read. - """ - # This Unsloth Zoo code section is licensed under AGPL3 - - # Not getattr(module, "forward").__func__: torch 2.14 Dynamo then raises AttributeError('__globals__'). - try: - forward = experts_module.__dict__.get("forward") - except AttributeError: - forward = None - if forward is None: - forward = getattr(type(experts_module), "forward", None) - forward = getattr(forward, "__func__", forward) - if getattr(forward, "__code__", None) is None: - return None - return _code_reaches_stash(forward) - - -# Dynamo calls this eagerly and bakes in the bool: tracing the scan itself graph-breaks (torch 2.11) or -# raises under fullgraph once a forward's globals hold an lru_cache wrapper Dynamo cannot getattr through. -@_assume_constant_result -def _code_reaches_stash(forward): - # This Unsloth Zoo code section is licensed under AGPL3 - code = forward.__code__ - # (code, shallowest depth) pairs matched by `is`: id() is untraceable on some torch versions, - # equality merges same-source functions with different __globals__, and a plain visited set - # makes a helper first reached at the depth limit hide its deeper callees (hash-seed dependent). - seen_code = [] - pending = [(code, getattr(forward, "__globals__", {}), 0)] - - def _visited_at_or_above(target, depth): - for seen, seen_depth in seen_code: - if seen is target: - return seen_depth <= depth - return False - - while pending: - current, namespace, depth = pending.pop() - if _visited_at_or_above(current, depth): - continue - seen_code.append((current, depth)) - names = set(getattr(current, "co_names", ())) - if names & _STASH_READ_MARKERS: - return True - for constant in getattr(current, "co_consts", ()): - if hasattr(constant, "co_names"): - pending.append((constant, namespace, depth)) - if depth >= _STASH_SCAN_MAX_DEPTH: - continue - for name in names: - called = namespace.get(name) - called = getattr(called, "__func__", called) - called_code = getattr(called, "__code__", None) - if called_code is None: - continue - if _visited_at_or_above(called_code, depth + 1): - continue - pending.append((called_code, getattr(called, "__globals__", namespace), depth + 1)) - return False - - -@contextlib.contextmanager -def _preserved_rng_for_probe(x): - """Run the probe forward without advancing any generator the real forward will read. - - `no_grad` only turns off the graph; every random draw inside the throwaway forward - still advances the generator, so a family with dropout in the experts path would get - different numbers out of the call that counts than it got before this probe existed. - Gradient checkpointing makes that a wrong-gradient bug rather than a reproducibility - one: non-reentrant checkpointing restores the RNG state at the start of the region and - replays it, so the original pass (probe plus real forward) and the recompute (verdict - cached, real forward only) would draw different masks for the same region. - - Three things this is careful about, all of them ways a preservation attempt could be - worse than none: - - * The BACKEND is named, not inferred. `fork_rng`'s `devices` identifies devices within - `device_type`, which it resolves from `torch.accelerator.current_accelerator()` and - falls back to "cuda" for, so handing it an XPU device without saying so can ask the - wrong module for a generator state. - * A failure here must not become the probe's answer. The caller turns any exception - into `None`, which caches no verdict and leaves the stash path selected, so a - `fork_rng` that raised would silently keep the LoRA unapplied on a family that - ignores the stash. Anything that goes wrong setting the fork up leaves the forward - to run unforked instead. - Reached only on the eager path: `_measure_moe_lora_stash_read` returns before it while - Dynamo is tracing, so a generator-based context manager never enters a captured graph. - The check is kept here as well because this helper is usable on its own. - """ - if torch.compiler.is_compiling(): - yield - return - device = x.device if isinstance(x, torch.Tensor) else None - fork = None - try: - if device is None or device.type in ("cpu", "meta"): - # The CPU generator is forked whatever `devices` says; an empty list is what - # keeps `fork_rng` from initialising every visible accelerator, which it does - # when `devices` is None, and warns about. - fork = torch.random.fork_rng(devices = [], device_type = "cpu") - else: - fork = torch.random.fork_rng(devices = [device], device_type = device.type) - fork.__enter__() - except Exception: - fork = None - try: - yield - finally: - if fork is not None: - try: - fork.__exit__(None, None, None) - except Exception: - pass - - -# Disabled, not just `is_compiling()`-guarded: fullgraph=False compiles this frame alone, so the guard never cached a verdict. -@torch.compiler.disable -def _measure_moe_lora_stash_read(wrapper, experts_module, parameter_name, x, args, kwargs) -> bool: - """Run one throwaway experts forward under `no_grad` and report whether it read the stash. - - Measured on a separate graph-free call rather than on the call that counts. The call - that counts then runs exactly one experts forward, on whichever path the verdict names, - and it runs that same single path again on a recompute. That is the property - non-reentrant gradient checkpointing requires: it replays the region and refuses, with - `CheckpointError: a different number of tensors was saved during the original forward - and recomputation`, if the replay does not save what the forward saved. Measuring in - band cannot have that property, because the measuring call has to run the forward once - to find out and then a second time through PEFT when the answer is no, while the - recompute already knows the answer and runs it once. transformers 5 enables gradient - checkpointing with `use_reentrant = False` by default - (`modeling_utils.py`: `gradient_checkpointing_kwargs = {"use_reentrant": False}`), and - Unsloth's own vision path selects it whenever the run is distributed, so that is the - common configuration and not a corner. - - `no_grad` and the discarded result are what make this cheap enough to do eagerly: no - graph is built, nothing is saved for backward, and it happens once per experts module - and parameter for the lifetime of the process. A stash-reading family pays one extra - forward on its first step and nothing afterwards; a family that ignores the stash now - pays two forwards on its first step where the in-band measurement paid three. - - Returns True, False, or None when the probe itself raised. None caches nothing and - leaves the call on the stash path, so the real call raises the real error with the real - traceback instead of a second, confusing one from PEFT's fold. - """ - # This Unsloth Zoo code section is licensed under AGPL3 - - if torch.compiler.is_compiling(): - # Never inside a captured graph. Dynamo has no notion of a one-time measurement: it - # traces this whole throwaway forward into the graph alongside the call that counts, - # so the compiled region ends up running the experts forward three times on EVERY - # invocation, not once during a single warm-up. Measured on PyTorch 2.12 with - # fullgraph: six expert einsums in one graph where two are correct. - # - # None, which is the "no evidence" answer: it caches no verdict and leaves the call - # on the stash path, which is exactly what this module did before the probe existed. - # An eager forward before or after compilation still establishes the verdict, and the - # cached one is consulted before this function is ever called. - return None - lora_data = _extract_lora_from_wrapper(wrapper) - lora_attr = moe_lora_stash_name(parameter_name) - if lora_data is not None: - setattr(experts_module, lora_attr, lora_data) - _reset_moe_lora_stash_read(experts_module, parameter_name) - try: - with _preserved_rng_for_probe(x): - with torch.no_grad(): - wrapper.base_layer(x, *args, **kwargs) - except Exception: - # Not evidence either way, and not ours to report. Do not cache a verdict. - return None - finally: - if hasattr(experts_module, lora_attr): - delattr(experts_module, lora_attr) - was_read = _moe_lora_stash_was_read(experts_module, parameter_name) - _record_moe_lora_forward_verdict(experts_module, parameter_name, was_read) - return was_read - - -def _moe_lora_folded_weight(self, param, active): - """`W` with the active adapters folded in, by the cheaper of PEFT's two routes. - - PEFT's own choice, mirrored: with ONE active adapter over many experts it keeps the - low-rank factors and folds with a single `baddbmm`, and only otherwise materialises a - `get_delta_weight` the size of the whole expert stack. The difference is real on the - families this path exists for, OlmoE having 64 experts: the dense route allocates a - delta as large as the parameter and then a second tensor for the sum, so mirroring the - factored route halves the peak of every fold. `get_delta_factors` is newer than the - oldest PEFT with `target_parameters`, so it is asked for rather than assumed. - - The autocast is disabled around the fold for PEFT's reason: a parametrization may not - change the dtype of the parameter, so the arithmetic stays in W's dtype. Float8 never - arrives here, since `_can_fold_moe_lora_through_peft` refuses it, so PEFT's low - precision add has no counterpart. - """ - # This Unsloth Zoo code section is licensed under AGPL3 - - factors = getattr(self, "get_delta_factors", None) - if factors is not None and len(active) == 1 and getattr(self, "num_experts", 1) > 1: - try: - lhs, rhs, scaling = factors(active[0]) - with torch.autocast(device_type = param.device.type, enabled = False): - return torch.baddbmm(param, lhs, rhs, alpha = scaling) - except Exception: - pass - delta = None - try: - for name in active: - contribution = self.get_delta_weight(name) - delta = contribution if delta is None else delta + contribution - except Exception: - return None - return None if delta is None else param + delta - - -def _fold_moe_lora_without_parametrization( - self, immediate_base_layer, experts_module, parameter_name, x, args, kwargs -): - """PEFT's fold, arithmetically, with the parametrization left out. None if not doable. - - Same folded weight PEFT would produce, by the same route it would choose (see - `_moe_lora_folded_weight`). The difference is only how the base forward gets to see - it: PEFT registers a parametrization on the stored parameter, - whose `set_` Dynamo rejects as a graph-input mutation, while this swaps the attribute - for `W + delta` for the duration of the call and puts the parameter back after. - Autograd is unaffected, since the sum is an ordinary op on both tensors. - - Only for the compiled path. Eagerly PEFT's own forward stays in charge, so adapter - bookkeeping this does not reproduce (variants, merge state, anything future PEFT does - inside `_activate_lora`) is only ever bypassed where the alternative is not running at - all. Returns None whenever anything is missing, and the caller then falls through to - the stash path rather than losing the call. - - Nested wrappers compose: the wrapper chain is down_proj -> gate_up_proj -> experts and - each level folds and restores its own parameter around the next, so both are folded - for the innermost forward. - - The swap is done inside `_parameters` and nowhere else. `nn.Module.__getattr__` reads - the value straight out of that dict, so a plain tensor placed there is what the forward - sees, and the module's `__setattr__`, which would reject a plain tensor where a - Parameter is registered, is never involved. The obvious alternative, shadowing the - name in the instance `__dict__`, is what this used to do and it does not trace: - `experts_module.__dict__` is an unknown type to Dynamo, so removing the shadow in the - `finally` raises `Unsupported: Dynamo does not know how to trace method 'pop' of class - ''` on torch 2.10, the floor this package supports. That turned the - fullgraph compile this function exists to make possible back into a hard error. - """ - # This Unsloth Zoo code section is licensed under AGPL3 - - try: - active = [name for name in self.active_adapters if name in self.lora_A] - except Exception: - return None - if not active: - return None - - parameters = getattr(experts_module, "_parameters", None) - if not isinstance(parameters, dict) or parameter_name not in parameters: - return None - if parameter_name in experts_module.__dict__: - # Something already shadows the registered parameter, so swapping `_parameters` - # would not reach the forward. Refuse rather than fold into a value nothing reads. - return None - original = parameters[parameter_name] - folded = _moe_lora_folded_weight(self, original, active) - if folded is None: - return None - - parameters[parameter_name] = folded - try: - return immediate_base_layer(x, *args, **kwargs) - finally: - parameters[parameter_name] = original - - -# PEFT names the float8 storage dtypes it must upcast in `peft.utils.UPCAST_DTYPES`. Track -# that list rather than restating it, so a dtype PEFT adds later is excluded here too, and -# fall back to resolving the names off torch for PEFT versions that do not export it. Every -# one of these raises on `param + delta`, including the fnuz pair ROCm uses. -def _resolve_float8_storage_dtypes(): - names = None - try: - from peft.utils import UPCAST_DTYPES as names - except Exception: - names = ( - "float8_e4m3fn", "float8_e4m3fnuz", - "float8_e5m2", "float8_e5m2fnuz", - "float8_e8m0fnu", - ) - resolved = [] - for name in names: - dtype = getattr(torch, name, None) - # Older torch does not define every float8 variant. - if isinstance(dtype, torch.dtype): - resolved.append(dtype) - return tuple(resolved) - - -_FLOAT8_STORAGE_DTYPES = _resolve_float8_storage_dtypes() - - -def _can_fold_moe_lora_through_peft(experts_module, parameter_name: str) -> bool: - """Whether handing this parameter back to PEFT would fold a delta into a real weight. - - PEFT folds by registering a parametrization that adds `delta_weight` to the stored - parameter, which is only defined for a plain floating-point tensor. A stacked expert - weight held as bitsandbytes `Params4bit`, an MXFP4 blocked tensor or an FP8 tensor is - none of those: the add is a dtype or shape error at best and silent corruption at - worst. `_forward_native_fp8_expert_loop` says the same thing from the other side, that - the reroute must not take charge of a quantized parameter, and enforces it by recording - its reads; this enforces it for the quantized forwards that have no such hook, the - MXFP4 GPT-OSS experts forward among them, where `_get_lora_wrapper_for_param` resolves - to None against PEFT's `target_parameters` layout and so records nothing. - - Refusing to reroute leaves that case exactly as it is on main: the stash is written and - the forward does what it does with it. That is not a fix for a quantized family, and it - is not meant to be one; it is a guarantee that this change cannot make one worse. - """ - # This Unsloth Zoo code section is licensed under AGPL3 - - if _original_param_wrapper_forward is None: - return False - param = getattr(experts_module, parameter_name, None) - if not isinstance(param, torch.Tensor): - return False - if not param.dtype.is_floating_point: - return False - # `is_floating_point` is True for every float8 storage dtype, and PEFT's own - # `get_delta_weight` special-cases them precisely because the add is not defined there: - # `param + delta` raises "Promotion for Float8 Types is not supported" for all five, and - # naming only the two CUDA ones let the fnuz pair that ROCm uses through. - if param.dtype in _FLOAT8_STORAGE_DTYPES: - return False - # bitsandbytes stores 4-bit weights in a float or uint8 tensor of the packed shape, and - # marks them with `quant_state`. Neither the dtype nor the shape gives it away. - if getattr(param, "quant_state", None) is not None: - return False - if HAS_BNB and Params4bit is not None and isinstance(param, Params4bit): - return False - if is_mxfp4_expert_param(param): - return False - return True - - -# Store original ParamWrapper.forward for fallback -_original_param_wrapper_forward = None - - - -def _gpt_oss_routed_wrapper_forward(wrapper, experts_module, x, args, kwargs): - """gpt-oss eval decode through the outermost expert ParamWrapper: only the routed experts, - with every adapter in the chain applied. None for anything else (callers run as before).""" - if torch.is_grad_enabled() or type(experts_module).__name__ != "GptOssExperts": - return None - if "adapter_names" in kwargs: # mixed-adapter batch: PEFT's own forward handles or rejects it - return None - if getattr(wrapper, "_unsloth_inner_expert_wrapper", False): - return None - inner = wrapper.base_layer - while hasattr(inner, "base_layer"): - # An inner wrapper sees only part of the adapter chain, so it must never route. - inner._unsloth_inner_expert_wrapper = True - inner = inner.base_layer - router_indices = kwargs.get("router_indices", args[0] if len(args) > 0 else None) - routing_weights = kwargs.get("routing_weights", args[1] if len(args) > 1 else None) - if router_indices is None or routing_weights is None: - return None - from unsloth_zoo.temporary_patches.gpt_oss_routed import ROUTED_MAX_SLOTS, routed_bf16_eligible, routed_bf16_forward - if router_indices.numel() > ROUTED_MAX_SLOTS or not routed_bf16_eligible(wrapper, x): - return None - return routed_bf16_forward(wrapper, x, router_indices, routing_weights) - -def _patched_param_wrapper_forward( - self, x: torch.Tensor, *args, **kwargs -) -> torch.Tensor: - """Patched ParamWrapper.forward for MoE separated LoRA. - - For MoE experts: bypass PEFT's _activate_lora and stash LoRA data by - parameter_name for forward_native_grouped_mm. For non-MoE: original forward. - """ - # This Unsloth Zoo code section is licensed under AGPL3 - - # Use self.base_layer (immediate parent), NOT get_base_layer() which recurses - # to the deepest layer; the wrapper chain down_proj -> gate_up_proj -> - # Qwen3MoeExperts must be preserved. - immediate_base_layer = self.base_layer - - # For stashing LoRA data we need the actual experts module (recursive lookup). - experts_module = self.get_base_layer() - - param_name = getattr(self, "parameter_name", None) - - routed = _gpt_oss_routed_wrapper_forward(self, experts_module, x, args, kwargs) - if routed is not None: - return routed - - if _wrapper_uses_separated_moe_lora(self, experts_module): - # MoE experts: bypass PEFT's _activate_lora, use separated computation. - if self.disable_adapters: - if self.merged: - self.unmerge() - return immediate_base_layer(x, *args, **kwargs) - - if self.merged: - return immediate_base_layer(x, *args, **kwargs) - - # Ensure wrapper.num_experts is set for LoRA weight reshaping. - if not hasattr(self, "num_experts"): - if hasattr(experts_module, "num_experts"): - self.num_experts = experts_module.num_experts - elif hasattr(experts_module, param_name): - p = getattr(experts_module, param_name) - if hasattr(p, "shape") and len(p.shape) >= 1: - self.num_experts = p.shape[0] - - # An experts forward that does not read the stash never sees this LoRA, so hand the - # wrapper back to PEFT, which folds the delta into the weight instead. Measured once - # per experts module and parameter, before the call that counts, so the call that - # counts always runs exactly one experts forward on the path the verdict names. - applies_stash = moe_lora_forward_applies_stash(experts_module, param_name) - if applies_stash is None and torch.compiler.is_compiling(): - # First invocation is a compiled one, so there is no verdict and no way to - # measure one: the probe cannot run inside a captured graph without being - # captured with it and re-run on every call. Answer statically instead, from - # whether the forward references the stash API at all. - # - # Both wrong answers are costly, which is why this is not a fixed assumption. - # Assuming unread sends a supported stash-reading family into PEFT's own - # ParamWrapper.forward, which registers and removes a parametrization while - # tracing and hard-fails under fullgraph with "Getting an inplace view on a - # graph input is not supported". Assuming read leaves the expert LoRA out of - # every output and gradient on a family that ignores the stash, silently and - # for the life of the captured graph. The bytecode says which family this is. - # - # Kept apart from the verdict (the first eager call still measures); the marker reads it. - applies_stash = _forward_statically_reads_stash(experts_module) - if not applies_stash and _interface_route_reads_stash(experts_module): - applies_stash = True - if applies_stash is not None: - _record_moe_lora_forward_verdict( - experts_module, param_name, applies_stash, _MOE_LORA_COMPILED_VERDICT_ATTR, - ) - if applies_stash is False: - _log_moe_lora_stash_unread_once(experts_module, param_name) - elif applies_stash is None: - applies_stash = _measure_moe_lora_stash_read( - self, experts_module, param_name, x, args, kwargs - ) - if applies_stash is False: - _log_moe_lora_stash_unread_once(experts_module, param_name) - if applies_stash is False and _can_fold_moe_lora_through_peft(experts_module, param_name): - if torch.compiler.is_compiling(): - # PEFT's own fold cannot be traced. `_activate_lora` registers a - # parametrization for the call and removes it after, and the `set_` that - # registration performs on the stored parameter is a graph-input mutation, - # so Dynamo refuses it: "Getting an inplace view on a graph input is not - # supported". Under fullgraph that is a hard error, which would mean these - # families cannot compile at all. Folding the same delta ourselves keeps - # the arithmetic and drops the parametrization, and traces. - folded = _fold_moe_lora_without_parametrization( - self, immediate_base_layer, experts_module, param_name, x, args, kwargs - ) - if folded is not None: - return folded - else: - return _original_param_wrapper_forward(self, x, *args, **kwargs) - - # Extract LoRA for this parameter and stash on the experts module - # (not base_layer): _unsloth_lora_gate_up_proj / _unsloth_lora_down_proj. - lora_data = _extract_lora_from_wrapper(self) - - if lora_data is not None and param_name: - lora_attr = moe_lora_stash_name(param_name) - setattr(experts_module, lora_attr, lora_data) - - try: - # Immediate base_layer preserves the wrapper chain. - result = immediate_base_layer(x, *args, **kwargs) - finally: - if param_name: - lora_attr = moe_lora_stash_name(param_name) - if hasattr(experts_module, lora_attr): - delattr(experts_module, lora_attr) - - return result - - # Unclaimed stacks (NemotronH up_proj / down_proj): PEFT's parametrization graph-breaks under compile. - if ( - torch.compiler.is_compiling() - and param_name - and not self.disable_adapters - and not self.merged - and getattr(getattr(experts_module, param_name, None), "ndim", 0) == 3 - and _can_fold_moe_lora_through_peft(experts_module, param_name) - ): - folded = _fold_moe_lora_without_parametrization( - self, immediate_base_layer, experts_module, param_name, x, args, kwargs - ) - if folded is not None: - return folded - - # Non-MoE: original PEFT forward with _activate_lora. - return _original_param_wrapper_forward(self, x, *args, **kwargs) - - -# Store original ParamWrapper.get_delta_weight for fallback -_original_param_wrapper_get_delta_weight = None - - -def _grouped_by_expert_delta_weight(wrapper, adapter_name): - """PEFT's ParamWrapper.get_delta_weight arithmetic with lora_B read the way the - separated MoE forward reads it: (out, num_experts, rank) rather than PEFT's - (out, rank, num_experts). Everything else, including which einsum the swapped - in/out orientation needs, is PEFT's own.""" - weight_A = wrapper.lora_A[adapter_name].weight - weight_B = wrapper.lora_B[adapter_name].weight - num_experts = int(wrapper.num_experts) - - # experts x rank x in_features, as in PEFT: lora_A is grouped by expert there too. - weight_A = weight_A.reshape(num_experts, -1, weight_A.shape[-1]) - # out_features x experts x rank. This is the one line that differs from PEFT. - weight_B = weight_B.reshape(weight_B.shape[0], num_experts, -1) - - scaling = wrapper.scaling[adapter_name] - if not getattr(wrapper, "_did_swap_in_out_features", False): - return torch.einsum("o e r, e r i -> e i o", weight_B, weight_A) * scaling - # for some MoE layers, the order is (experts, out_features, in_features) - return torch.einsum("o e r, e r i -> e o i", weight_B, weight_A) * scaling - - -def _cast_delta_weight_like_param(delta_weight, param): - """PEFT's own tail: move the delta to the parameter, and match its dtype unless that - dtype is a low precision one PEFT deliberately adds in higher precision (float8 and - friends). Probed from PEFT when it publishes the set, so this follows PEFT rather - than hardcoding a dtype list that a new release would make wrong.""" - try: - from peft.tuners.lora.layer import ALLOWED_COMPUTE_DTYPES - allowed = param.dtype in ALLOWED_COMPUTE_DTYPES - except Exception: - # PEFT 0.18 has no such set and casts unconditionally, 0.19.0 onwards do; every dtype - # 0.18 accepts for a 3D expert parameter is an ordinary floating point one. - allowed = param.dtype in (torch.float16, torch.bfloat16, torch.float32, torch.float64) - if allowed: - return delta_weight.to(param.device, param.dtype) - return delta_weight.to(param.device) - - -def _patched_param_wrapper_get_delta_weight(self, adapter_name, *args, **kwargs): - """PEFT's `get_delta_weight` for a legacy, grouped-by-expert fused MoE expert LoRA. - - PEFT reconstructs a fused expert delta with lora_B packed rank-major, which is what - it trained and what every consumer reads, so by default this does nothing at all and - PEFT's own code runs. It only takes over when the caller has explicitly declared that - the adapters in this process were packed the pre-fix way, with - `UNSLOTH_MOE_LORA_B_LAYOUT=grouped_by_expert`; that same variable is what makes the - separated forward read them that way (unsloth_zoo#1269), and without this patch the - forward and PEFT's `merge_and_unload` / `merge_adapter` / `unmerge` would disagree for - exactly the checkpoints that switch exists to support (unsloth#6930). - - Gating on the declaration rather than on "is this wrapper routed to the separated - forward" is the point: the packing of a stored adapter is a property of the adapter, - and reading it off this process's routing would scramble a correctly packed adapter - merged inside an Unsloth process.""" - # This Unsloth Zoo code section is licensed under AGPL3 - - if not _legacy_lora_b_layout_requested(): - return _original_param_wrapper_get_delta_weight(self, adapter_name, *args, **kwargs) - if moe_lora_b_layout_for_wrapper(self, adapter_name) != LORA_B_LAYOUT_GROUPED_BY_EXPERT: - return _original_param_wrapper_get_delta_weight(self, adapter_name, *args, **kwargs) - - delta_weight = _grouped_by_expert_delta_weight(self, adapter_name) - return _cast_delta_weight_like_param(delta_weight, self.get_param()) - - -def _moe_utils_logger(): - """The shared patch logger, imported lazily and absolutely: this file is also copied - into unsloth_compiled_cache and executed there as a top level module, where a - relative import has no package to resolve against.""" - try: - from unsloth_zoo.temporary_patches.common import logger - return logger - except Exception: - return None - - -# The adapter_config.json keys. The flat one is what a converter branches on -# (unsloth#6930 asks for exactly this key); the nested one carries the detail, because a -# checkpoint can in principle hold one fused parameter trained through the separated -# forward and another PEFT kept for itself. -FUSED_EXPERT_LORA_LAYOUT_KEY = "lora_B_layout" -FUSED_EXPERT_LORA_DETAIL_KEY = "unsloth_fused_expert_lora" - - -def _default_adapter_name(peft_model) -> str: - """The adapter `fused_expert_lora_layout` describes when the caller names none: the - model's active one, or PEFT's own default name. Both attributes are properties on - some PEFT versions and can raise on a partly built model, so neither is trusted.""" - # delete_adapter leaves active_adapter a list, and active_adapters is always one. - for attribute in ("active_adapter", "active_adapters"): - try: - active = getattr(peft_model, attribute, None) - except Exception: - continue - if isinstance(active, str) and active: - return active - if isinstance(active, (list, tuple)) and len(active) != 0: - if isinstance(active[0], str) and active[0]: - return active[0] - return "default" - - -def fused_expert_lora_layout(peft_model, adapter_name = None) -> Optional[dict]: - """How ONE adapter's fused MoE expert LoRA on `peft_model` packs its lora_A and - lora_B, or None when that adapter has no fused expert LoRA at all. - - lora_A is grouped by expert in both stacks. lora_B is the one that differs, and it is - reported per PEFT parameter name, since that is what decides whether Unsloth's - separated forward or PEFT's own forward owns the wrapper. - - The answer is per adapter because a PeftModel can hold several. A fused expert - adapter and a dense adapter side by side share the fused expert wrappers, so a scan - that only asks "does this model have a fused expert LoRA anywhere" answers yes for - the dense adapter too, and that answer would be written into the dense adapter's own - adapter_config.json. `adapter_name` defaults to the model's active adapter, so the - single adapter case is unchanged.""" - if adapter_name is None: - adapter_name = _default_adapter_name(peft_model) - - parameters = {} - try: - modules = list(peft_model.modules()) - except Exception: - return None - - for module in modules: - parameter_name = getattr(module, "parameter_name", None) - if not parameter_name or not hasattr(module, "lora_A"): - continue - if not _wrapper_has_adapter(module, adapter_name): - # A wrapper another adapter owns. Its packing is not this adapter's business. - continue - num_experts = int(getattr(module, "num_experts", 1) or 1) - if num_experts <= 1: - # Not a fused expert stack, so there is no expert axis to pack. - continue - layout = moe_lora_b_layout_for_wrapper(module, adapter_name) - entry = parameters.get(parameter_name) - if entry is None: - parameters[parameter_name] = { - "lora_B_layout": layout, - "num_experts": num_experts, - } - elif entry["lora_B_layout"] != layout: - entry["lora_B_layout"] = "mixed" - - if not parameters: - return None - - detail = { - "lora_A_layout": LORA_B_LAYOUT_GROUPED_BY_EXPERT, - "parameters": parameters, - } - layouts = {entry["lora_B_layout"] for entry in parameters.values()} - # The flat key is what a converter branches on, so it is only written when every - # fused expert parameter agrees on one of the two real names. "mixed" is a report, - # not a layout, and a converter testing `== "grouped_by_expert"` would read it as - # rank_major, so it stays in the nested detail where it has to be read deliberately. - if len(layouts) == 1: - layout = layouts.pop() - if layout in (LORA_B_LAYOUT_GROUPED_BY_EXPERT, LORA_B_LAYOUT_RANK_MAJOR): - detail["lora_B_layout"] = layout - return detail - - -def _fused_expert_lora_adapter_config_paths(peft_model, save_directory, selected_adapters): - """Where PEFT just wrote an adapter_config.json: the root for `default`, a - subdirectory named after any other adapter. Paired with the adapter each file - belongs to, because the marker is per adapter.""" - if selected_adapters is None: - try: - selected_adapters = list(peft_model.peft_config.keys()) - except Exception: - selected_adapters = ["default"] - paths = [] - for adapter_name in selected_adapters: - directory = save_directory - if adapter_name != "default": - directory = os.path.join(save_directory, adapter_name) - path = os.path.join(directory, "adapter_config.json") - if os.path.isfile(path): - paths.append((adapter_name, path)) - return paths - - -def write_fused_expert_lora_layout(peft_model, save_directory, selected_adapters = None): - """Record the fused MoE expert lora_B packing in the adapter_config.json PEFT has - just written, and return the files updated. PEFT ignores keys it does not know, so - the checkpoint stays loadable by any PEFT version; this only stops a downstream - converter, or PEFT's own merge in a differently configured process, from having to - guess (unsloth#6930). - - The packing is recomputed for each adapter being saved, and an adapter with no fused - expert LoRA is left exactly as PEFT wrote it: not written, and not stripped either, - since a key this never put there is not this function's to remove. A model holding a - fused expert adapter and a dense one would otherwise describe the fused adapter's - experts in the dense adapter's config, which is worse than saying nothing.""" - written = [] - for adapter_name, path in _fused_expert_lora_adapter_config_paths( - peft_model, save_directory, selected_adapters, - ): - detail = fused_expert_lora_layout(peft_model, adapter_name) - if detail is None: - continue - with open(path, "r", encoding = "utf-8") as f: - config = json.load(f) - if not isinstance(config, dict): - continue - config[FUSED_EXPERT_LORA_DETAIL_KEY] = detail - if "lora_B_layout" in detail: - config[FUSED_EXPERT_LORA_LAYOUT_KEY] = detail["lora_B_layout"] - elif config.get(FUSED_EXPERT_LORA_LAYOUT_KEY) in ( - LORA_B_LAYOUT_GROUPED_BY_EXPERT, LORA_B_LAYOUT_RANK_MAJOR, - ): - # An earlier save of this adapter wrote a flat layout and the parameters no - # longer agree on one, so that key is now false. Only a value this could have - # written is dropped: anything else is somebody else's key. - config.pop(FUSED_EXPERT_LORA_LAYOUT_KEY, None) - # Byte for byte how PeftConfigMixin.save_pretrained writes it, so adding the marker - # does not also reformat the file PEFT produced. - _atomic_write_text(path, json.dumps(config, indent = 2, sort_keys = True)) - written.append(path) - return written - - -_LORA_INTERNAL_NAMES = frozenset(( - "base_layer", "lora_A", "lora_B", "lora_dropout", "lora_embedding_A", - "lora_embedding_B", "lora_magnitude_vector", -)) - - -def _lora_wrapped_module_names(base_model, adapter_name): - """(module names wrapped by `adapter_name`, whether it also has a fused expert parameter LoRA).""" - wrapped, has_param_lora = [], False - for name, module in base_model.named_modules(): - # nn.Embedding LoRA lives in lora_embedding_A, its lora_A stays empty. - if not any( - hasattr(d, "keys") and adapter_name in d - for d in (getattr(module, "lora_A", None), getattr(module, "lora_embedding_A", None)) - ): - continue - if getattr(module, "parameter_name", None): - has_param_lora = True - continue - if hasattr(module, "base_layer"): - wrapped.append(name) - return wrapped, has_param_lora - - -def _target_regex_for(names, all_names): - """Fullmatch regex selecting exactly `names`; layer indices collapse to \\d+ only if that adds nothing.""" - import re - wanted = set(names) - collapsed = sorted({re.sub(r"\\\.\d+(?=\\\.|$)", r"\\.\\d+", re.escape(n)) for n in wanted}) - regex = "(?:" + "|".join(collapsed) + ")" - if {n for n in all_names if re.fullmatch(regex, n)} == wanted: - return regex - return "(?:" + "|".join(re.escape(n) for n in sorted(wanted)) + ")" - - -# Pre-fusion expert names PEFT's v5 MoE conversion probes under a fused expert container. -_PRE_FUSION_EXPERT_NAMES = ("gate_proj", "up_proj", "down_proj", "w1", "w2", "w3") - - -def _all_target_names(base_model): - return [ - n for n, _ in base_model.named_modules() - if n and not (set(n.split(".")) & _LORA_INTERNAL_NAMES) - ] - - -def _portable_regex_target_modules(base_model, adapter_name, regex): - """Exact regex over the wrapped modules when `regex` also fullmatches `.gate_proj`. - - Unsloth's get_peft_regex matches `mlp.experts.gate_proj` / `up_proj`, which PEFT's v5 MoE conversion - reads as a v4 adapter and doubles gate_up_proj's rank, so plain PEFT fails to load (unsloth#6930). - """ - import re - wrapped, has_param_lora = _lora_wrapped_module_names(base_model, adapter_name) - if not has_param_lora: - return None - containers = set() - for name, module in base_model.named_modules(): - if getattr(module, "parameter_name", None): - while name.endswith(".base_layer"): - name = name[: -len(".base_layer")] - containers.add(name) - try: - if not any( - re.fullmatch(regex, f"{container}.{leaf}") - for container in containers for leaf in _PRE_FUSION_EXPERT_NAMES - ): - return None - except re.error: - return None - if not wrapped: - return [] - return _target_regex_for(wrapped, _all_target_names(base_model)) - - -def portable_lora_target_modules(peft_model, adapter_name, target_modules): - """`target_modules` that plain PEFT loads as Unsloth trained, or None if already fine. - - Gemma 4: LoRA sits on the inner `.linear`; plain PEFT would pick the unsupported wrapper, so write a regex. - Fused experts: unmatched gate_proj / up_proj make PEFT's v5 MoE conversion double gate_up_proj's rank; drop them. - """ - if not target_modules or not isinstance(target_modules, (str, list, tuple, set)): - return None - try: - base_model = peft_model.get_base_model() - except Exception: - base_model = getattr(getattr(peft_model, "base_model", None), "model", None) - if base_model is None: - return None - if isinstance(target_modules, str): - return _portable_regex_target_modules(base_model, adapter_name, target_modules) - wrapped, has_param_lora = _lora_wrapped_module_names(base_model, adapter_name) - targets = [t for t in target_modules if isinstance(t, str)] - if len(targets) != len(target_modules): - return None - - def hit(key, target): - return key == target or key.endswith("." + target) - - redirected = [ - n for n in wrapped - if n.endswith(".linear") and any(hit(n[: -len(".linear")], t) for t in targets) - ] - if redirected: - return _target_regex_for(wrapped, _all_target_names(base_model)) - if has_param_lora: - kept = [t for t in targets if any(hit(n, t) for n in wrapped)] - if len(kept) != len(targets): - return sorted(kept) - return None - - -def _redirected_patterns(peft_model, adapter_name, patterns): - """Add a `(?:key)\\.linear` twin for each rank/alpha pattern key that matched a redirected - wrapper, since PEFT fullmatches `(.*\\.)?(key)$` against the inner `.linear` on reload.""" - import re - if not isinstance(patterns, dict) or not patterns: - return patterns - wrapped, _ = _lora_wrapped_module_names(peft_model.get_base_model(), adapter_name) - inner = [n[: -len(".linear")] for n in wrapped if n.endswith(".linear")] - out = dict(patterns) - for key, value in patterns.items(): - twin = f"(?:{key})\\.linear" - if twin not in out and any(re.match(rf"(.*\.)?({key})$", n) for n in inner): - out[twin] = value - return out - - -def write_portable_target_modules(peft_model, save_directory, selected_adapters = None): - """Rewrite saved `target_modules` per `portable_lora_target_modules`; returns paths written.""" - written = [] - for adapter_name, path in _fused_expert_lora_adapter_config_paths( - peft_model, save_directory, selected_adapters, - ): - with open(path, "r", encoding = "utf-8") as f: - config = json.load(f) - if not isinstance(config, dict): - continue - replacement = portable_lora_target_modules( - peft_model, adapter_name, config.get("target_modules"), - ) - if replacement is None: - continue - config["target_modules"] = replacement - if isinstance(replacement, str): - for field in ("rank_pattern", "alpha_pattern"): - config[field] = _redirected_patterns(peft_model, adapter_name, config.get(field)) - # PEFT rejects layer selectors beside a str target; the regex already names only those layers. - config["layers_to_transform"] = None - config["layers_pattern"] = None - _atomic_write_text(path, json.dumps(config, indent = 2, sort_keys = True)) - written.append(path) - return written - - -def _atomic_write_text(path, text): - """Replace `path` with `text`, or leave it exactly as it was. - - Opening the real file "w" truncates it before the first byte is written, so a write - that fails part way (a full disk, an erroring network mount) leaves a truncated - adapter_config.json behind. `_patched_peft_model_save_pretrained` then swallows the - exception and reports a successful save, so the checkpoint is unloadable and nothing - says so. Writing a sibling temporary file first means a failure destroys only that - file, and `os.replace` is atomic on POSIX and on Windows, so no reader can observe a - half-written config either. - - The temporary file is a sibling, not a tempdir entry, because `os.replace` across - filesystems raises. It is removed on failure so a failed save leaves no litter.""" - # This Unsloth Zoo code section is licensed under AGPL3 - - directory = os.path.dirname(path) or "." - handle, temporary = tempfile.mkstemp( - dir = directory, prefix = os.path.basename(path) + ".", suffix = ".tmp", - ) - try: - # mkstemp creates the file 0600 and os.replace keeps the NEW inode's mode, so - # without this the config silently drops from (say) 0644 to 0600 on every save - # and a checkpoint shared with other users or a serving process stops being - # readable, while its weight files stay readable. Carry the existing mode over; - # for a config PEFT has not written yet there is nothing to carry, so leave - # mkstemp's private mode rather than inventing a laxer one. - try: - os.chmod(temporary, stat.S_IMODE(os.stat(path).st_mode)) - except OSError: - pass - with os.fdopen(handle, "w", encoding = "utf-8") as f: - f.write(text) - f.flush() - os.fsync(f.fileno()) - os.replace(temporary, path) - except BaseException: - try: - os.unlink(temporary) - except OSError: - pass - raise - - -# Store original PeftModel.save_pretrained for fallback -_original_peft_model_save_pretrained = None - - -def _save_pretrained_argument(name, default, args, kwargs): - """One of PeftModel.save_pretrained's arguments, whether the caller passed it by - keyword or positionally. Resolved against the real signature rather than a copy of - it, so a PEFT release that adds an argument cannot silently shift the index. `args` - excludes self and save_directory.""" - if name in kwargs: - return kwargs[name] - try: - import inspect - parameters = list( - inspect.signature(_original_peft_model_save_pretrained).parameters - ) - index = parameters.index(name) - 2 - if 0 <= index < len(args): - return args[index] - except Exception: - pass - return default - - -def _stamp_fused_expert_lora_layout(peft_model, selected_adapters): - """Put the layout in `auto_mapping` too, which PEFT writes with the config itself: the - post-save marker write can fail, and a marker-less Unsloth adapter reads as pre-2026.9.6.""" - try: - names = selected_adapters or list(peft_model.peft_config.keys()) - for adapter_name in names: - auto_mapping = getattr(peft_model.peft_config.get(adapter_name), "auto_mapping", None) - if not isinstance(auto_mapping, dict): - continue - detail = fused_expert_lora_layout(peft_model, adapter_name) - auto_mapping.pop(FUSED_EXPERT_LORA_LAYOUT_KEY, None) - auto_mapping.pop(FUSED_EXPERT_LORA_DETAIL_KEY, None) - if detail is not None: - auto_mapping[FUSED_EXPERT_LORA_DETAIL_KEY] = detail - except Exception: - pass - - -def _patched_peft_model_save_pretrained(self, save_directory, *args, **kwargs): - """Add the fused expert layout marker to what PEFT saved. Never fails a save: the - adapter itself is already on disk by the time this runs.""" - # This Unsloth Zoo code section is licensed under AGPL3 - - _stamp_fused_expert_lora_layout( - self, _save_pretrained_argument("selected_adapters", None, args, kwargs), - ) - result = _original_peft_model_save_pretrained(self, save_directory, *args, **kwargs) - if _save_pretrained_argument("is_main_process", True, args, kwargs): - try: - write_fused_expert_lora_layout( - self, save_directory, - selected_adapters = _save_pretrained_argument( - "selected_adapters", None, args, kwargs, - ), - ) - except Exception as exception: - logger = _moe_utils_logger() - if logger is not None: - logger.warning( - f"Unsloth: could not record the fused MoE expert LoRA layout in " - f"adapter_config.json ({type(exception).__name__}: {exception}). The " - f"adapter itself saved correctly." - ) - try: - write_portable_target_modules( - self, save_directory, - selected_adapters = _save_pretrained_argument( - "selected_adapters", None, args, kwargs, - ), - ) - except Exception as exception: - logger = _moe_utils_logger() - if logger is not None: - logger.warning( - f"Unsloth: could not rewrite target_modules for plain PEFT in " - f"adapter_config.json ({type(exception).__name__}: {exception}). The " - f"adapter itself saved correctly." - ) - return result - - -def _patch_peft_save_pretrained_for_moe_layout(): - # This Unsloth Zoo code section is licensed under AGPL3 - - global _original_peft_model_save_pretrained - - try: - from peft import PeftModel - except Exception: - return False - - if getattr(PeftModel.save_pretrained, "_unsloth_moe_layout_patched", False): - return True - - if _original_peft_model_save_pretrained is None: - _original_peft_model_save_pretrained = PeftModel.save_pretrained - - patched = wraps(_original_peft_model_save_pretrained)( - _patched_peft_model_save_pretrained, - ) - patched._unsloth_moe_layout_patched = True - PeftModel.save_pretrained = patched - return True - - -# Families whose experts forward read lora_B grouped by expert from #450 to #1269 (2026.9.6); -# every other family stayed on PEFT's forward, i.e. rank_major. -_LEGACY_GROUPED_LORA_B_FAMILIES = frozenset(( - "deepseek_v3", "ernie4_5_moe", "ernie4_5_vl_moe", "gemma4", "glm4_moe", "glm4_moe_lite", - "gpt_oss", "lfm2_moe", "mixtral", "qwen3_5_moe", "qwen3_moe", "qwen3_next", "qwen3_vl_moe", -)) -_legacy_lora_b_converted_logged = set() -_original_peft_model_load_adapter = None - - -def _experts_family(module): - """`qwen3_moe` for Qwen3MoeExperts from `transformers.models.qwen3_moe.modeling_qwen3_moe` - or from Unsloth's compiled copy, `unsloth_compiled_module_qwen3_moe`.""" - if type(module).__name__ == "Mxfp4GptOssExperts": - return "gpt_oss" - name = type(module).__module__ - leaf = name.rsplit(".", 1)[-1] - if leaf.startswith("unsloth_compiled_module_"): - return leaf[len("unsloth_compiled_module_"):] - parts = name.split(".") - if len(parts) > 2 and parts[:2] == ["transformers", "models"]: - return parts[2] - return None - - -def _declared_lora_b_layouts(source) -> Optional[dict]: - detail = source.get(FUSED_EXPERT_LORA_DETAIL_KEY) - if isinstance(detail, dict) and isinstance(detail.get("parameters"), dict): - return { - name: entry.get("lora_B_layout") - for name, entry in detail["parameters"].items() - if isinstance(entry, dict) and entry.get("lora_B_layout") in _LORA_B_LAYOUTS - } - marker = source.get(FUSED_EXPERT_LORA_LAYOUT_KEY) - if marker in _LORA_B_LAYOUTS: - return {name: marker for name in _SEPARATED_MOE_LORA_PARAMETER_NAMES} - return None - - -def saved_lora_b_layouts(config) -> dict: - """{parameter_name: layout} from the #1232 marker (top level, or the copy in - `auto_mapping`), else {None: grouped_by_expert} for a marker-less Unsloth save - (`unsloth_fixed`) with fused expert targets: Unsloth only wrote `target_parameters` - from #450 (2026-02-05), so that is the 2026.9.5-or-older window.""" - if not isinstance(config, dict): - return {} - auto_mapping = config.get("auto_mapping") - auto_mapping = auto_mapping if isinstance(auto_mapping, dict) else {} - for source in (config, auto_mapping): - declared = _declared_lora_b_layouts(source) - if declared is not None: - return declared - # UNSLOTH_MOE_LORA_MERGED=1 trained on PEFT's own forward (rank_major); a process - # still set that way most likely loads such an adapter. - if config.get("target_parameters") and auto_mapping.get("unsloth_fixed") and _should_use_separated_lora(): - return {None: LORA_B_LAYOUT_GROUPED_BY_EXPERT} - return {} - - -def convert_legacy_lora_b(peft_model, adapter_name, config) -> int: - """Repack `adapter_name`'s grouped-by-expert fused expert lora_B into PEFT's rank_major, - in place, per `saved_lora_b_layouts(config)`. Lossless: column `e * r + j` moves to - `j * num_experts + e`. Returns how many tensors moved.""" - # This Unsloth Zoo code section is licensed under AGPL3 - - layouts = saved_lora_b_layouts(config) - if not layouts: - return 0 - targets = [] - for module in peft_model.modules(): - parameter_name = getattr(module, "parameter_name", None) - if parameter_name not in _SEPARATED_MOE_LORA_PARAMETER_NAMES: - continue - num_experts = int(getattr(module, "num_experts", 1) or 1) - if num_experts <= 1 or not _wrapper_has_adapter(module, adapter_name): - continue - layout = layouts.get(parameter_name) - if layout is None and None in layouts: - experts = module.get_base_layer() - while hasattr(experts, "base_layer"): - experts = experts.base_layer - if _experts_family(experts) in _LEGACY_GROUPED_LORA_B_FAMILIES: - layout = layouts[None] - if layout != LORA_B_LAYOUT_GROUPED_BY_EXPERT: - continue - weight = module.lora_B[adapter_name].weight - rank = weight.shape[1] // num_experts - if rank <= 1 or rank * num_experts != weight.shape[1]: - continue - targets.append((weight, num_experts, rank)) - if any(weight.device.type == "meta" for weight, _, _ in targets): - # Accelerate holds disk-offloaded weights elsewhere; repacking only some would mix layouts. - raise RuntimeError( - "fused expert lora_B is offloaded, so it cannot be repacked in place; set " - "UNSLOTH_MOE_LORA_B_LAYOUT=grouped_by_expert to load this pre-2026.9.6 adapter" - ) - for weight, num_experts, rank in targets: - permutation = _moe_lora_b_column_permutation(num_experts, rank, weight.device, invert = True) - with torch.no_grad(): - weight.copy_(weight.index_select(1, permutation)) - # A resumed Trainer restores optimizer state in the old column order after this. - weight._unsloth_lora_b_repacked = (num_experts, rank) - return len(targets) - - -def repack_optimizer_state_for_legacy_lora_b(optimizer) -> Tuple[int, int]: - """Move the optimizer state of lora_B tensors `convert_legacy_lora_b` repacked into the - same column order; state that cannot be permuted (8-bit blockwise) is dropped and - re-initialised on the next step. Returns (moved, reset).""" - # This Unsloth Zoo code section is licensed under AGPL3 - - moved = reset = 0 - for group in getattr(optimizer, "param_groups", ()): - for param in group["params"]: - repacked = getattr(param, "_unsloth_lora_b_repacked", None) - if repacked is None: - continue - del param._unsloth_lora_b_repacked - state = optimizer.state.get(param) - if not state: - continue - permutation = _moe_lora_b_column_permutation(*repacked, param.device, invert = True) - tensors = [v for v in state.values() if torch.is_tensor(v) and v.shape == param.shape] - if tensors and all(v.is_floating_point() for v in tensors) and not any( - torch.is_tensor(v) and v.dim() > 0 and v.shape != param.shape for v in state.values() - ): - with torch.no_grad(): - for value in tensors: - value.copy_(value.index_select(1, permutation.to(value.device))) - moved += 1 - else: - optimizer.state.pop(param, None) - reset += 1 - return moved, reset - - -_original_trainer_load_optimizer_and_scheduler = None - - -def _patched_trainer_load_optimizer_and_scheduler(self, checkpoint, *args, **kwargs): - result = _original_trainer_load_optimizer_and_scheduler(self, checkpoint, *args, **kwargs) - try: - moved, reset = repack_optimizer_state_for_legacy_lora_b(self.optimizer) if checkpoint else (0, 0) - except Exception: - moved = reset = 0 - if moved or reset: - print( - f"Unsloth: resumed optimizer state for repacked fused MoE expert lora_B: {moved} " - f"tensors moved to the new column order, {reset} reset (8-bit state cannot be permuted)." - ) - return result - - -def _patch_trainer_resume_for_legacy_lora_b(): - # This Unsloth Zoo code section is licensed under AGPL3 - - global _original_trainer_load_optimizer_and_scheduler - try: - from transformers import Trainer - except Exception: - return False - method = getattr(Trainer, "_load_optimizer_and_scheduler", None) - if method is None or getattr(method, "_unsloth_moe_layout_patched", False): - return method is not None - if _original_trainer_load_optimizer_and_scheduler is None: - _original_trainer_load_optimizer_and_scheduler = method - patched = wraps(method)(_patched_trainer_load_optimizer_and_scheduler) - patched._unsloth_moe_layout_patched = True - Trainer._load_optimizer_and_scheduler = patched - return True - - -def _has_fused_expert_lora(peft_model, adapter_name) -> bool: - return any( - getattr(module, "parameter_name", None) in _SEPARATED_MOE_LORA_PARAMETER_NAMES - and int(getattr(module, "num_experts", 1) or 1) > 1 - and _wrapper_has_adapter(module, adapter_name) - for module in peft_model.modules() - ) - - -def _read_adapter_config(model_id, kwargs): - subfolder = kwargs.get("subfolder") or "" - path = os.path.join(str(model_id), subfolder, "adapter_config.json") - if not os.path.isfile(path): - from huggingface_hub import hf_hub_download - # The Hub options PEFT was given (cache_dir, local_files_only, ...), so this finds - # the same file PEFT just loaded. - accepted = inspect.signature(hf_hub_download).parameters - download = {k: v for k, v in kwargs.items() if k in accepted} - if download.get("token") is None and kwargs.get("use_auth_token") is not None: - download["token"] = kwargs["use_auth_token"] - download.update(repo_id = str(model_id), filename = "adapter_config.json", subfolder = subfolder or None) - try: - path = hf_hub_download(**{**download, "local_files_only": True}) - except Exception: - path = hf_hub_download(**download) - with open(path, "r", encoding = "utf-8") as f: - return json.load(f) - - -def _patched_peft_model_load_adapter(self, model_id, adapter_name, *args, **kwargs): - """PEFT's load_adapter, then repack a pre-2026.9.6 Unsloth fused expert lora_B into - rank_major so the forward, merge and the next save all agree (unsloth#6930). An explicit - `UNSLOTH_MOE_LORA_B_LAYOUT` keeps the caller's setting and skips this.""" - # This Unsloth Zoo code section is licensed under AGPL3 - - result = _original_peft_model_load_adapter(self, model_id, adapter_name, *args, **kwargs) - if os.environ.get("UNSLOTH_MOE_LORA_B_LAYOUT") or not _has_fused_expert_lora(self, adapter_name): - return result - try: - converted = convert_legacy_lora_b(self, adapter_name, _read_adapter_config(model_id, kwargs)) - except Exception as exception: - converted = 0 - logger = _moe_utils_logger() - if logger is not None: - logger.warning( - f"Unsloth: could not check {model_id}'s fused MoE expert LoRA layout " - f"({type(exception).__name__}: {exception})." - ) - if converted and (str(model_id), adapter_name) not in _legacy_lora_b_converted_logged: - _legacy_lora_b_converted_logged.add((str(model_id), adapter_name)) - print( - f"Unsloth: {model_id} was trained before unsloth_zoo 2026.9.6 with the old fused " - f"MoE expert LoRA packing; repacked {converted} lora_B tensors to PEFT's layout. " - f"Save it again to update the file (UNSLOTH_MOE_LORA_B_LAYOUT=rank_major loads " - f"it unchanged)." - ) - return result - - -def _patch_peft_load_adapter_for_moe_layout(): - # This Unsloth Zoo code section is licensed under AGPL3 - - global _original_peft_model_load_adapter - - try: - from peft import PeftModel - except Exception: - return False - - if getattr(PeftModel.load_adapter, "_unsloth_moe_layout_patched", False): - return True - - if _original_peft_model_load_adapter is None: - _original_peft_model_load_adapter = PeftModel.load_adapter - - patched = wraps(_original_peft_model_load_adapter)(_patched_peft_model_load_adapter) - patched._unsloth_moe_layout_patched = True - PeftModel.load_adapter = patched - return True - - -def patch_param_wrapper_for_moe(): - """Patch PEFT's ParamWrapper.forward for MoE separated LoRA (call after PEFT import).""" - # This Unsloth Zoo code section is licensed under AGPL3 - - global _original_param_wrapper_forward - global _original_param_wrapper_get_delta_weight - - module = _load_cached_moe_utils_module() - if module is not None and hasattr(module, "patch_param_wrapper_for_moe"): - try: - return module.patch_param_wrapper_for_moe() - except Exception: - pass - - try: - from peft.tuners.lora.layer import ParamWrapper - - if _original_param_wrapper_forward is None: - _original_param_wrapper_forward = ParamWrapper.forward - - ParamWrapper.forward = _patched_param_wrapper_forward - - # The forward is only half of it: the merge path reads lora_B through - # get_delta_weight, which is PEFT's and packs it the other way. - if not getattr(ParamWrapper.get_delta_weight, "_unsloth_moe_layout_patched", False): - if _original_param_wrapper_get_delta_weight is None: - _original_param_wrapper_get_delta_weight = ParamWrapper.get_delta_weight - patched = wraps(_original_param_wrapper_get_delta_weight)( - _patched_param_wrapper_get_delta_weight, - ) - patched._unsloth_moe_layout_patched = True - ParamWrapper.get_delta_weight = patched - - _patch_peft_save_pretrained_for_moe_layout() - _patch_peft_load_adapter_for_moe_layout() - _patch_trainer_resume_for_legacy_lora_b() - _patch_peft_get_peft_model_for_moe() - - return True - except ImportError: - return False - - -# Gate gradient via the inner-product identity dGate = / gate, instead of -# which pins the down-projection output Y on the tape solely for that -# gradient. Output unchanged; on by default to save memory. The runtime gate below -# still auto-disables it for fp16, down-bias models and frozen routers. Set -# UNSLOTH_MOE_GATEGRAD=0 to revert to the standard path. - - -_WEIGHTED_UNPERMUTE = None - - -def _resolve_weighted_unpermute(): - global _WEIGHTED_UNPERMUTE - if _WEIGHTED_UNPERMUTE is None: - try: - from unsloth_zoo.temporary_patches.moe_triton_kernels import weighted_unpermute as _wu - except Exception: # now runs at import too: a broken triton must not fail importing zoo - _wu = False - _WEIGHTED_UNPERMUTE = _wu - return _WEIGHTED_UNPERMUTE - - -def weighted_unpermute(*args, **kwargs): - """Fused Triton combine, or None so the caller falls back to the eager reduction.""" - if _resolve_weighted_unpermute() is False: - return None - return _WEIGHTED_UNPERMUTE(*args, **kwargs) - - -# Resolved at import: a first-call assignment inside a compiled MoE block guards on the global being -# None, so the next call recompiles every compiled decoder layer once. -_resolve_weighted_unpermute() - - -@lru_cache(maxsize=1) -def _moe_gategrad_enabled() -> bool: - """Whether the MoE gate-gradient identity path is active (on by default). - - Cached with maxsize=1 since UNSLOTH_MOE_GATEGRAD is read once per process. Any - code or test that toggles the env var at runtime must call - _moe_gategrad_enabled.cache_clear() afterwards for the change to take effect. - """ - return os.environ.get("UNSLOTH_MOE_GATEGRAD", "1") != "0" - - -class _MoEGateGradIdentity(torch.autograd.Function): - """Identity over the pre-down activation ``inter``; backward derives the gate - gradient as ``dGate = / gate`` instead of ````. - - The incoming dA equals ``gate * (dOut @ W2_eff.T)`` for the effective down - weight (base + LoRA), so `` / gate`` is exactly ```` - without ever materialising Y for the gradient. The caller passes a - sign-floored gate and multiplies Y by that same value, so the identity holds - for any routing weight, including exact zeros. Exact for any linear down - projection; NOT valid with a post-matmul down bias, so callers must - disable it when one is present. - """ - - @staticmethod - def forward(ctx, inter, permuted_weights): - ctx.save_for_backward(inter, permuted_weights) - return inter - - @staticmethod - @torch.autograd.function.once_differentiable - def backward(ctx, grad_inter): - # once_differentiable: the reconstruction is first-order exact but its - # graph is not the original one, so create_graph must raise, not - # silently return wrong second derivatives. - # grad_inter is None when nothing downstream needs inter's gradient - # (e.g. a frozen down projection); there is nothing to propagate then. - if grad_inter is None: - return None, None - inter, gate = ctx.saved_tensors - grad_gate = None - # Skip the gate gradient when the routing weight is frozen (e.g. a frozen - # router under LoRA), where its gradient is never consumed. - if ctx.needs_input_grad[1]: - dgate = (inter.to(torch.float32) * grad_inter.to(torch.float32)).sum(dim=-1) - # The call site multiplies Y by this same sign-floored gate, so the - # floor cancels exactly and the quotient equals for any - # gate, including exact zeros. The re-clamp is a no-op for the safe - # gate passed by forward_native_grouped_mm; it only protects direct - # callers from dividing by a raw zero. - gate_f = gate.to(torch.float32) - safe_gate = torch.where( - gate_f >= 0, gate_f.clamp_min(1e-12), gate_f.clamp_max(-1e-12) - ) - grad_gate = (dgate / safe_gate).to(gate.dtype) - # inter passes through unchanged, so its gradient is grad_inter. - return grad_inter, grad_gate - - -_FLAG_MISSING = object() - - -def _module_flag(module, name, default = None): - value = module.__dict__.get(name, _FLAG_MISSING) - if value is _FLAG_MISSING: - value = getattr(type(module), name, default) - return value - - -def _gate_up_is_interleaved(module) -> bool: - if "GptOssExperts" in module.__class__.__name__: - return True - return _module_flag(module, "is_concatenated", True) is False - - -def _uses_own_apply_gate(module) -> bool: - """Own-gate flag on the class or (FP8) instance; hot path, so no nn.Module.__getattr__.""" - return bool( - module.__dict__.get("_unsloth_own_apply_gate", False) - or getattr(type(module), "_unsloth_own_apply_gate", False) - ) - - -def forward_native_grouped_mm( - self, - hidden_states: torch.Tensor, - top_k_index: torch.Tensor, - top_k_weights: torch.Tensor, -) -> torch.Tensor: - """Native PyTorch grouped-GEMM MoE forward via torch._grouped_mm (no Triton; needs runtime support).""" - # This Unsloth Zoo code section is licensed under AGPL3 - - # Runtime safety check (defense in depth). - if not _check_torch_grouped_mm_supported(): - # Compute Capability is CUDA-only; on XPU it would mask this message. - if torch.cuda.is_available(): - major, minor = torch.cuda.get_device_capability(torch.cuda.current_device()) - where = f"this device (Compute Capability {major}.{minor})" - else: - where = "this device" - raise RuntimeError( - f"torch._grouped_mm is not supported on {where}. " - f"Set UNSLOTH_MOE_BACKEND='unsloth_triton' or 'native_torch' to use a compatible backend." - ) - - is_2d_input = hidden_states.dim() == 2 - if is_2d_input: - sequence_length, hidden_dim = hidden_states.shape - batch_size = 1 - else: - batch_size, sequence_length, hidden_dim = hidden_states.shape - - hidden_states = hidden_states.view(-1, hidden_dim) - # torch._grouped_mm is not autocast-cast, and float32 operands hit its slow per-group fallback. - if hidden_states.dtype == torch.float32: - _stack = self._parameters.get("gate_up_proj", self._parameters.get("gate_proj")) - if _stack is not None and _stack.dtype in (torch.float16, torch.bfloat16): - hidden_states = hidden_states.to(_stack.dtype) - elif _stack is not None and _stack.dtype != torch.float32: - hidden_states = hidden_states.to(moe_compute_dtype(hidden_states)) - - # Routing: count tokens per expert, sort to group by expert, gather inputs. - flat_top_k = top_k_index.view(-1) - num_tokens_per_expert, sorted_indices, offsets = moe_routing_sort(top_k_index, self.num_experts) - permuted_input = gather_routed_rows(hidden_states, sorted_indices, top_k_index.shape[-1]) - - # Gate + Up projection with optional separated LoRA (default). - use_separated_lora = _should_use_separated_lora() - gate_up_lora = None - - # Prefer LoRA injected by the patched ParamWrapper; fall back to the parameter. - _stashed_gate_up_lora = take_moe_lora_stash(self, "gate_up_proj") - if _stashed_gate_up_lora is not None: - gate_up_lora = _stashed_gate_up_lora[:3] # (first, second, scaling) - elif ( - use_separated_lora - and hasattr(self, "gate_up_proj") - and _has_lora_adapters(self.gate_up_proj) - ): - gate_up_lora = _extract_lora_weights( - self.gate_up_proj, num_experts=self.num_experts, experts_module=self - ) - - own_gate_up = None - fused_inter = None - selective_plan = None - if hasattr(self, "gate_up_proj"): - model_type = getattr(self, "_unsloth_model_type", None) - _gate_up_src = self.gate_up_proj - - # Provider re-derives the base weight on demand so Fix 3 can recompute it in - # backward instead of pinning it (grouped_mm needs contiguous weights). - def _gate_up_provider(experts=None, _src=_gate_up_src, _mt=model_type, _h=hidden_dim, _dt=hidden_states.dtype, _mod=self, _n=num_tokens_per_expert): - return preprocess_weight(_get_base_weight(_src, _dt, _n, experts), "gate_up", _h, _mt, experts_module=_mod) - _gate_up_layout = _mxfp4_expert_layout(_gate_up_src, "gate_up", hidden_dim, model_type, self) - if _gate_up_layout is not None: - mm1_out = _mxfp4_base_grouped_mm( - permuted_input, offsets, num_tokens_per_expert, _gate_up_provider, - _moe_recompute_enabled(_gate_up_src, dtype=hidden_states.dtype), _gate_up_layout, - ) - else: - mm1_out = _nf4_fused_base_mm(permuted_input, offsets, _gate_up_src, "gate_up", hidden_dim, model_type, self) - if _gate_up_layout is None and mm1_out is None: - _recompute = _moe_recompute_enabled(_gate_up_src, dtype=hidden_states.dtype) - if _recompute: - selective_plan = _selective_dequant_plan( - (_gate_up_src, getattr(self, "down_proj", None)), num_tokens_per_expert, permuted_input.shape[0], - ) - mm1_out = _base_grouped_mm( - permuted_input, offsets, _gate_up_provider, _recompute, selective_plan, - ) - - # Separated LoRA: + ((X @ first) @ second) * scaling. - if gate_up_lora is not None: - first_weight, second_weight, scaling = gate_up_lora - - # Cast to input dtype (LoRA is float32) and make contiguous for grouped_mm. - first_weight = first_weight.to(permuted_input.dtype).contiguous() - second_weight = second_weight.to(permuted_input.dtype).contiguous() - first_weight, second_weight = _pad_lora_rank_for_grouped_mm(first_weight, second_weight, permuted_input) - - try: - lora_out = _grouped_mm_with_backward_fix(permuted_input, first_weight, offsets) - lora_out = lora_out.contiguous() - except RuntimeError as e: - raise e - - # Second matmul; pad an unaligned output dim or fall back on failure. - try: - if second_weight.shape[-1] % 8 != 0: - pad_size = 8 - (second_weight.shape[-1] % 8) - second_weight_padded = F.pad( - second_weight, (0, pad_size) - ).contiguous() - lora_delta = _grouped_mm_with_backward_fix( - lora_out, second_weight_padded, offsets - ) - lora_delta = lora_delta[:, :-pad_size] - else: - lora_delta = _grouped_mm_with_backward_fix( - lora_out, second_weight, offsets - ) - except RuntimeError: - # Manual loop fallback on grouped_mm failure (e.g. stride alignment). - lora_delta = torch.empty( - (lora_out.shape[0], second_weight.shape[-1]), - dtype=lora_out.dtype, - device=lora_out.device, - ) - cpu_offsets = offsets.cpu().tolist() - prev_offset = 0 - for i, end in enumerate(cpu_offsets): - if prev_offset < end: - lora_delta[prev_offset:end] = torch.matmul( - lora_out[prev_offset:end], second_weight[i] - ) - prev_offset = end - - mm1_out = mm1_out + lora_delta * scaling - - if hasattr(self, "gate_up_proj_bias") and self.gate_up_proj_bias is not None: - # repeat_interleave without output_size= D2H-syncs. sorted_indices is - # already in expert order, so gathering by expert id matches it, sync-free - # (74 us -> 8 us host at E=128 x 16384 rows). - sorted_expert_ids = flat_top_k[sorted_indices] - bias_expanded = _gather_expert_bias(self.gate_up_proj_bias, sorted_expert_ids) - mm1_out = mm1_out + bias_expanded.to(mm1_out.dtype) - - _interleaved = _gate_up_is_interleaved(self) - if _interleaved: - gate = mm1_out[..., ::2] - up = mm1_out[..., 1::2] - else: - gate, up = mm1_out.chunk(2, dim=-1) - own_gate_up = mm1_out if _uses_own_apply_gate(self) else None - if own_gate_up is None: - fused_inter = _fused_moe_activation(self, mm1_out, _interleaved) - - elif hasattr(self, "w1") and hasattr(self, "w3"): - # Separate w1/w3 weights (older models). - w1_base = _get_base_weight(self.w1) - w3_base = _get_base_weight(self.w3) - - w1 = w1_base.transpose(-2, -1) - w3 = w3_base.transpose(-2, -1) - - gate = _grouped_mm_with_backward_fix(permuted_input, w1, offsets) - up = _grouped_mm_with_backward_fix(permuted_input, w3, offsets) - - # Add LoRA for w1 and w3 separately if present. - if use_separated_lora: - if _has_lora_adapters(self.w1): - w1_lora = _extract_lora_weights(self.w1, experts_module=self) - if w1_lora is not None: - lora_A, lora_B, scaling = w1_lora - lora_A_t, lora_B_t = _pad_lora_rank_for_grouped_mm( - lora_A.transpose(-2, -1), lora_B.transpose(-2, -1), permuted_input - ) - lora_A_out = _grouped_mm_with_backward_fix( - permuted_input, lora_A_t, offsets - ) - lora_B_out = _grouped_mm_with_backward_fix(lora_A_out, lora_B_t, offsets) - gate = gate + lora_B_out * scaling - - if _has_lora_adapters(self.w3): - w3_lora = _extract_lora_weights(self.w3, experts_module=self) - if w3_lora is not None: - lora_A, lora_B, scaling = w3_lora - lora_A_t, lora_B_t = _pad_lora_rank_for_grouped_mm( - lora_A.transpose(-2, -1), lora_B.transpose(-2, -1), permuted_input - ) - lora_A_out = _grouped_mm_with_backward_fix( - permuted_input, lora_A_t, offsets - ) - lora_B_out = _grouped_mm_with_backward_fix(lora_A_out, lora_B_t, offsets) - up = up + lora_B_out * scaling - else: - raise AttributeError("MoE layer must have 'gate_up_proj' or 'w1'/'w3'.") - - # Activation - if own_gate_up is not None: - inter = self._apply_gate(own_gate_up) - elif fused_inter is not None: - inter = fused_inter - elif "GptOssExperts" in self.__class__.__name__: - # Custom GptOss activation. - limit = getattr(self, "limit", 7.0) - alpha = getattr(self, "alpha", 1.702) - - gate = gate.clamp(min=None, max=limit) - up = up.clamp(min=-limit, max=limit) - glu = gate * torch.sigmoid(gate * alpha) - inter = (up + 1.0) * glu - elif hasattr(self, 'act_fn') and callable(self.act_fn): - inter = self.act_fn(gate) * up - else: - inter = F.silu(gate) * up - - # Env-gated gate-grad identity; disabled when a down bias exists (identity - # assumes a linear down), the router is frozen (nothing to synthesize, and - # the Function would needlessly keep inter saved), or inter is float16. For - # fp16, the gradient into the identity carries the tiny floored gate as a - # factor (grad_inter = gate * (W2 @ dOut)); with a near-zero gate that product - # underflows fp16 to 0 before backward runs, dropping the router gradient for - # near-zero routes. fp16 therefore keeps the standard multiply (correct - # gradient, saves Y); bf16/fp32 hold the tiny value and are unaffected. - _gategrad = ( - _moe_gategrad_enabled() - and getattr(self, "down_proj_bias", None) is None - and top_k_weights.requires_grad - and inter.dtype != torch.float16 - ) - permuted_weights = None - if _gategrad: - raw_weights = top_k_weights.reshape(-1)[sorted_indices] - # Sign-floor the gate away from zero (straight-through, so the synthesized - # gradient still reaches top_k_weights). The multiply below and the - # identity's divide use the SAME floored value, so they cancel exactly and - # the gate gradient survives even a routing weight of exactly - # zero. An fp16 gate cannot represent eps=1e-12 (its smallest normal is - # ~6e-5), so upcast it to float32; inter is non-fp16 here, so the floored - # value survives in grad_inter. The forward then changes by at most - # 1e-12 * |Y| for |gate| < 1e-12. - if raw_weights.dtype == torch.float16: - raw_weights = raw_weights.to(torch.float32) - eps = 1e-12 - floored = torch.where( - raw_weights >= 0, - raw_weights.clamp(min=eps), - raw_weights.clamp(max=-eps), - ) - permuted_weights = raw_weights + (floored - raw_weights).detach() - inter = _MoEGateGradIdentity.apply(inter, permuted_weights) - - # Down projection with optional separated LoRA (default). - down_lora = None - - # Prefer LoRA injected by the patched ParamWrapper; fall back to the parameter. - _stashed_down_lora = take_moe_lora_stash(self, "down_proj") - if _stashed_down_lora is not None: - down_lora = _stashed_down_lora[:3] # (first, second, scaling) - elif ( - use_separated_lora - and hasattr(self, "down_proj") - and _has_lora_adapters(self.down_proj) - ): - down_lora = _extract_lora_weights(self.down_proj, num_experts=self.num_experts, experts_module=self) - - if hasattr(self, "down_proj"): - model_type = getattr(self, "_unsloth_model_type", None) - _down_src = self.down_proj - def _down_provider(experts=None, _src=_down_src, _mt=model_type, _h=hidden_dim, _dt=hidden_states.dtype, _mod=self, _n=num_tokens_per_expert): - return preprocess_weight(_get_base_weight(_src, _dt, _n, experts), "down", _h, _mt, experts_module=_mod) - _down_layout = _mxfp4_expert_layout(_down_src, "down", hidden_dim, model_type, self) - if _down_layout is not None: - mm2_out = _mxfp4_base_grouped_mm( - inter, offsets, num_tokens_per_expert, _down_provider, - _moe_recompute_enabled(_down_src, dtype=hidden_states.dtype), _down_layout, - ) - else: - mm2_out = _nf4_fused_base_mm(inter, offsets, _down_src, "down", hidden_dim, model_type, self) - if _down_layout is None and mm2_out is None: - _recompute = _moe_recompute_enabled(_down_src, dtype=hidden_states.dtype) - mm2_out = _base_grouped_mm( - inter, offsets, _down_provider, _recompute, selective_plan if _recompute else None, - ) - - if down_lora is not None: - first_weight, second_weight, scaling = down_lora - - # Cast to input dtype (LoRA is float32) and make contiguous for grouped_mm. - first_weight = first_weight.to(inter.dtype).contiguous() - second_weight = second_weight.to(inter.dtype).contiguous() - first_weight, second_weight = _pad_lora_rank_for_grouped_mm(first_weight, second_weight, inter) - - lora_out = _grouped_mm_with_backward_fix(inter, first_weight, offsets) - lora_out = lora_out.contiguous() - - try: - lora_delta = _grouped_mm_with_backward_fix(lora_out, second_weight, offsets) - except RuntimeError: - # Manual loop fallback. - lora_delta = torch.empty( - (lora_out.shape[0], second_weight.shape[-1]), - dtype=lora_out.dtype, - device=lora_out.device, - ) - cpu_offsets = offsets.cpu().tolist() - prev_offset = 0 - for i, end in enumerate(cpu_offsets): - if prev_offset < end: - lora_delta[prev_offset:end] = torch.matmul( - lora_out[prev_offset:end], second_weight[i] - ) - prev_offset = end - - mm2_out = mm2_out + lora_delta * scaling - - if hasattr(self, "down_proj_bias") and self.down_proj_bias is not None: - # Capture-safe gather; see gate_up_proj_bias above. - sorted_expert_ids = flat_top_k[sorted_indices] - bias_expanded = _gather_expert_bias(self.down_proj_bias, sorted_expert_ids).to(mm2_out.device) - mm2_out = mm2_out + bias_expanded.to(mm2_out.dtype) - - elif hasattr(self, "w2"): - w2_base = _get_base_weight(self.w2) - w2 = w2_base.transpose(-2, -1) - mm2_out = _grouped_mm_with_backward_fix(inter, w2, offsets) - - if use_separated_lora and _has_lora_adapters(self.w2): - w2_lora = _extract_lora_weights(self.w2, experts_module=self) - if w2_lora is not None: - lora_A, lora_B, scaling = w2_lora - lora_A_t, lora_B_t = _pad_lora_rank_for_grouped_mm( - lora_A.transpose(-2, -1), lora_B.transpose(-2, -1), inter - ) - lora_A_out = _grouped_mm_with_backward_fix(inter, lora_A_t.contiguous(), offsets) - lora_B_t = lora_B_t.contiguous() - lora_B_out = _grouped_mm_with_backward_fix(lora_A_out, lora_B_t, offsets) - mm2_out = mm2_out + lora_B_out * scaling - else: - raise AttributeError("MoE layer must have 'down_proj' or 'w2'.") - - # Apply routing weights and scatter-add (reduce). - if _gategrad: - # Gate grad comes from the identity; detach so the multiply does not pin Y. - permuted_weights = permuted_weights.detach() - else: - flat_weights = top_k_weights.reshape(-1) - permuted_weights = flat_weights[sorted_indices] - - final_hidden_states = weighted_unpermute( - mm2_out, - sorted_indices, - permuted_weights, - batch_size * sequence_length, - top_k_index.shape[-1], - out_dtype = hidden_states.dtype, - ) - if final_hidden_states is None: - mm2_out = mm2_out * permuted_weights.unsqueeze(-1) - final_hidden_states = combine_permuted_moe_outputs( - mm2_out, - sorted_indices, - batch_size * sequence_length, - top_k_index.shape[-1], - out_dtype = hidden_states.dtype, - ) - - if is_2d_input: - return final_hidden_states - - return final_hidden_states.view(batch_size, sequence_length, hidden_dim) - - -def _experts_are_input_major(module, name, in_dim) -> bool: - if bool(_module_flag(module, "_unsloth_grouped_mm_format", False)): - return True - declared = _module_flag(module, "is_transposed", None) - if isinstance(declared, bool): - return declared - return getattr(module, name).shape[-1] != in_dim - - -TRITON_MOE_COMPUTE_DTYPE_STATS = {"cast": 0, "kept_fp32": 0} - - -def _triton_moe_compute_dtype(module, hidden_states): - """bf16 for float32 activations into bf16 experts (matches autocast); float16 is left alone (overflow).""" - if hidden_states.dtype != torch.float32: - return hidden_states.dtype - for name in ("gate_up_proj", "down_proj"): - weight = getattr(module, name, None) - if not isinstance(weight, torch.Tensor) or weight.dtype != torch.bfloat16: - return hidden_states.dtype - if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1": - return hidden_states.dtype - keep = os.environ.get("UNSLOTH_MOE_TRITON_FP32", "0") == "1" - # Eager only: Dynamo would recompile on every counter value. - if not torch.compiler.is_compiling(): - TRITON_MOE_COMPUTE_DTYPE_STATS["kept_fp32" if keep else "cast"] += 1 - return hidden_states.dtype if keep else torch.bfloat16 - - -@torch.compiler.disable -def _moe_weight_dtype_kwargs(tuner, module, name, dtype, device = None): - """weight_dtype kwarg when expert weights differ from activations (bf16 W, float32 X), and the - operands' CUDA device: on a mixed-GPU host the configs must fit the card the GEMM launches on.""" - try: - parameters = inspect.signature(tuner).parameters - except (TypeError, ValueError): - return {} - kwargs = {} - weight = getattr(module, name, None) - weight_dtype = getattr(weight, "dtype", None) - if weight_dtype in (torch.float32, torch.bfloat16, torch.float16) and weight_dtype != dtype \ - and "weight_dtype" in parameters: - kwargs["weight_dtype"] = weight_dtype - # Older unsloth has no device argument and reads the current device. - if device is not None and device.type == "cuda" and "device" in parameters: - kwargs["device"] = device - return kwargs - - -# tl.dot needs tile dims >= 16 and the kernels static_assert K % BLOCK_SIZE_K == 0 (N likewise); the tile -# fitter stops at 16, so a hidden or MoE intermediate size that is not a multiple of 16 cannot compile. -_TRITON_GROUPED_GEMM_DIM_MULTIPLE = 16 - - -def _triton_grouped_gemm_dims_supported(module, hidden_dim) -> bool: - try: - gate_up = module.gate_up_proj - if _experts_are_input_major(module, "gate_up_proj", hidden_dim): - intermediate_dim = gate_up.shape[-1] // 2 - else: - intermediate_dim = gate_up.shape[1] // 2 - except Exception: - return True - m = _TRITON_GROUPED_GEMM_DIM_MULTIPLE - return hidden_dim % m == 0 and intermediate_dim % m == 0 - - -def forward_triton_grouped_gemm( - self, - hidden_states: torch.Tensor, - top_k_index: torch.Tensor, - top_k_weights: torch.Tensor, -) -> torch.Tensor: - """Grouped-GEMM MoE forward via Triton kernels (torch.compile-compatible, mode="max-autotune").""" - # This Unsloth Zoo code section is licensed under AGPL3 - - if not _triton_grouped_gemm_dims_supported(self, hidden_states.shape[-1]): - # Toy shapes (hidden_size 8): the same routing on grouped_mm, else the eager loop; neither has the tile limit. - if _grouped_mm_supported_on(hidden_states.device): - return forward_native_grouped_mm(self, hidden_states, top_k_index, top_k_weights) - return forward_native_moe_loop(self, hidden_states, top_k_index, top_k_weights) - - from unsloth.kernels.moe.grouped_gemm.interface import grouped_gemm - from unsloth.kernels.moe.autotune_cache import get_or_autotune_moe_kernels - - if not hasattr(self, "_unsloth_moe_configs"): - self._unsloth_moe_configs = None - - use_separated_lora = _should_use_separated_lora() - - # gate_up LoRA from the patched ParamWrapper (mirrors the down block below). - gate_up_lora = None - _stashed_gate_up_lora = take_moe_lora_stash(self, "gate_up_proj") - if _stashed_gate_up_lora is not None: - gate_up_lora = _stashed_gate_up_lora[:3] - elif ( - use_separated_lora - and hasattr(self, "gate_up_proj") - and _has_lora_adapters(self.gate_up_proj) - ): - gate_up_lora = _extract_lora_weights( - self.gate_up_proj, num_experts=self.num_experts - ) - - # Flatten 3D inputs (batch_size, seq_len, hidden_dim). - is_3d = hidden_states.dim() == 3 - if is_3d: - batch_size, seq_len, hidden_dim = hidden_states.shape - hidden_states = hidden_states.view(-1, hidden_dim) - num_tokens = batch_size * seq_len - if top_k_index.dim() == 3: - top_k_index = top_k_index.view(-1, top_k_index.shape[-1]) - if top_k_weights.dim() == 3: - top_k_weights = top_k_weights.view(-1, top_k_weights.shape[-1]) - else: - num_tokens, hidden_dim = hidden_states.shape - - top_k = top_k_index.shape[1] - compute_dtype = _triton_moe_compute_dtype(self, hidden_states) - # The gate_up GEMM reads X in compute_dtype too, as plain bf16 training does: a float32 X runs it in - # float32, whose default tile does not fit L4 shared memory. float16 and the FP32 flags keep float32. - # LoRA and the routing sum keep the caller's dtype. - gemm_hidden_states = hidden_states if hidden_states.dtype == compute_dtype else hidden_states.to(compute_dtype) - - # Cache model dims and kernel configs on first call. - if self._unsloth_moe_configs is None: - if _experts_are_input_major(self, "gate_up_proj", hidden_dim): - intermediate_dim = self.gate_up_proj.shape[-1] // 2 - else: - intermediate_dim = self.gate_up_proj.shape[1] // 2 - - # The autotuner doubles intermediate_dim itself: a doubled N can tune a BLOCK_SIZE_N the real N does not divide. - gemm1_configs = get_or_autotune_moe_kernels( - num_experts=self.num_experts, - hidden_dim=hidden_dim, - intermediate_dim=intermediate_dim, - top_k=top_k, - dtype=gemm_hidden_states.dtype, - **_moe_weight_dtype_kwargs( - get_or_autotune_moe_kernels, self, "gate_up_proj", gemm_hidden_states.dtype, gemm_hidden_states.device, - ), - ) - - # Second GEMM: K = intermediate_dim, N = hidden_dim. - gemm2_configs = get_or_autotune_moe_kernels( - num_experts=self.num_experts, - hidden_dim=intermediate_dim, - intermediate_dim=hidden_dim // 2 if hidden_dim % 2 == 0 else hidden_dim, - top_k=top_k, - dtype=compute_dtype, - **_moe_weight_dtype_kwargs( - get_or_autotune_moe_kernels, self, "down_proj", compute_dtype, gemm_hidden_states.device, - ), - ) - - self._unsloth_moe_configs = (intermediate_dim, gemm1_configs, gemm2_configs) - torch.cuda.empty_cache() - - intermediate_dim, gemm1_configs, gemm2_configs = self._unsloth_moe_configs - fwd_config_1, bwd_dX_config_1, bwd_dW_config_1 = gemm1_configs - fwd_config_2, bwd_dX_config_2, bwd_dW_config_2 = gemm2_configs - - with torch.no_grad(): - token_counts_by_expert, gather_indices, offsets = moe_routing_sort(top_k_index, self.num_experts) - - if _experts_are_input_major(self, "gate_up_proj", hidden_dim): - w1 = self.gate_up_proj.transpose(-2, -1).contiguous() - else: - w1 = self.gate_up_proj - - # First grouped GEMM: gate_up projection. - first_gemm_output = grouped_gemm( - X=gemm_hidden_states, - W=w1, - m_sizes=token_counts_by_expert, - topk=top_k, - gather_indices=gather_indices, - permute_x=True, - permute_y=False, - autotune=False, # cached configs - kernel_config_fwd=fwd_config_1, - kernel_config_bwd_dX=bwd_dX_config_1, - kernel_config_bwd_dW=bwd_dW_config_1, - is_first_gemm=True, - ) - - # Separated LoRA for gate_up. grouped_gemm ran permute_x=True so first_gemm_output - # is expert-sorted; _apply_lora_grouped_mm wants pre-permuted input, so gather via - # gather_indices // top_k (expert-sorted row -> originating token row). - if gate_up_lora is not None: - first_weight, second_weight, scaling = gate_up_lora - first_weight = first_weight.to(hidden_states.dtype) - second_weight = second_weight.to(hidden_states.dtype) - permuted_hidden = gather_routed_rows(hidden_states, gather_indices, top_k) - gate_up_lora_delta = _apply_lora_grouped_mm( - permuted_hidden, - first_weight, - second_weight, - offsets, - scaling, - grouped_mm_func=native_moe_grouped_mm, - ) - first_gemm_output = first_gemm_output + gate_up_lora_delta - - if first_gemm_output.dtype != compute_dtype: - first_gemm_output = first_gemm_output.to(compute_dtype) - - # Activation + gate*up. - fused_intermediate = None - if not _uses_own_apply_gate(self): - fused_intermediate = _fused_moe_activation(self, first_gemm_output, False, allow_gptoss = False) - if _uses_own_apply_gate(self): - intermediate = self._apply_gate(first_gemm_output) - elif fused_intermediate is not None: - intermediate = fused_intermediate - elif hasattr(self, 'act_fn') and callable(self.act_fn): - gate, up = first_gemm_output.chunk(2, dim=-1) - intermediate = self.act_fn(gate) * up - else: - intermediate = _silu_and_mul(first_gemm_output) - - # Grouped GEMM 2: down projection. - down_lora = None - _stashed_down_lora = take_moe_lora_stash(self, "down_proj") - if _stashed_down_lora is not None: - down_lora = _stashed_down_lora[:3] - elif ( - use_separated_lora - and hasattr(self, "down_proj") - and _has_lora_adapters(self.down_proj) - ): - down_lora = _extract_lora_weights(self.down_proj, num_experts=self.num_experts) - - if _experts_are_input_major(self, "down_proj", intermediate.shape[-1]): - w2 = self.down_proj.transpose(-2, -1).contiguous() - else: - w2 = self.down_proj - - second_gemm_output = grouped_gemm( - X=intermediate, - W=w2, - m_sizes=token_counts_by_expert, - topk=top_k, - gather_indices=gather_indices, - permute_x=False, - permute_y=True, - autotune=False, # cached configs - kernel_config_fwd=fwd_config_2, - kernel_config_bwd_dX=bwd_dX_config_2, - kernel_config_bwd_dW=bwd_dW_config_2, - is_first_gemm=False, - ) - - # Separated LoRA for down (intermediate already permuted from step 1, same offsets). - if down_lora is not None: - first_weight, second_weight, scaling = down_lora - - first_weight = first_weight.to(intermediate.dtype) - second_weight = second_weight.to(intermediate.dtype) - - lora_delta = _apply_lora_grouped_mm( - intermediate, - first_weight, - second_weight, - offsets, - scaling, - grouped_mm_func=native_moe_grouped_mm - ) - - # permute_y=True put second_gemm_output in token order; lora_delta is still - # expert-sorted, so scatter it rather than adding it row-for-row. - if lora_delta.dtype == second_gemm_output.dtype: - second_gemm_output.index_add_(0, gather_indices, lora_delta) - else: - # index_add_ requires matching dtypes; the plain add this replaced promoted. - promoted = torch.promote_types(second_gemm_output.dtype, lora_delta.dtype) - second_gemm_output = second_gemm_output.to(promoted).index_add( - 0, gather_indices, lora_delta.to(promoted) - ) - - # Apply routing weights and sum across top_k: (num_tokens, top_k, hidden) -> (num_tokens, hidden). - top_k_weights_casted = top_k_weights.to(hidden_states.dtype) - final_hidden_states = ( - second_gemm_output.view(num_tokens, top_k, hidden_dim) - * top_k_weights_casted[..., None] - ) - final_hidden_states = final_hidden_states.sum(dim=1) - - if is_3d: - final_hidden_states = final_hidden_states.view(batch_size, seq_len, hidden_dim) - - return final_hidden_states - - -@torch.compiler.disable -def forward_native_moe_loop( - self, - hidden_states: torch.Tensor, - top_k_index: torch.Tensor, - top_k_weights: torch.Tensor, -) -> torch.Tensor: - """Loop over experts with routed tokens; torch.compile-disabled to avoid graph breaks on dynamic control flow.""" - # This Unsloth Zoo code section is licensed under AGPL3 - final_hidden_states = torch.zeros_like(hidden_states) - use_separated_lora = _should_use_separated_lora() - - gate_up_lora = take_moe_lora_stash(self, "gate_up_proj") - if gate_up_lora is not None: - gate_up_lora = gate_up_lora[:3] - elif ( - use_separated_lora - and hasattr(self, "gate_up_proj") - and _has_lora_adapters(self.gate_up_proj) - ): - gate_up_lora = _extract_lora_weights( - self.gate_up_proj, num_experts=self.num_experts, experts_module=self - ) - # Pre-cast LoRA factors to the activation dtype once (avoid per-expert .to()). - # `scaling` is left alone: a Python float is a no-op, a tensor broadcasts. - if gate_up_lora is not None: - _gate_up_first, _gate_up_second, _gate_up_scaling = gate_up_lora - gate_up_lora = ( - _gate_up_first.to(hidden_states.dtype), - _gate_up_second.to(hidden_states.dtype), - _gate_up_scaling, - ) - - down_lora = take_moe_lora_stash(self, "down_proj") - if down_lora is not None: - down_lora = down_lora[:3] - elif ( - use_separated_lora - and hasattr(self, "down_proj") - and _has_lora_adapters(self.down_proj) - ): - down_lora = _extract_lora_weights( - self.down_proj, num_experts=self.num_experts, experts_module=self - ) - if down_lora is not None: - _down_first, _down_second, _down_scaling = down_lora - down_lora = ( - _down_first.to(hidden_states.dtype), - _down_second.to(hidden_states.dtype), - _down_scaling, - ) - - # Stable sort keeps torch.where(one_hot)'s (slot, token) row order: bit-identical results. - with torch.no_grad(): - n_tokens = top_k_index.shape[0] - slot_major = top_k_index.t().reshape(-1) # position = top_k_pos * n_tokens + token - expert_counts, order, _ = moe_routing_sort(slot_major, self.num_experts) - routed_token_idx = order % n_tokens - routed_top_k_pos = order // n_tokens - expert_counts = expert_counts.tolist() - - # A declared (E, in, out) layout (_unsloth_grouped_mm_format) wins over the shape test, which a square stack defeats. - grouped_mm_format = bool(_module_flag(self, "_unsloth_grouped_mm_format", False)) or ( - _module_flag(self, "is_transposed", None) is True - ) - - # Interleaved storage is a layout flag; GPT-OSS's clamped swiglu is its own activation: keep them apart. - interleaved = _gate_up_is_interleaved(self) - is_gpt_oss = "GptOssExperts" in self.__class__.__name__ - own_apply_gate = _uses_own_apply_gate(self) - - row_end = 0 - for expert_idx, count in enumerate(expert_counts): - row_start, row_end = row_end, row_end + count - if count == 0: - continue - token_idx = routed_token_idx[row_start:row_end] - top_k_pos = routed_top_k_pos[row_start:row_end] - current_state = hidden_states[token_idx] - - # gate_up projection for this expert ('gate_up_proj' or 'w1'/'w3'). - if hasattr(self, "gate_up_proj"): - gate_up_weight = self.gate_up_proj[expert_idx] - if grouped_mm_format or gate_up_weight.shape[-1] != current_state.shape[-1]: - gate_up_weight = gate_up_weight.T - gate_up = F.linear(current_state, gate_up_weight) - if gate_up_lora is not None: - first_weight, second_weight, scaling = gate_up_lora - lora_delta = current_state @ first_weight[expert_idx] - lora_delta = lora_delta @ second_weight[expert_idx] - gate_up = gate_up + lora_delta * scaling - gate_up_bias = getattr(self, "gate_up_proj_bias", None) - if gate_up_bias is not None: - gate_up = gate_up + gate_up_bias[expert_idx].to(gate_up.dtype) - if interleaved: - gate = gate_up[..., ::2] - up = gate_up[..., 1::2] - else: - gate, up = gate_up.chunk(2, dim=-1) - else: - gate = F.linear(current_state, self.w1[expert_idx]) - up = F.linear(current_state, self.w3[expert_idx]) - gate_up = None - - if own_apply_gate: - current_hidden_states = self._apply_gate( - gate_up if gate_up is not None else torch.cat((gate, up), dim=-1) - ) - elif is_gpt_oss: - limit = getattr(self, "limit", 7.0) - alpha = getattr(self, "alpha", 1.702) - gate = gate.clamp(min=None, max=limit) - up = up.clamp(min=-limit, max=limit) - current_hidden_states = (up + 1.0) * (gate * torch.sigmoid(gate * alpha)) - elif hasattr(self, "act_fn") and callable(self.act_fn): - current_hidden_states = self.act_fn(gate) * up - else: - current_hidden_states = F.silu(gate) * up - - # down projection for this expert. - if hasattr(self, "down_proj"): - down_weight = self.down_proj[expert_idx] - # Mirror gate_up: prefer the flag over the shape heuristic (unsafe at square dims). - if grouped_mm_format or down_weight.shape[-1] != current_hidden_states.shape[-1]: - down_weight = down_weight.T - down = F.linear(current_hidden_states, down_weight) - if down_lora is not None: - first_weight, second_weight, scaling = down_lora - lora_delta = current_hidden_states @ first_weight[expert_idx] - lora_delta = lora_delta @ second_weight[expert_idx] - down = down + lora_delta * scaling - down_bias = getattr(self, "down_proj_bias", None) - if down_bias is not None: - down = down + down_bias[expert_idx].to(down.dtype) - current_hidden_states = down - else: - current_hidden_states = F.linear(current_hidden_states, self.w2[expert_idx]) - - current_hidden_states = ( - current_hidden_states * top_k_weights[token_idx, top_k_pos, None] - ) - - final_hidden_states.index_add_( - 0, token_idx, current_hidden_states.to(final_hidden_states.dtype) - ) - - return final_hidden_states diff --git a/unsloth_compiled_cache/unsloth_compiled_module_gemma4.py b/unsloth_compiled_cache/unsloth_compiled_module_gemma4.py deleted file mode 100644 index 4e4737ff5cacd77d76bc118b30372c9a9b157147..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/unsloth_compiled_module_gemma4.py +++ /dev/null @@ -1,2127 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - - -import os -import sys -import torch -import importlib.util -import math -if importlib.util.find_spec("unsloth_studio") is None: - UNSLOTH_STUDIO_ENABLED = False -else: - UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" -pass -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -import math - -UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" -UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" -UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) -UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") -if UNSLOTH_COMPILE_LOCATION not in sys.path: - sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) - -import logging -logger_compiler = logging.getLogger(__name__) -if UNSLOTH_ENABLE_LOGGING: - logger_compiler.setLevel(logging.DEBUG) - -global INFERENCE_RUNS -INFERENCE_RUNS = 0 - -try: - import torch._dynamo.eval_frame as torch_dynamo_eval_frame - torch_dynamo_eval_frame._stance.stance - torch_compiler_set_stance = torch.compiler.set_stance -except: - torch_dynamo_eval_frame = None - torch_compiler_set_stance = None -pass - -# The stance is process wide, so whether Unsloth owns it lives in unsloth_zoo, shared by -# every generated module: inference in one model's module, training in another's. -from unsloth_zoo.temporary_patches.utils import UNSLOTH_EAGER_STANCE_OWNED, unsloth_claim_eager_stance, unsloth_owns_stance, unsloth_release_stance - -def unsloth_training_stance(): - # eager_on_recompile is global: left on by generate() or eval forwards, every recompile - # in a later training step (new sequence length, model.train() without for_training) - # would silently run the eager body. Training forwards restore the default stance, - # unless the user chose eager_on_recompile themselves. Every set_stance call installs a - # fresh stance object, so ours is still active only if the current one `is` it; a user - # set_stance since then (even to eager_on_recompile again) is theirs and stays. - global INFERENCE_RUNS - INFERENCE_RUNS = 0 - if len(UNSLOTH_EAGER_STANCE_OWNED) != 0 and torch_dynamo_eval_frame is not None: - current = torch_dynamo_eval_frame._stance - # Ownership is only dropped once ours is reset: a temporary stance scope's exit restores - # one of our exact objects, which a later training forward must still recognise. - if current.stance == "eager_on_recompile" and unsloth_owns_stance(current): - # Only the reset object: an outer scope may still restore another one of ours. - # Put back what the caller had before Unsloth switched (force_eager, a backend, ...). - previous = unsloth_release_stance(current) - if previous is not None and hasattr(torch_dynamo_eval_frame, "_set_stance"): - torch_dynamo_eval_frame._set_stance(previous) - else: - torch_compiler_set_stance(stance = "default", skip_guard_eval_unsafe = False) -pass - -from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT -from unsloth_zoo.temporary_patches.utils import UNSLOTH_DECODE_COMPILE - - -from unsloth_zoo.loss_utils import ( - HAS_CUT_CROSS_ENTROPY, - fused_linear_cross_entropy, - unsloth_fused_ce_loss, -) -from unsloth_zoo.fused_losses.forward_adapter import unsloth_fused_lm_head_loss -from unsloth_zoo.fused_losses.forward_install import _can_fuse_loss -from unsloth_zoo.compiler import lm_head_on_old_cuda_arch -from unsloth_zoo.fused_losses.cross_entropy_loss import unsloth_count_aware_cross_entropy, unsloth_loss_count_kwargs, unsloth_fsdp_unshard -from unsloth_zoo.fused_losses.aux_loss import unsloth_ga_scale_aux_loss - -scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention -@torch.compiler.disable(recursive = False) -def disable_compile_scaled_dot_product_attention(*args, **kwargs): - return scaled_dot_product_attention(*args, **kwargs) -pass - - -from transformers.modeling_flash_attention_utils import is_flash_attn_available - -if is_flash_attn_available(): - try: - from transformers.modeling_flash_attention_utils import flash_attn_supports_top_left_mask - except: - flash_attn_supports_top_left_mask = None - try: - from transformers.modeling_flash_attention_utils import _flash_attention_forward - except: - _flash_attention_forward = None - try: - from transformers.modeling_flash_attention_utils import FlashAttentionKwargs - except: - FlashAttentionKwargs = None - try: - from transformers.modeling_flash_attention_utils import flash_attn_varlen_func - except: - flash_attn_varlen_func = None -else: - flash_attn_supports_top_left_mask = None - _flash_attention_forward = None - FlashAttentionKwargs = None - flash_attn_varlen_func = None -pass - - -torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 32, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': False, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} - -from torch.nn import CrossEntropyLoss -from unsloth_zoo.temporary_patches.utils import torch_compile_with_fallback - -@torch_compile_with_fallback(fullgraph = True, dynamic = True, options = torch_compile_options) -def normal_cross_entropy_loss(self, hidden_states, labels): - logits = self.lm_head(hidden_states) - logits = logits.float() - # Shift so that tokens < n predict n - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - # Flatten the tokens - loss_fct = CrossEntropyLoss() - shift_logits = shift_logits.view(-1, self.config.vocab_size) - shift_labels = shift_labels.view(-1) - # Enable model parallelism - shift_labels = shift_labels.to(shift_logits.device) - loss = loss_fct(shift_logits, shift_labels) - return loss, logits -pass - -# We need an empty logits flag to warn people logits will not be returned anymore unless asked ie -# os.environ['UNSLOTH_RETURN_LOGITS'] = '1' -LOGITS_ERROR_STRING = \ - "Unsloth: Logits are empty from 2024.11 onwards. To get raw logits again, please "\ - 'set the environment variable `UNSLOTH_RETURN_LOGITS` to `"1" BEFORE starting to train ie before `trainer.train()`. For example:\n'\ - "```\nimport os\n"\ - "os.environ['UNSLOTH_RETURN_LOGITS'] = '1'\n"\ - "trainer.train()\n```\n"\ - "No need to restart your console - just add `os.environ['UNSLOTH_RETURN_LOGITS'] = '1'` before trainer.train() and re-run the cell!" - -def raise_logits_error(*args, **kwargs): raise NotImplementedError(LOGITS_ERROR_STRING) -def return_none(*args, **kwargs): return None -class EmptyLogits: - def __init__(self): return - def raise_getattr_error(self, attr): - if attr == "to": return return_none - # A catch-all __getattr__ makes hasattr() true for every name, dunders included, so the - # sentinel answers yes to protocol probes it cannot honour. torch.distributed's output cast - # tests `hasattr(x, "__dataclass_fields__")` and then calls `dataclasses.replace(x)` on - # whatever said yes, so with FSDP2 mixed precision every step died in - # `TypeError: replace() should be called on dataclass instances` (unsloth#409, reached - # through `_fsdp_state._cast_output_dtype`). Protocol probes get an honest AttributeError; - # ordinary attribute access still gets the callable that explains UNSLOTH_RETURN_LOGITS. - if len(attr) > 4 and attr.startswith("__") and attr.endswith("__"): - raise AttributeError(f"{type(self).__name__!r} object has no attribute {attr!r}") - return raise_logits_error - __getitem__ = raise_logits_error - __getattr__ = raise_getattr_error - def __repr__(self): return LOGITS_ERROR_STRING - def __str__ (self): return LOGITS_ERROR_STRING - # Stateless pickling so accelerate gather_object works on the sentinel - def __reduce__(self): return (type(self), ()) - # Gathered copies must compare equal in accelerate debug mode - def __eq__(self, other): return type(other).__name__ == "EmptyLogits" - __hash__ = object.__hash__ -pass -EMPTY_LOGITS = EmptyLogits() -functions = dir(torch.Tensor) -for j, function in enumerate(functions): - if function.startswith("__") and function.endswith("__"): - exec(f"def raise_{j}(*args, **kwargs): print('{function}')", globals(), locals()) - try: exec(f"EMPTY_LOGITS.{function} = raise_{j}", globals(), locals()) - except: continue -pass -# The loop above stomps pickle hooks with stubs returning None; restore them. -for function in ("__reduce__", "__reduce_ex__", "__getstate__", "__setstate__"): - try: delattr(EMPTY_LOGITS, function) - except Exception: pass -pass - - -def mask_attention_mask_out(labels = None, attention_mask = None): - # Only a mask over the labels' own positions: MusicGen's attention_mask is the text encoder's, - # while its labels are decoder codebook frames. - if ( - labels is not None and attention_mask is not None - and tuple(attention_mask.shape) == tuple(labels.shape) - ): - attention_mask = attention_mask.to(device = labels.device) - labels[attention_mask == 0] = -100 - return labels -pass - - -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from unsloth_zoo.temporary_patches.utils import torch_compile_with_fallback -from unsloth_zoo.temporary_patches.utils import torch_compiler_disable_unless_decode -from unsloth_zoo.temporary_patches.common import torch_compile -import transformers -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from transformers.models.gemma4.modeling_gemma4 import (F, math, Callable, Optional, torch, nn, init, ACT2FN, Cache, PreTrainedConfig, GenerationMixin, create_bidirectional_mask, create_causal_mask, create_sliding_window_causal_mask, FlashAttentionKwargs, BaseModelOutputWithPast, ModelOutput, CausalLMOutputWithPast, ROPE_INIT_FUNCTIONS, dynamic_rope_update, ALL_ATTENTION_FUNCTIONS, PreTrainedModel, Unpack, TransformersKwargs, can_return_tuple, maybe_autocast, capture_outputs, Gemma4AudioConfig, Gemma4Config, Gemma4TextConfig, Gemma4VisionConfig, Gemma4CausalLMOutputWithPast, Gemma4AudioCausalConv1d, Gemma4VisionEncoderLayer, Gemma4PreTrainedModel, Gemma4TextModel, Gemma4ForCausalLM, Gemma4VisionModel, Gemma4Model, Gemma4ForConditionalGeneration, create_masks_for_generate, _to_boundary_dtype, _linear_boundary_dtype, Gemma4MultimodalEmbedder, Gemma4VisionPooler, Gemma4TextExperts, Gemma4VisionPatchEmbedder, Gemma4VisionEncoder, create_bidirectional_mask, create_causal_mask, create_masks_for_generate, create_sliding_window_causal_mask) - -@torch_compile_with_fallback(fullgraph = False, dynamic = True, options = torch_compile_options) -def Gemma4ClippableLinear_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - if self.use_clipped_linears: - hidden_states = torch.clamp(hidden_states, self.input_min, self.input_max) - - hidden_states = self.linear(hidden_states) - - if self.use_clipped_linears: - hidden_states = torch.clamp(hidden_states, self.output_min, self.output_max) - - return hidden_states - -class Gemma4ClippableLinear(nn.Module): - def __init__( - self, - config: Gemma4VisionConfig | Gemma4AudioConfig, - in_features: int, - out_features: int, - ) -> None: - super().__init__() - self.use_clipped_linears = config.use_clipped_linears - self.linear = nn.Linear(in_features, out_features, bias=False) - - if self.use_clipped_linears: - self.register_buffer("input_min", torch.tensor(-float("inf"))) - self.register_buffer("input_max", torch.tensor(float("inf"))) - self.register_buffer("output_min", torch.tensor(-float("inf"))) - self.register_buffer("output_max", torch.tensor(float("inf"))) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - return Gemma4ClippableLinear_forward(self, hidden_states=hidden_states) - - -@torch_compile_with_fallback(fullgraph = True, dynamic = True, options = torch_compile_options) -def Gemma4RMSNorm_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - normed_output = self._norm(hidden_states.float()) - if self.with_scale: - normed_output = normed_output * self.weight.float() - return normed_output.type_as(hidden_states) - -class Gemma4RMSNorm(nn.Module): - def __init__(self, dim: int, eps: float = 1e-6, with_scale: bool = True): - super().__init__() - self.eps = eps - self.with_scale = with_scale - - if self.with_scale: - self.weight = nn.Parameter(torch.ones(dim), requires_grad=True) - - def _norm(self, hidden_states: torch.Tensor): - mean_squared = hidden_states.pow(2).mean(-1, keepdim=True) + self.eps - # Use torch.pow() (over torch.sqrt() or torch.rsqrt()) to addess compiler differences between Torch and JAX - return hidden_states * torch.pow(mean_squared, -0.5) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - return Gemma4RMSNorm_forward(self, hidden_states=hidden_states) - - -@torch_compile_with_fallback(fullgraph = True, dynamic = True, options = torch_compile_options) -@torch.no_grad() -def Gemma4AudioRelPositionalEncoding_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - position_ids = torch.arange(12, -1, -1, device=hidden_states.device) - position_ids = position_ids[..., None] - scaled_time = position_ids * self.inv_timescales.to(device=hidden_states.device) - pos_embed = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=-1) - return pos_embed.to(dtype=hidden_states.dtype) - -class Gemma4AudioRelPositionalEncoding(nn.Module): - """Sinusoidal relative positional encoding for the audio encoder. - - Produces position embeddings of shape [1, 2*context_size - 1, hidden_size] with - concatenated [sin..., cos...] layout matching the original Gemma4 convention. - """ - - inv_timescales: torch.Tensor - - def __init__(self, config: Gemma4AudioConfig): - super().__init__() - self.hidden_size = config.hidden_size - self.context_size = ( - config.attention_chunk_size + config.attention_context_left - 1 + config.attention_context_right - ) - min_timescale = 1.0 - max_timescale = 10000.0 - num_timescales = self.hidden_size // 2 - log_timescale_increment = math.log(max_timescale / min_timescale) / max(num_timescales - 1, 1) - inv_timescales = min_timescale * torch.exp(torch.arange(num_timescales) * -log_timescale_increment) - self.register_buffer("inv_timescales", inv_timescales.unsqueeze(0).unsqueeze(0), persistent=False) - - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - return Gemma4AudioRelPositionalEncoding_forward(self, hidden_states=hidden_states) - - -@torch_compile_with_fallback(fullgraph = False, dynamic = True, options = torch_compile_options) -def Gemma4AudioAttention_forward( - self, - hidden_states: torch.Tensor, - position_embeddings: torch.Tensor, - attention_mask: torch.BoolTensor | None = None, -) -> tuple[torch.Tensor, None]: - batch_size, seq_length, _ = hidden_states.shape - hidden_shape = (batch_size, seq_length, self.num_heads, self.head_dim) - - query_states = self.q_proj(hidden_states).float().view(hidden_shape) - key_states = self.k_proj(hidden_states).float().view(hidden_shape) - value_states = self.v_proj(hidden_states).float().view(hidden_shape) - - query_states = query_states * self.q_scale * F.softplus(self.per_dim_scale) - key_states = key_states * self.k_scale - - query_states = self._convert_to_block(query_states) - key_states = self._extract_block_context(key_states) - value_states = self._extract_block_context(value_states) - num_blocks = query_states.shape[1] - - relative_key_states = self.relative_k_proj(position_embeddings) - relative_key_states = relative_key_states.view(-1, self.num_heads, self.head_dim) - relative_key_states = relative_key_states.to(dtype=query_states.dtype) - - queries = query_states.permute(0, 3, 1, 2, 4) - matrix_ac = queries @ key_states.permute(0, 3, 1, 4, 2) - - queries_flat = queries.reshape(batch_size, self.num_heads, -1, self.head_dim) - matrix_bd = queries_flat @ relative_key_states.permute(1, 2, 0) - matrix_bd = matrix_bd.reshape(batch_size, self.num_heads, num_blocks, self.chunk_size, -1) - matrix_bd = self._rel_shift(matrix_bd) - - attn_weights = matrix_ac + matrix_bd - attn_weights = attn_weights / self.softcap - attn_weights = torch.tanh(attn_weights) - attn_weights = attn_weights * self.softcap - - if attention_mask is not None: - attn_weights = attn_weights.masked_fill( - attention_mask.logical_not(), self.config.attention_invalid_logits_value - ) - - attn_weights = F.softmax(attn_weights, dim=-1, dtype = torch.float32).to(attn_weights.dtype).to(value_states.dtype) - attn_output = attn_weights @ value_states.permute(0, 3, 1, 2, 4) - attn_output = attn_output.permute(0, 2, 3, 1, 4).reshape(batch_size, num_blocks * self.chunk_size, -1) - attn_output = attn_output[:, :seq_length].contiguous() - attn_output = self.post(attn_output.to(dtype=self.post.linear.weight.dtype)) - - return attn_output, attn_weights - -class Gemma4AudioAttention(nn.Module): - """Chunked local attention with relative position bias""" - - def __init__(self, config: Gemma4AudioConfig, layer_idx: int): - super().__init__() - self.config = config - self.layer_idx = layer_idx - self.attention_logits_soft_cap = config.attention_logit_cap - self.head_dim = config.hidden_size // config.num_attention_heads - self.num_heads = config.num_attention_heads - - self.q_scale = (self.head_dim**-0.5) / math.log(2) - self.k_scale = math.log(1 + math.e) / math.log(2) - - self.chunk_size = config.attention_chunk_size - self.max_past_horizon = config.attention_context_left - 1 - self.max_future_horizon = config.attention_context_right - self.context_size = self.chunk_size + self.max_past_horizon + self.max_future_horizon - - self.q_proj = Gemma4ClippableLinear(config, config.hidden_size, self.num_heads * self.head_dim) - self.k_proj = Gemma4ClippableLinear(config, config.hidden_size, self.num_heads * self.head_dim) - self.v_proj = Gemma4ClippableLinear(config, config.hidden_size, self.num_heads * self.head_dim) - self.post = Gemma4ClippableLinear(config, config.hidden_size, config.hidden_size) - - self.relative_k_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False) - self.per_dim_scale = nn.Parameter(torch.zeros(self.head_dim)) - - self.register_buffer("softcap", torch.tensor(self.attention_logits_soft_cap), persistent=False) - - def _convert_to_block(self, hidden_states: torch.Tensor) -> torch.Tensor: - """Splits a `(batch_size, seq_len, num_heads, head_dim)` tensor into non-overlapping blocks of `chunk_size` along the sequence dim.""" - batch_size, seq_len, num_heads, head_dim = hidden_states.shape - num_blocks = (seq_len + self.chunk_size - 1) // self.chunk_size - pad = num_blocks * self.chunk_size - seq_len - hidden_states = F.pad(hidden_states, (0, 0, 0, 0, 0, pad)) - return hidden_states.reshape(batch_size, num_blocks, self.chunk_size, num_heads, head_dim).contiguous() - - def _extract_block_context(self, hidden_states: torch.Tensor) -> torch.Tensor: - """Extracts overlapping context windows of `context_size` for every block, strided by `chunk_size`.""" - batch_size, seq_len, num_heads, head_dim = hidden_states.shape - hidden_states = F.pad( - hidden_states, (0, 0, 0, 0, self.max_past_horizon, self.max_future_horizon + self.chunk_size - 1) - ) - hidden_states = hidden_states.unfold(1, self.context_size, self.chunk_size) - hidden_states = torch.movedim(hidden_states, -1, 2) - return hidden_states.contiguous() - - def _rel_shift(self, x: torch.Tensor) -> torch.Tensor: - """Relative position shift for blocked attention. See appendix B of https://huggingface.co/papers/1901.02860.""" - batch_size, num_heads, num_blocks, block_size, position_length = x.shape - context_size = self.context_size - x = F.pad(x, (0, context_size + 1 - position_length)) - x = x.view(batch_size, num_heads, num_blocks, block_size * (context_size + 1)) - x = x[..., : block_size * context_size] - return x.view(batch_size, num_heads, num_blocks, block_size, context_size) - - def forward( - self, - hidden_states: torch.Tensor, - position_embeddings: torch.Tensor, - attention_mask: torch.BoolTensor | None = None, - ) -> tuple[torch.Tensor, None]: - return Gemma4AudioAttention_forward(self, hidden_states=hidden_states, position_embeddings=position_embeddings, attention_mask=attention_mask) - - -@torch_compile_with_fallback(fullgraph = True, dynamic = True, options = torch_compile_options) -def Gemma4AudioSubSampleConvProjectionLayer_forward(self, hidden_states: torch.Tensor, mask: torch.Tensor | None = None): - if mask is not None: - mask = mask.to(device=hidden_states.device) - hidden_states = hidden_states * mask[:, None, :, None] - - hidden_states = self.conv(hidden_states.to(self.conv.weight.dtype)) - hidden_states = self.act(self.norm(hidden_states.permute(0, 2, 3, 1)).permute(0, 3, 1, 2).contiguous()) - - if mask is not None: - mask = mask[:, ::2] - - return hidden_states, mask - -class Gemma4AudioSubSampleConvProjectionLayer(nn.Module): - def __init__(self, in_channels, out_channels, norm_eps): - super().__init__() - self.conv = nn.Conv2d( - in_channels=in_channels, - out_channels=out_channels, - kernel_size=(3, 3), - stride=(2, 2), - padding=1, - bias=False, - ) - self.norm = nn.LayerNorm(out_channels, eps=norm_eps, elementwise_affine=True, bias=False) - self.act = nn.ReLU() - - def forward(self, hidden_states: torch.Tensor, mask: torch.Tensor | None = None): - return Gemma4AudioSubSampleConvProjectionLayer_forward(self, hidden_states=hidden_states, mask=mask) - - -@torch_compile_with_fallback(fullgraph = False, dynamic = True, options = torch_compile_options) -def Gemma4AudioSubSampleConvProjection_forward( - self, - input_features: torch.Tensor, - input_features_mask: torch.Tensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - hidden_states = input_features.unsqueeze(1) - hidden_states, mask = self.layer0(hidden_states, input_features_mask) - hidden_states, mask = self.layer1(hidden_states, mask) - - batch_size, _, seq_len, _ = hidden_states.shape - hidden_states = hidden_states.permute(0, 2, 3, 1).contiguous().reshape(batch_size, seq_len, -1) - return self.input_proj_linear(hidden_states), mask - -class Gemma4AudioSubSampleConvProjection(nn.Module): - def __init__(self, config: Gemma4AudioConfig): - super().__init__() - self.layer0 = Gemma4AudioSubSampleConvProjectionLayer( - in_channels=1, - out_channels=config.subsampling_conv_channels[0], - norm_eps=config.rms_norm_eps, - ) - self.layer1 = Gemma4AudioSubSampleConvProjectionLayer( - in_channels=config.subsampling_conv_channels[0], - out_channels=config.subsampling_conv_channels[1], - norm_eps=config.rms_norm_eps, - ) - proj_input_dim = (config.subsampling_conv_channels[0] // 4) * config.subsampling_conv_channels[1] - self.input_proj_linear = nn.Linear(proj_input_dim, config.hidden_size, bias=False) - - def forward( - self, - input_features: torch.Tensor, - input_features_mask: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - return Gemma4AudioSubSampleConvProjection_forward(self, input_features=input_features, input_features_mask=input_features_mask) - - -@torch_compile_with_fallback(fullgraph = False, dynamic = True, options = torch_compile_options) -def Gemma4AudioFeedForward_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - # This is needed to avoid any underflow/overflow issues when clipping - gradient_clipping = min(self.gradient_clipping, torch.finfo(self.ffw_layer_1.linear.weight.dtype).max) - - residual = hidden_states - hidden_states = torch.clamp(hidden_states, -gradient_clipping, gradient_clipping) - hidden_states = self.pre_layer_norm(hidden_states) - - hidden_states = self.ffw_layer_1(hidden_states) - hidden_states = self.act_fn(hidden_states) - hidden_states = self.ffw_layer_2(hidden_states) - - hidden_states = torch.clamp(hidden_states, -gradient_clipping, gradient_clipping) - hidden_states = self.post_layer_norm(hidden_states) - hidden_states *= self.post_layer_scale - hidden_states += residual - - return hidden_states - -class Gemma4AudioFeedForward(nn.Module): - def __init__(self, config: Gemma4AudioConfig): - super().__init__() - self.config = config - - self.ffw_layer_1 = Gemma4ClippableLinear(config, config.hidden_size, config.hidden_size * 4) - self.ffw_layer_2 = Gemma4ClippableLinear(config, config.hidden_size * 4, config.hidden_size) - - self.pre_layer_norm = Gemma4RMSNorm(config.hidden_size) - self.post_layer_norm = Gemma4RMSNorm(config.hidden_size) - self.act_fn = ACT2FN[config.hidden_act] - - self.gradient_clipping = config.gradient_clipping - self.post_layer_scale = config.residual_weight - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - return Gemma4AudioFeedForward_forward(self, hidden_states=hidden_states) - - -@torch_compile_with_fallback(fullgraph = False, dynamic = True, options = torch_compile_options) -def Gemma4AudioLightConv1d_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - residual = hidden_states - - hidden_states = self.pre_layer_norm(hidden_states) - hidden_states = self.linear_start(hidden_states) - hidden_states = nn.functional.glu(hidden_states, dim=-1) - - hidden_states = self.depthwise_conv1d(hidden_states.transpose(1, 2)).transpose(1, 2) - - # This is needed to avoid any underflow/overflow issues when clipping - gradient_clipping = min(self.gradient_clipping, torch.finfo(self.linear_start.linear.weight.dtype).max) - hidden_states = torch.clamp(hidden_states, -gradient_clipping, gradient_clipping) - hidden_states = self.conv_norm(hidden_states) - - hidden_states = self.act_fn(hidden_states) - hidden_states = self.linear_end(hidden_states) - hidden_states += residual - return hidden_states - -class Gemma4AudioLightConv1d(nn.Module): - def __init__(self, config: Gemma4AudioConfig): - super().__init__() - self.config = config - - self.linear_start = Gemma4ClippableLinear(config, config.hidden_size, config.hidden_size * 2) - self.linear_end = Gemma4ClippableLinear(config, config.hidden_size, config.hidden_size) - self.depthwise_conv1d = Gemma4AudioCausalConv1d( - in_channels=config.hidden_size, - out_channels=config.hidden_size, - kernel_size=config.conv_kernel_size, - groups=config.hidden_size, - bias=False, - ) - - self.pre_layer_norm = Gemma4RMSNorm(config.hidden_size, eps=config.rms_norm_eps, with_scale=True) - self.conv_norm = Gemma4RMSNorm(config.hidden_size, eps=config.rms_norm_eps, with_scale=True) - self.act_fn = ACT2FN[config.hidden_act] - - self.gradient_clipping = config.gradient_clipping - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - return Gemma4AudioLightConv1d_forward(self, hidden_states=hidden_states) - - -@torch_compiler_disable_unless_decode -def Gemma4VisionPatchEmbedder_forward( - self, pixel_values: torch.Tensor, pixel_position_ids: torch.Tensor, padding_positions: torch.Tensor -) -> torch.Tensor: - # Gemma4 applies no normalization and instead scales in model code - pixel_values = 2 * (pixel_values - 0.5) - hidden_states = self.input_proj(pixel_values.to(self.input_proj.weight.dtype)) - position_embeddings = self._position_embeddings(pixel_position_ids, padding_positions) - return hidden_states + position_embeddings - -class Gemma4VisionPatchEmbedder(nn.Module): - def __init__(self, config: Gemma4VisionConfig): - super().__init__() - self.config = config - self.hidden_size = config.hidden_size - self.patch_size = config.patch_size - self.position_embedding_size = config.position_embedding_size - - self.input_proj = nn.Linear(3 * self.patch_size**2, self.hidden_size, bias=False) - self.position_embedding_table = nn.Parameter(torch.ones(2, self.position_embedding_size, self.hidden_size)) - - def _position_embeddings(self, pixel_position_ids: torch.Tensor, padding_positions: torch.Tensor) -> torch.Tensor: - """Prepare patch positions map for matmul with positon embedding table.""" - # Expanding and permute patch positions to (batch_size, num_patches, 2, position_embedding_size) for matmul. - clamped_positions = pixel_position_ids.clamp(min=0) - one_hot = F.one_hot(clamped_positions, num_classes=self.position_embedding_size) - one_hot = one_hot.permute(0, 2, 1, 3).to(self.position_embedding_table) - # Compute positional embeddings and sum across x and y. - position_embeddings = one_hot @ self.position_embedding_table - position_embeddings = position_embeddings.sum(dim=1) - # Zero out embeddings for any padding patches. - position_embeddings = torch.where(padding_positions.unsqueeze(-1), 0.0, position_embeddings) - return position_embeddings - - def forward( - self, pixel_values: torch.Tensor, pixel_position_ids: torch.Tensor, padding_positions: torch.Tensor - ) -> torch.Tensor: - return Gemma4VisionPatchEmbedder_forward(self, pixel_values=pixel_values, pixel_position_ids=pixel_position_ids, padding_positions=padding_positions) - - -@torch_compile_with_fallback(fullgraph = False, dynamic = True, options = torch_compile_options) -def Gemma4VisionMLP_forward(self, x): - down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) - return down_proj - -class Gemma4VisionMLP(nn.Module): - def __init__(self, config: Gemma4VisionConfig): - super().__init__() - self.config = config - self.hidden_size = config.hidden_size - self.intermediate_size = config.intermediate_size - self.gate_proj = Gemma4ClippableLinear(config, self.hidden_size, self.intermediate_size) - self.up_proj = Gemma4ClippableLinear(config, self.hidden_size, self.intermediate_size) - self.down_proj = Gemma4ClippableLinear(config, self.intermediate_size, self.hidden_size) - self.act_fn = ACT2FN[config.hidden_activation] - - def forward(self, x): - return Gemma4VisionMLP_forward(self, x=x) - - -@torch_compile_with_fallback(fullgraph = False, dynamic = True, options = torch_compile_options) -@torch.no_grad() -@dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) -def Gemma4VisionRotaryEmbedding_forward(self, x, position_ids): - inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) - device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" - - # Multidimensional positions: [batch, num_patches, ndim]. Apply rotations to each spatial dim separately - all_cos, all_sin = [], [] - for i in range(2): - dim_position_ids = position_ids[:, :, i] - dim_position_ids_expanded = dim_position_ids[:, None, :].float() - - with maybe_autocast(device_type=device_type, enabled=False): # Force float32 - freqs = (inv_freq_expanded.float() @ dim_position_ids_expanded.float()).transpose(1, 2) - emb = torch.cat((freqs, freqs), dim=-1) - cos = emb.cos() * self.attention_scaling - sin = emb.sin() * self.attention_scaling - all_cos.append(cos) - all_sin.append(sin) - - cos = torch.cat(all_cos, dim=-1).to(dtype=x.dtype) - sin = torch.cat(all_sin, dim=-1).to(dtype=x.dtype) - return cos, sin - -class Gemma4VisionRotaryEmbedding(nn.Module): - inv_freq: torch.Tensor # fix linting for `register_buffer` - - def __init__(self, config: Gemma4VisionConfig, device=None): - super().__init__() - self.max_seq_len_cached = config.max_position_embeddings - self.original_max_seq_len = config.max_position_embeddings - - self.config = config - - self.rope_type = self.config.rope_parameters["rope_type"] - rope_init_fn: Callable = self.compute_default_rope_parameters - if self.rope_type != "default": - rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] - inv_freq, self.attention_scaling = rope_init_fn(self.config, device) - - self.register_buffer("inv_freq", inv_freq, persistent=False) - self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False) - - @staticmethod - def compute_default_rope_parameters( - config: Gemma4VisionConfig | None = None, - device: torch.device | None = None, - seq_len: int | None = None, - ) -> tuple["torch.Tensor", float]: - """ - Computes the inverse frequencies according to the original RoPE implementation - Args: - config ([`~transformers.PreTrainedConfig`]): - The model configuration. - device (`torch.device`): - The device to use for initialization of the inverse frequencies. - seq_len (`int`, *optional*): - The current sequence length. Unused for this type of RoPE. - Returns: - Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the - post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). - """ - base = config.rope_parameters["rope_theta"] - dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads - - # The reference implementation computes RoPE frequencies INDEPENDENTLY - # for each spatial dimension using the partitioned head_dim (head_dim // ndim), - # so both x and y dimensions get identical frequency ranges. - # This is different from splitting the global inv_freq between dimensions. - spatial_dim = dim // 2 - - attention_factor = 1.0 # Unused in this type of RoPE - inv_freq = 1.0 / ( - base - ** (torch.arange(0, spatial_dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / spatial_dim) - ) - return inv_freq, attention_factor - - - def forward(self, x, position_ids): - return Gemma4VisionRotaryEmbedding_forward(self, x=x, position_ids=position_ids) - - -@torch_compile_with_fallback(fullgraph = True, dynamic = True, options = torch_compile_options) -def rotate_half(x): - """Rotates half the hidden dims of the input.""" - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -@torch_compile_with_fallback(fullgraph = True, dynamic = True, options = torch_compile_options) -def apply_rotary_pos_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, unsqueeze_dim: int = 1): - """Applies Rotary Position Embedding to the query and key tensors. - - Args: - x (`torch.Tensor`): The tensor to embed. - cos (`torch.Tensor`): The cosine part of the rotary embedding. - sin (`torch.Tensor`): The sine part of the rotary embedding. - unsqueeze_dim (`int`, *optional*, defaults to 1): - The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and - sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes - cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have - the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. - Returns: - `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. - """ - cos = cos.unsqueeze(unsqueeze_dim) - sin = sin.unsqueeze(unsqueeze_dim) - return (x * cos) + (rotate_half(x) * sin) - - -@torch_compile_with_fallback(fullgraph = True, dynamic = True, options = torch_compile_options) -def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: - """ - This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, - num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) - """ - batch, num_key_value_heads, slen, head_dim = hidden_states.shape - if n_rep == 1: - return hidden_states - hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) - return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) - - -@torch_compile_with_fallback(fullgraph = True, dynamic = True, options = torch_compile_options) -def eager_attention_forward( - module: nn.Module, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attention_mask: torch.Tensor | None, - dropout: float | int = 0.0, - scaling: float | None = None, - softcap: float | None = None, - **kwargs, -) -> tuple[torch.Tensor, torch.Tensor]: - if scaling is None: - scaling = module.head_dim**-0.5 - - key_states = repeat_kv(key, module.num_key_value_groups) - value_states = repeat_kv(value, module.num_key_value_groups) - - attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling - - if softcap is not None: - attn_weights = attn_weights / softcap - attn_weights = torch.tanh(attn_weights) - attn_weights = attn_weights * softcap - if attention_mask is not None: - - if isinstance(attention_mask, dict): - - attention_mask = attention_mask.get(getattr(module, 'layer_type', None), None) - - if attention_mask is not None: - - attn_weights = attn_weights + attention_mask - - # upcast attention to fp32 - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype = torch.float32).to(attn_weights.dtype).to(query.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) - attn_output = torch.matmul(attn_weights, value_states) - attn_output = attn_output.transpose(1, 2).contiguous() - return attn_output, attn_weights - - -@torch_compile_with_fallback(fullgraph = True, dynamic = True, options = torch_compile_options) -def apply_multidimensional_rope( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - position_ids: torch.Tensor, - unsqueeze_dim: int = 2, -) -> torch.Tensor: - """Applies multidimensional RoPE to inputs. - - Args: - x (`torch.Tensor`): The tensor to embed. - cos (`torch.Tensor`): The cosine part of the rotary embedding. - sin (`torch.Tensor`): The sine part of the rotary embedding. - position_ids (`torch.Tensor`, *optional*): - If position_ids.ndim + 2 == x.ndim, then this function passes through to `apply_rotary_pos_emb()`. - Otherwise, position_ids is used to split the inputs, x, into multiple pieces, where each piece is fed to - `apply_rotary_pos_emb()`, and then concatenated back together. - unsqueeze_dim (`int`, *optional*, defaults to 1): - The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and - sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes - cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have - the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. - - Returns: - Tensor of shape [B, L, N, H] with RoPE applied. - """ - ndim = position_ids.shape[-1] - num_input_channels = x.shape[-1] - num_rotated_channels_per_dim = 2 * (num_input_channels // (2 * ndim)) - - if num_rotated_channels_per_dim <= 0: - raise ValueError( - "Invalid configuration: num_rotated_channels_per_dim must be > 0, got" - f" {num_rotated_channels_per_dim} (num_input_channels={num_input_channels}," - f" ndim={ndim})" - ) - - # Correctly split the input tensor into ndim parts - split_sizes = [num_rotated_channels_per_dim] * ndim - x_parts = torch.split(x, split_sizes, dim=-1) - cos_parts = torch.split(cos, split_sizes, dim=-1) - sin_parts = torch.split(sin, split_sizes, dim=-1) - y_parts = [ - apply_rotary_pos_emb( - x=x_parts[k], - cos=cos_parts[k], - sin=sin_parts[k], - unsqueeze_dim=unsqueeze_dim, - ) - for k in range(ndim) - ] - return torch.cat(y_parts, dim=-1) - - -@torch_compiler_disable_unless_decode -def Gemma4VisionAttention_forward( - self, - hidden_states: torch.Tensor, - position_embeddings: torch.Tensor = None, - attention_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - **kwargs: Unpack[TransformersKwargs], -) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: - input_shape = hidden_states.shape[:-1] - hidden_shape = (*input_shape, -1, self.head_dim) - - cos, sin = position_embeddings - - query_states = self.q_proj(hidden_states).view(hidden_shape) - query_states = self.q_norm(query_states) - query_states = apply_multidimensional_rope(query_states, cos, sin, position_ids) - query_states = query_states.transpose(1, 2) - - key_states = self.k_proj(hidden_states).view(hidden_shape) - key_states = self.k_norm(key_states) - key_states = apply_multidimensional_rope(key_states, cos, sin, position_ids) - key_states = key_states.transpose(1, 2) - - value_states = self.v_proj(hidden_states).view(hidden_shape) - value_states = self.v_norm(value_states) - value_states = value_states.transpose(1, 2) - - attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( - self.config._attn_implementation, eager_attention_forward - ) - - attn_output, attn_weights = attention_interface( - self, - query_states, - key_states, - value_states, - attention_mask, - dropout=self.attention_dropout if self.training else 0.0, - scaling=self.scaling, - **kwargs, - ) - - attn_output = attn_output.reshape(*input_shape, -1).contiguous() - attn_output = self.o_proj(attn_output) - return attn_output, attn_weights - -class Gemma4VisionAttention(nn.Module): - """Multi-headed attention from 'Attention Is All You Need' paper""" - - def __init__(self, config: Gemma4VisionConfig, layer_idx: int): - super().__init__() - self.layer_type = config.layer_types[layer_idx] if hasattr(config, "layer_types") else None - self.config = config - self.layer_idx = layer_idx - self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) - self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads - self.scaling = 1.0 - self.attention_dropout = self.config.attention_dropout - self.is_causal = False - self.q_proj = Gemma4ClippableLinear(config, config.hidden_size, config.num_attention_heads * self.head_dim) - self.k_proj = Gemma4ClippableLinear(config, config.hidden_size, config.num_key_value_heads * self.head_dim) - self.v_proj = Gemma4ClippableLinear(config, config.hidden_size, config.num_key_value_heads * self.head_dim) - self.o_proj = Gemma4ClippableLinear(config, config.num_attention_heads * self.head_dim, config.hidden_size) - - self.q_norm = Gemma4RMSNorm(dim=config.head_dim, eps=config.rms_norm_eps) - self.k_norm = Gemma4RMSNorm(dim=config.head_dim, eps=config.rms_norm_eps) - self.v_norm = Gemma4RMSNorm(self.head_dim, eps=config.rms_norm_eps, with_scale=False) - - def forward( - self, - hidden_states: torch.Tensor, - position_embeddings: torch.Tensor = None, - attention_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - **kwargs: Unpack[TransformersKwargs], - ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: - return Gemma4VisionAttention_forward(self, hidden_states=hidden_states, position_embeddings=position_embeddings, attention_mask=attention_mask, position_ids=position_ids, **kwargs) - - -@torch_compiler_disable_unless_decode -def Gemma4VisionEncoder_forward( - self, - inputs_embeds: torch.Tensor, - attention_mask: torch.Tensor, - pixel_position_ids: torch.LongTensor | None = None, - **kwargs: Unpack[TransformersKwargs], -) -> BaseModelOutputWithPast: - r""" - pixel_position_ids (torch.Tensor): - Patch positions as (x, y) coordinates in the image as [batch, num_patches, 2]. - """ - attention_mask = create_bidirectional_mask( - config=self.config, - inputs_embeds=inputs_embeds, - attention_mask=attention_mask, - ) - - # embed positions - hidden_states = inputs_embeds - position_embeddings = self.rotary_emb(hidden_states, pixel_position_ids) - - # decoder layers - for decoder_layer in self.layers[: self.config.num_hidden_layers]: - hidden_states = decoder_layer( - hidden_states, - attention_mask=attention_mask, - position_embeddings=position_embeddings, - position_ids=pixel_position_ids, - **kwargs, - ) - - return BaseModelOutputWithPast(last_hidden_state=hidden_states) - -class Gemma4VisionEncoder(nn.Module): - def __init__(self, config: Gemma4VisionConfig): - super().__init__() - self.config = config - self.num_layers = config.num_hidden_layers - self.rotary_emb = Gemma4VisionRotaryEmbedding(config) - self.layers = nn.ModuleList( - [Gemma4VisionEncoderLayer(config=config, layer_idx=i) for i in range(self.num_layers)] - ) - - def forward( - self, - inputs_embeds: torch.Tensor, - attention_mask: torch.Tensor, - pixel_position_ids: torch.LongTensor | None = None, - **kwargs: Unpack[TransformersKwargs], - ) -> BaseModelOutputWithPast: - return Gemma4VisionEncoder_forward(self, inputs_embeds=inputs_embeds, attention_mask=attention_mask, pixel_position_ids=pixel_position_ids, **kwargs) - - -@torch_compile_with_fallback(fullgraph = False, dynamic = True, options = torch_compile_options) -def Gemma4TextMLP_forward(self, x): - # Forced float32 full finetuning: x leaves the RMSNorm as fp16, the Linears are fp32 and - # nothing autocasts, so meet the weight. Identity whenever they already agree. - x = _to_boundary_dtype(x, _linear_boundary_dtype(self, "gate_proj")) - gate = self.gate_proj(x) - # Check matmul output dtype to catch autocast / PEFT fp16 casts. - if gate.dtype != torch.float16: - return self.down_proj(self.act_fn(gate) * self.up_proj(x)) - product = self.act_fn(gate.float()) * self.up_proj(x).float() - product = torch.clamp(product, min=-_SAFE_FP16, max=_SAFE_FP16) - out = self.down_proj(product.to(gate.dtype)) - # Zero overflows so the residual identity path survives. - return torch.nan_to_num(out, nan=0.0, posinf=0.0, neginf=0.0) - -class Gemma4TextMLP(nn.Module): - def __init__(self, config: Gemma4TextConfig, layer_idx: int): - super().__init__() - first_kv_shared_layer_idx = config.num_hidden_layers - config.num_kv_shared_layers - is_kv_shared_layer = layer_idx >= first_kv_shared_layer_idx > 0 - use_double_wide_mlp = config.use_double_wide_mlp and is_kv_shared_layer - self.config = config - self.hidden_size = config.hidden_size - self.intermediate_size = config.intermediate_size * (2 if use_double_wide_mlp else 1) - self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) - self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) - self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) - self.act_fn = ACT2FN[config.hidden_activation] - - def forward(self, x): - down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) - return down_proj - - -@torch_compile_with_fallback(fullgraph = False, dynamic = True, options = torch_compile_options) -@torch.no_grad() -@dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) -def Gemma4TextRotaryEmbedding_forward(self, x, position_ids, layer_type=None): - inv_freq = getattr(self, f"{layer_type}_inv_freq") - attention_scaling = getattr(self, f"{layer_type}_attention_scaling") - - inv_freq_expanded = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) - position_ids_expanded = position_ids[:, None, :].float() - - device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" - with maybe_autocast(device_type=device_type, enabled=False): # Force float32 - freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) - emb = torch.cat((freqs, freqs), dim=-1) - cos = emb.cos() * attention_scaling - sin = emb.sin() * attention_scaling - - return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) - -class Gemma4TextRotaryEmbedding(nn.Module): - inv_freq: torch.Tensor # fix linting for `register_buffer` - - def __init__(self, config: Gemma4TextConfig, device=None, layer_type=None): - super().__init__() - self.max_seq_len_cached = config.max_position_embeddings - self.original_max_seq_len = config.max_position_embeddings - - self.config = config - self.layer_types = set(config.layer_types) - self.rope_init_fns: dict[str, Callable[..., tuple[torch.Tensor, float]]] = {} - self.rope_type: dict[str, str] = {} - - for layer_type in self.layer_types: - rope_params = self.config.rope_parameters[layer_type] - if rope_params is None: - continue - - if (rope_type := rope_params["rope_type"]) != "default": - rope_init_fn = ROPE_INIT_FUNCTIONS[rope_type] - else: - rope_init_fn = self.compute_default_rope_parameters - - self.rope_init_fns[layer_type] = rope_init_fn - self.rope_type[layer_type] = rope_type - - rope_init_fn_kwargs = {"device": device, "layer_type": layer_type} - if layer_type == "full_attention" and rope_type == "proportional": - rope_init_fn_kwargs["head_dim_key"] = "global_head_dim" - - curr_inv_freq, curr_attention_scaling = rope_init_fn(self.config, **rope_init_fn_kwargs) - self.register_buffer(f"{layer_type}_inv_freq", curr_inv_freq, persistent=False) - self.register_buffer(f"{layer_type}_original_inv_freq", curr_inv_freq.clone(), persistent=False) - setattr(self, f"{layer_type}_attention_scaling", curr_attention_scaling) - - @staticmethod - def compute_default_rope_parameters( - config: Gemma4TextConfig | None = None, - device: Optional["torch.device"] = None, - seq_len: int | None = None, - layer_type: str | None = None, - ) -> tuple["torch.Tensor", float]: - """ - Computes the inverse frequencies according to the original RoPE implementation - Args: - config ([`~transformers.PreTrainedConfig`]): - The model configuration. - device (`torch.device`): - The device to use for initialization of the inverse frequencies. - seq_len (`int`, *optional*): - The current sequence length. Unused for this type of RoPE. - layer_type (`str`, *optional*): - The current layer type if the model has different RoPE parameters per type. - Should not be used unless `config.layer_types is not None` - - Returns: - Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the - post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). - """ - # For backward compatibility standardize the `rope_parameters_dict` if it uses old format - base = config.rope_parameters[layer_type]["rope_theta"] - dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads - - attention_factor = 1.0 # Unused in this type of RoPE - - # Compute the inverse frequencies - inv_freq = 1.0 / ( - base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) - ) - return inv_freq, attention_factor - - - def forward(self, x, position_ids, layer_type=None): - return Gemma4TextRotaryEmbedding_forward(self, x=x, position_ids=position_ids, layer_type=layer_type) - - -@torch_compiler_disable_unless_decode -def Gemma4TextAttention_forward(self, *args, **kwargs): - carrier = getattr(self, "_unsloth_shared_kv_carrier", None) - if carrier is not None: - # past_key_values is the 4th positional arg (after self) or a kwarg. - if "past_key_values" in kwargs: - if kwargs["past_key_values"] is None: - kwargs["past_key_values"] = carrier - elif len(args) >= 4: - if args[3] is None: - args = args[:3] + (carrier,) + args[4:] - else: - kwargs["past_key_values"] = carrier - return _orig_attn_forward(self, *args, **kwargs) - -class Gemma4TextAttention(nn.Module): - """Multi-headed attention from 'Attention Is All You Need' paper""" - - def __init__(self, config: Gemma4TextConfig, layer_idx: int): - super().__init__() - self.layer_type = config.layer_types[layer_idx] if hasattr(config, "layer_types") else None - self.config = config - self.layer_idx = layer_idx - self.is_sliding = self.layer_type == "sliding_attention" - self.sliding_window = config.sliding_window if self.is_sliding else None - - self.head_dim = config.global_head_dim if not self.is_sliding and config.global_head_dim else config.head_dim - self.use_alternative_attention = config.attention_k_eq_v and not self.is_sliding - num_key_value_heads = ( - config.num_global_key_value_heads if self.use_alternative_attention else config.num_key_value_heads - ) - self.num_key_value_groups = config.num_attention_heads // num_key_value_heads - self.scaling = 1.0 - self.attention_dropout = self.config.attention_dropout - self.is_causal = config.use_bidirectional_attention != "all" - - # Shared kv cache - first_kv_shared_layer_idx = self.config.num_hidden_layers - getattr(self.config, "num_kv_shared_layers", 0) - self.is_kv_shared_layer = layer_idx >= first_kv_shared_layer_idx > 0 - prev_layers = config.layer_types[:first_kv_shared_layer_idx] - if self.is_kv_shared_layer: - # For shared layers, find the last non-shared layer of the same type before sharing starts - self.kv_shared_layer_index = len(prev_layers) - 1 - prev_layers[::-1].index(config.layer_types[layer_idx]) - self.store_full_length_kv = False - else: - self.kv_shared_layer_index = None - # For non-shared layers, store full-length kv if this is the last non-shared layer of its type - self.store_full_length_kv = layer_idx == len(prev_layers) - 1 - prev_layers[::-1].index( - config.layer_types[layer_idx] - ) - - self.q_norm = Gemma4RMSNorm(dim=self.head_dim, eps=config.rms_norm_eps) - self.k_norm = Gemma4RMSNorm(dim=self.head_dim, eps=config.rms_norm_eps) - self.v_norm = Gemma4RMSNorm(self.head_dim, eps=config.rms_norm_eps, with_scale=False) - - self.k_proj = nn.Linear(config.hidden_size, num_key_value_heads * self.head_dim, bias=config.attention_bias) - self.q_proj = nn.Linear( - config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias - ) - self.v_proj = ( - nn.Linear(config.hidden_size, num_key_value_heads * self.head_dim, bias=config.attention_bias) - if not self.use_alternative_attention - else None - ) - self.o_proj = nn.Linear( - config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias - ) - - def forward( - self, - hidden_states: torch.Tensor, - position_embeddings: torch.Tensor, - attention_mask: torch.Tensor | None, - past_key_values: Cache | None = None, - **kwargs: Unpack[FlashAttentionKwargs], - ) -> tuple[torch.Tensor, torch.Tensor | None]: - input_shape = hidden_states.shape[:-1] - hidden_shape = (*input_shape, -1, self.head_dim) - - cos, sin = position_embeddings - - query_states = self.q_proj(hidden_states).view(hidden_shape) - query_states = self.q_norm(query_states) - query_states = apply_rotary_pos_emb(query_states, cos, sin, unsqueeze_dim=2) - query_states = query_states.transpose(1, 2) - - # For layers with shared KV (from kv sharing point onwards), we reuse the same keys/values states as the last non-sharing layer - if self.is_kv_shared_layer and past_key_values is not None: - key_states, value_states = past_key_values.shared_layers[self.kv_shared_layer_index] - # Device of past layer may be different from current one - key_states = key_states.to(query_states.device) - value_states = value_states.to(query_states.device) - else: - key_states = self.k_proj(hidden_states).view(hidden_shape) - value_states = self.v_proj(hidden_states).view(hidden_shape) if self.v_proj is not None else key_states - - key_states = self.k_norm(key_states) - key_states = apply_rotary_pos_emb(key_states, cos, sin, unsqueeze_dim=2) - key_states = key_states.transpose(1, 2) - - value_states = self.v_norm(value_states) - value_states = value_states.transpose(1, 2) - - if past_key_values is not None: - if not self.is_kv_shared_layer: - key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) - if self.store_full_length_kv: - if not hasattr(past_key_values, "shared_layers"): - past_key_values.shared_layers = {} - past_key_values.shared_layers[self.layer_idx] = key_states, value_states - - attention_interface: Callable = eager_attention_forward - if self.config._attn_implementation != "eager": - attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] - - attn_output, attn_weights = attention_interface( - self, - query_states, - key_states, - value_states, - attention_mask, - dropout=self.attention_dropout if self.training else 0.0, - scaling=self.scaling, - sliding_window=self.sliding_window, - **kwargs, - ) - - attn_output = attn_output.reshape(*input_shape, -1).contiguous() - attn_output = self.o_proj(attn_output) - return attn_output, attn_weights - - - -def _gemma4_experts_forward(self, hidden_states, top_k_index, top_k_weights): - # Gemma4TextRouter.forward already folds per_expert_scale into - # top_k_weights, so dispatch straight to the grouped-GEMM backend. - return _moe_backend(self, hidden_states, top_k_index, top_k_weights) - -class Gemma4TextExperts(nn.Module): - """Collection of expert weights stored as 3D tensors.""" - - def __init__(self, config: Gemma4TextConfig): - super().__init__() - self.num_experts = config.num_experts - self.hidden_dim = config.hidden_size - self.intermediate_dim = config.moe_intermediate_size - self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim)) - self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim)) - self.act_fn = ACT2FN[config.hidden_activation] - - def forward( - self, - hidden_states: torch.Tensor, - top_k_index: torch.Tensor, - top_k_weights: torch.Tensor, - ) -> torch.Tensor: - return _gemma4_experts_forward(self, hidden_states=hidden_states, top_k_index=top_k_index, top_k_weights=top_k_weights) - - -@torch_compile_with_fallback(fullgraph = False, dynamic = True, options = torch_compile_options) -def Gemma4TextRouter_forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - hidden_states = self.norm(hidden_states) - hidden_states = hidden_states * self.scale * self.scalar_root_size - - expert_scores = self.proj(hidden_states) # [B*S, E] - router_probabilities = nn.functional.softmax(expert_scores, dim=-1, dtype = torch.float32).to(expert_scores.dtype) - - # topk returns both values (probabilities) and indices directly - top_k_weights, top_k_index = torch.topk( - router_probabilities, - k=self.config.top_k_experts, - dim=-1, - ) # both [B*S, K] - - # Normalize the top-k weights so they sum to 1 per token - top_k_weights /= top_k_weights.sum(dim=-1, keepdim=True) - - # Apply per-expert scale directly to the weights - top_k_weights = top_k_weights * self.per_expert_scale[top_k_index] - - return router_probabilities, top_k_weights, top_k_index - -class Gemma4TextRouter(nn.Module): - def __init__(self, config: Gemma4TextConfig): - super().__init__() - self.config = config - self.hidden_size = config.hidden_size - self.scalar_root_size = self.hidden_size**-0.5 - self.eps = config.rms_norm_eps - - self.norm = Gemma4RMSNorm(self.hidden_size, eps=self.eps, with_scale=False) - self.proj = nn.Linear(config.hidden_size, config.num_experts, bias=False) - self.scale = nn.Parameter(torch.ones(self.hidden_size)) - self.per_expert_scale = nn.Parameter(torch.ones(config.num_experts)) - - def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - return Gemma4TextRouter_forward(self, hidden_states=hidden_states) - - -@torch_compiler_disable_unless_decode -@can_return_tuple -def Gemma4ForCausalLM_forward( - self, - input_ids: torch.LongTensor | None = None, - attention_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - past_key_values: Cache | None = None, - inputs_embeds: torch.FloatTensor | None = None, - labels: torch.LongTensor | None = None, - use_cache: bool | None = None, - logits_to_keep: int | torch.Tensor = 0, - **kwargs: Unpack[TransformersKwargs], -) -> CausalLMOutputWithPast: - r""" - Example: - - ```python - >>> from transformers import AutoTokenizer, Gemma4ForCausalLM - - >>> model = Gemma4ForCausalLM.from_pretrained("google/gemma-2-9b") - >>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b") - - >>> prompt = "What is your favorite condiment?" - >>> inputs = tokenizer(prompt, return_tensors="pt") - - >>> # Generate - >>> generate_ids = model.generate(inputs.input_ids, max_length=30) - >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] - "What is your favorite condiment?" - ```""" - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) - if not torch.compiler.is_compiling() and torch.is_grad_enabled() and getattr(self, "training", False): - unsloth_training_stance() - outputs: BaseModelOutputWithPast = self.model( - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - inputs_embeds=inputs_embeds, - use_cache=use_cache, - **kwargs, - ) - - hidden_states = outputs.last_hidden_state - # Only compute necessary logits, and do not upcast them to float if we are not computing the loss - slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep - logits = self.lm_head(hidden_states[:, slice_indices, :]) if os.environ.get('UNSLOTH_RETURN_LOGITS', '0') == '1' else EMPTY_LOGITS - loss = None - NOT_RETURN_LOGITS = os.environ.get('UNSLOTH_RETURN_LOGITS', '0') == '0' - RETURN_HIDDEN_STATES = os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES", "0") == "1" - - n_items = None - if (kwargs) != () and type(kwargs) is dict: - n_items = (kwargs).get("num_items_in_batch", None) - if n_items is None: n_items = (kwargs).get("n_items", None) - if n_items is None: - all_locals = locals() - if 'loss_kwargs' in all_locals: - __kwargs = all_locals['loss_kwargs'] - if type(__kwargs) is dict: - n_items = __kwargs.get("num_items_in_batch", None) - if n_items is None: n_items = __kwargs.get("n_items", None) - if n_items is None and 'kwargs' in all_locals: - __kwargs = all_locals['kwargs'] - if type(__kwargs) is dict: - n_items = __kwargs.get("num_items_in_batch", None) - if n_items is None: n_items = __kwargs.get("n_items", None) - if n_items is None: - all_locals = all_locals.values() - for __kwargs in all_locals: - if type(__kwargs) is dict: - n_items = __kwargs.get("num_items_in_batch", None) - if n_items is None: n_items = __kwargs.get("n_items", None) - break - pass - # Stock ForCausalLMLoss uses caller pre-shifted targets (context parallel, padding-free) as is. - explicit_shift_labels = None - if (kwargs) != () and type(kwargs) is dict: - explicit_shift_labels = (kwargs).get("shift_labels", None) - if not torch.is_tensor(explicit_shift_labels): explicit_shift_labels = None - - if labels is not None: unsloth_fsdp_unshard(self.lm_head) - requires_grad_ = self.lm_head.weight.requires_grad - requires_grad_ = requires_grad_ or self.lm_head.weight.dtype == torch.float32 - - if RETURN_HIDDEN_STATES: - logits = hidden_states[:, slice_indices, :] - elif labels is None: - - - # Set compiler stance to fail on recompiles for inference - global INFERENCE_RUNS - # Skipped while tracing (set_stance raises there, and the counter would guard every step) - # and around a compiled decode step, which eager_on_recompile would otherwise freeze. - if not torch.compiler.is_compiling() and not UNSLOTH_DECODE_COMPILE[0]: - if torch_dynamo_eval_frame is not None: - old_stance = torch_dynamo_eval_frame._stance.stance - else: - old_stance = None - if old_stance is not None and INFERENCE_RUNS == 1: - # Skip guards and return to eager -> we still need guards! - # Already eager (another model's inference, or the user's choice): leave that stance - # object, and so its ownership, as it is. Otherwise switch and claim the new object. - if old_stance != "eager_on_recompile": - previous = getattr(torch_dynamo_eval_frame, "_stance", None) - torch_compiler_set_stance(stance = "eager_on_recompile", skip_guard_eval_unsafe = False) - # The stance object itself, or True where torch keeps no such object. - unsloth_claim_eager_stance(getattr(torch_dynamo_eval_frame, "_stance", True), previous) - if UNSLOTH_ENABLE_LOGGING: - logger_compiler.info( - f"Unsloth: Removing compiler guards after 1 inference run. " \ - f"DYNAMO_STANCE.stance = {torch_dynamo_eval_frame._stance.stance} " \ - f"DYNAMO_STANCE.skip_guard_eval_unsafe = {torch_dynamo_eval_frame._stance.skip_guard_eval_unsafe}" - ) - elif old_stance == "eager_on_recompile": - pass - elif old_stance == "default" and INFERENCE_RUNS > 1: - # Reset compiler stance - torch_compiler_set_stance(stance = "default", skip_guard_eval_unsafe = False) - if UNSLOTH_ENABLE_LOGGING: - logger_compiler.info( - f"Unsloth: Reseting guards. " \ - f"DYNAMO_STANCE.stance = {torch_dynamo_eval_frame._stance.stance} " \ - f"DYNAMO_STANCE.skip_guard_eval_unsafe = {torch_dynamo_eval_frame._stance.skip_guard_eval_unsafe}" - ) - INFERENCE_RUNS = 0 - INFERENCE_RUNS += 1 - - logits = self.lm_head(hidden_states[:, slice_indices, :]) - # Inference returns these logits to the caller, so they must carry the same - # scale and softcap transforms the training branches below apply. Leaving - # them off does not change greedy decoding, since tanh is monotonic, but it - # does change the distribution: sampling, returned logprobs and any scoring - # that reads the logits all see uncapped values. - if () != (): - logits = logits * () - if () != (): - logits = logits / () - if (self.config.final_logit_softcapping) not in (None, (),): - logits = logits / (self.config.final_logit_softcapping) - logits = torch.tanh(logits) - logits = logits * (self.config.final_logit_softcapping) - elif (() == () and () == ()) and (UNSLOTH_ENABLE_CCE and HAS_CUT_CROSS_ENTROPY) and not lm_head_on_old_cuda_arch(self.lm_head.weight) and NOT_RETURN_LOGITS and _can_fuse_loss(self.loss_function) and labels is not None and not requires_grad_ and explicit_shift_labels is None and getattr(self.lm_head, "bias", None) is None: - loss = fused_linear_cross_entropy( - hidden_states = hidden_states[:, slice_indices, :], - lm_weight = self.lm_head.weight, - labels = labels.to(self.lm_head.weight.device), - num_items_in_batch = n_items, - logit_softcapping = None if (self.config.final_logit_softcapping) == () else (self.config.final_logit_softcapping), - ) - elif _can_fuse_loss(self.loss_function) and labels is not None and NOT_RETURN_LOGITS: - lm_head_weight = self.lm_head.weight - lm_head_bias = getattr(self.lm_head, "bias", None) - - # ========= NEW fused ========= - _hidden_states = hidden_states[:, slice_indices, :] - _labels = labels if explicit_shift_labels is None else explicit_shift_labels - torch._dynamo.mark_dynamic(_hidden_states, 1) - torch._dynamo.mark_dynamic(_labels, _labels.dim() - 1) - loss = unsloth_fused_ce_loss( - trainer = None, - hidden_states = _hidden_states, - lm_head_weight = lm_head_weight, - lm_head_bias = lm_head_bias, - labels = _labels, - mask = None, - n_items = n_items, - scaling = getattr(self, "accelerator_scaler", None), - target_gb = None, - torch_compile = not UNSLOTH_COMPILE_DISABLE and not lm_head_on_old_cuda_arch(lm_head_weight), - logit_scale_multiply = () if () != () else 0, - logit_scale_divide = () if () != () else 0, - logit_softcapping = (self.config.final_logit_softcapping) if (self.config.final_logit_softcapping) != () else 0, - shift_labels = explicit_shift_labels is None, - ) - elif _can_fuse_loss(self.loss_function) and labels is not None and not NOT_RETURN_LOGITS: - # UNSLOTH_RETURN_LOGITS=1 path. Prepended `logits = self.lm_head(...)` - # already materialised the full lm_head matmul; apply the captured logit - # scale/softcap transforms and route loss through self.loss_function on - # those logits instead of letting unsloth_fused_ce_loss redo the matmul. - if () != (): - logits = logits * () - if () != (): - logits = logits / () - if (self.config.final_logit_softcapping) not in (None, (),): - logits = logits / (self.config.final_logit_softcapping) - logits = torch.tanh(logits) - logits = logits * (self.config.final_logit_softcapping) - loss = self.loss_function(logits, labels.to(self.lm_head.weight.device), vocab_size=self.vocab_size, **kwargs) - else: - logits = self.lm_head(hidden_states[:, slice_indices, :]) - if () != (): - logits = logits * () - if () != (): - logits = logits / () - if (self.config.final_logit_softcapping) not in (None, (),): - logits = logits / (self.config.final_logit_softcapping) - logits = torch.tanh(logits) - logits = logits * (self.config.final_logit_softcapping) - loss = self.loss_function(logits, labels.to(self.lm_head.weight.device), vocab_size=self.vocab_size, **kwargs) - - - return CausalLMOutputWithPast( - loss=loss, - logits=logits, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - -class Gemma4ForCausalLM(Gemma4PreTrainedModel, GenerationMixin): - _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} - _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} - config: Gemma4TextConfig - base_model_prefix = "model" - - def __init__(self, config: Gemma4TextConfig): - super().__init__(config) - self.model = Gemma4TextModel(config) - self.vocab_size = config.vocab_size - self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) - - # Initialize weights and apply final processing - self.post_init() - - - def forward( - self, - input_ids: torch.LongTensor | None = None, - attention_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - past_key_values: Cache | None = None, - inputs_embeds: torch.FloatTensor | None = None, - labels: torch.LongTensor | None = None, - use_cache: bool | None = None, - logits_to_keep: int | torch.Tensor = 0, - **kwargs: Unpack[TransformersKwargs], - ) -> CausalLMOutputWithPast: - return Gemma4ForCausalLM_forward(self, input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, labels=labels, use_cache=use_cache, logits_to_keep=logits_to_keep, **kwargs) - -Gemma4ForCausalLM.__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__ = True - - -def sliding_window_mask_function(sliding_window: tuple[int, int]) -> Callable: - """ - This creates uni/bidirectional attention mask with sliding window. - """ - - def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool: - left_window_size, right_window_size = sliding_window - - dist = q_idx - kv_idx - left_mask = (dist >= 0) & (dist < left_window_size) - right_mask = (dist < 0) & (-dist < right_window_size) - return left_mask | right_mask - - return inner_mask - - -@torch_compiler_disable_unless_decode -@capture_outputs - -def Gemma4VisionModel_forward( - self, - pixel_values: torch.FloatTensor, - pixel_position_ids: torch.LongTensor, - **kwargs: Unpack[TransformersKwargs], -) -> BaseModelOutputWithPast: - r""" - pixel_values (`torch.FloatTensor` or `list[torch.FloatTensor]`): - The images to encode. Either a single `[batch, channels, height, width]` tensor - (all images same size) or a list of `[1, channels, height, width]` tensors (different sizes). - pixel_position_ids (`torch.LongTensor` of shape `(batch_size, max_patches, 2)`): - The patch positions as (x, y) coordinates in the image. Padding patches are indicated by (-1, -1). - """ - pooling_kernel_size = self.config.pooling_kernel_size - output_length = pixel_values.shape[-2] // (pooling_kernel_size * pooling_kernel_size) - - padding_positions = (pixel_position_ids == -1).all(dim=-1) - inputs_embeds = self.patch_embedder(pixel_values, pixel_position_ids, padding_positions) - output = self.encoder( - inputs_embeds=inputs_embeds, - attention_mask=~padding_positions, # encoder expects True=valid, padding_positions is True=padding - pixel_position_ids=pixel_position_ids, - **kwargs, - ) - - hidden_states, pooler_mask = self.pooler( - hidden_states=output.last_hidden_state, - pixel_position_ids=pixel_position_ids, - padding_positions=padding_positions, - output_length=output_length, - ) - - # Strip padding tokens. pooler_mask is True = valid, False = padding. - hidden_states = hidden_states[pooler_mask] - - if self.config.standardize: - hidden_states = (hidden_states - self.std_bias) * self.std_scale - - return BaseModelOutputWithPast(last_hidden_state=hidden_states) - -class Gemma4VisionModel(Gemma4PreTrainedModel): - """The Gemma 4 Vision Encoder.""" - - config = Gemma4VisionConfig - _can_record_outputs = { - "hidden_states": Gemma4VisionEncoderLayer, - "attentions": Gemma4VisionAttention, - } - - def __init__(self, config: Gemma4VisionConfig): - super().__init__(config) - self.patch_embedder = Gemma4VisionPatchEmbedder(config) - self.encoder = Gemma4VisionEncoder(config) - self.pooler = Gemma4VisionPooler(config) - - if self.config.standardize: - self.register_buffer("std_bias", torch.empty(self.config.hidden_size)) - self.register_buffer("std_scale", torch.empty(self.config.hidden_size)) - - self.post_init() - - - def forward( - self, - pixel_values: torch.FloatTensor, - pixel_position_ids: torch.LongTensor, - **kwargs: Unpack[TransformersKwargs], - ) -> BaseModelOutputWithPast: - return Gemma4VisionModel_forward(self, pixel_values=pixel_values, pixel_position_ids=pixel_position_ids, **kwargs) - - -@torch_compiler_disable_unless_decode -def Gemma4MultimodalEmbedder_forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor: - old_dtype = inputs_embeds.dtype - emb_norm = _Gemma4MultimodalEmbedder_RMSNorm_forward(self.embedding_pre_projection_norm, inputs_embeds) - # Call the module, never `.weight`: a PEFT lora.Linear applies its delta - # only in forward, so reading the weight trains the adapter to no effect. - projection = self.embedding_projection - # Match the weight dtype: only a float32 projector computes in float32. - # Promoting a half one needs an autocast several backends refuse, or a - # Parameter rebuild Dynamo cannot trace. - weight = getattr(projection, "weight", None) - compute_dtype = torch.float32 if weight is None else weight.dtype - # An enclosing bf16 autocast downcasts even float32 operands. Guarded - # because autocast rejects meta and unregistered custom backends. - try: - autocast_off = torch.autocast(device_type = emb_norm.device.type, enabled = False) - except (RuntimeError, AssertionError): - autocast_off = contextlib.nullcontext() - with autocast_off: - emb_norm_proj = projection(emb_norm.to(compute_dtype)) - return emb_norm_proj.to(old_dtype) - -class Gemma4MultimodalEmbedder(nn.Module): - """Embeds token ids or soft tokens for multimodal content into language model space.""" - - def __init__( - self, - multimodal_config: Gemma4AudioConfig | Gemma4VisionConfig, - text_config: Gemma4TextConfig, - ): - super().__init__() - - self.multimodal_hidden_size = getattr(multimodal_config, "output_proj_dims", multimodal_config.hidden_size) - self.eps = multimodal_config.rms_norm_eps - self.text_hidden_size = text_config.hidden_size - self.embedding_projection = nn.Linear(self.multimodal_hidden_size, self.text_hidden_size, bias=False) - self.embedding_pre_projection_norm = Gemma4RMSNorm(self.multimodal_hidden_size, eps=self.eps, with_scale=False) - - def forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor: - """Embeds token ids or soft tokens for multimodal content into language model space. - Args: - inputs_embeds: A torch.Tensor containing the soft tokens to embed. - Returns: - A torch.Tensor of embeddings with shape `[batch_size, seq_len, self.config.text_config.hidden_size]`. - """ - embs_normed = self.embedding_pre_projection_norm(inputs_embeds) - return self.embedding_projection(embs_normed) - - -def token_type_ids_mask_function( - token_type_ids: torch.Tensor | None, - image_group_ids: torch.Tensor | None, -) -> Callable | None: - """ - This function adds the correct offsets to the `q_idx` and `kv_idx` as the torch API can only accept lengths, - not start and end indices. - """ - # Do not return an additional mask in this case - if token_type_ids is None: - return None - - def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool: - seq_length = image_group_ids.shape[-1] - - # clamp indices because with static cache they can go beyond `image_group_ids.shape[-1]` - q_idx_clamped = q_idx.clamp(max=seq_length - 1) - kv_idx_clamped = kv_idx.clamp(max=seq_length - 1) - - # Unmask if the q and kv come from same group which is not -1 (i.e. non-text) - q_group = image_group_ids[batch_idx, q_idx_clamped] - kv_group = image_group_ids[batch_idx, kv_idx_clamped] - q_group = torch.where(q_idx < seq_length, q_group, -1) - kv_group = torch.where(kv_idx < seq_length, kv_group, -1) - return (q_group == kv_group) & (q_group >= 0) - - return inner_mask - - -def create_causal_mask_mapping( - config: PreTrainedConfig, - inputs_embeds: torch.Tensor, - attention_mask: torch.Tensor | None, - past_key_values: Cache | None, - position_ids: torch.Tensor | None, - mm_token_type_ids: torch.Tensor | None = None, - pixel_values: torch.FloatTensor | None = None, - is_training: bool = False, - is_first_iteration: bool | None = None, - **kwargs, -) -> dict: - """ - Overwrites the base `create_masks_for_generate` with `token_type_ids` masking to create the causal mask mapping - for all kinds of forward passes. Gemma4 uses a bidirectional mask for images. - - Uses `pixel_values` as an optional input to disambiguate edge cases. - """ - if is_training and mm_token_type_ids is None: - raise ValueError("`mm_token_type_ids` is required as a model input when training") - - mask_kwargs = { - "config": config.get_text_config(), - "inputs_embeds": inputs_embeds, - "attention_mask": attention_mask, - "past_key_values": past_key_values, - "position_ids": position_ids, - } - sliding_mask_kwargs = mask_kwargs.copy() - - # NOTE: this `may_have_image_input` logic is not flawless, it fails when we're using a cache eagerly initialized - # (e.g. compiled prefill) AND `pixel_values` are not provided (i.e. the image data is provided through other - # means). Determining prefill in that case requires checking data values, which is not compile-compatible. - is_first_iteration = ( - is_first_iteration - if is_first_iteration is not None - else (past_key_values is None or not past_key_values.is_initialized or pixel_values is not None) - ) - if mm_token_type_ids is not None and is_first_iteration: - # We need to pass an additional mask function to account for token type ids, and it needs to be an `or` (to - # undo the causal masking) - - # First find where a new vision block starts. Vision tokens cannot attend to - # future vision tokens, but can attend to all prev tokens and to itself bidirectionally - is_vision = (mm_token_type_ids == 1) | (mm_token_type_ids == 2) - is_prev_vision = torch.roll(is_vision, shifts=1, dims=-1) - is_prev_vision[..., 0] = False - new_vision_starts = is_vision & ~is_prev_vision - vision_group_ids = torch.cumsum(new_vision_starts.int(), dim=1) - 1 - vision_group_ids = torch.where(is_vision, vision_group_ids, -1) - sliding_mask_kwargs["or_mask_function"] = token_type_ids_mask_function( - mm_token_type_ids.to(inputs_embeds.device), vision_group_ids - ) - - return { - "full_attention": create_causal_mask(**mask_kwargs), - "sliding_attention": create_sliding_window_causal_mask(**sliding_mask_kwargs), - } - - -@torch_compiler_disable_unless_decode -@can_return_tuple -def Gemma4ForConditionalGeneration_forward( - self, - input_ids: torch.LongTensor | None = None, - pixel_values: torch.FloatTensor | None = None, - pixel_values_videos: torch.FloatTensor | None = None, - input_features: torch.FloatTensor | None = None, - attention_mask: torch.Tensor | None = None, - input_features_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - image_position_ids: torch.LongTensor | None = None, - video_position_ids: torch.LongTensor | None = None, - past_key_values: Cache | None = None, - mm_token_type_ids: torch.LongTensor | None = None, - inputs_embeds: torch.FloatTensor | None = None, - labels: torch.LongTensor | None = None, - use_cache: bool | None = None, - logits_to_keep: int | torch.Tensor = 0, - **kwargs: Unpack[TransformersKwargs], -) -> Gemma4CausalLMOutputWithPast: - r""" - input_features_mask (`torch.FloatTensor]` of shape `(num_images, seq_length)`): - The attention mask for the input audio. - image_position_ids (`torch.LongTensor` of shape `(batch_size, max_patches, 2)`, *optional*): - 2D patch position coordinates from the image processor, with `(-1, -1)` indicating padding. - Passed through to the vision encoder for positional embedding computation. - video_position_ids (`torch.LongTensor` of shape `(num_videos, num_frames, max_patches, 2)`, *optional*): - 2D patch position coordinates from the video processor, with `(-1, -1)` indicating padding. - Passed through to the vision encoder for positional embedding computation. - """ - if not torch.compiler.is_compiling() and torch.is_grad_enabled() and getattr(self, "training", False): - unsloth_training_stance() - outputs = self.model( - input_ids=input_ids, - pixel_values=pixel_values, - pixel_values_videos=pixel_values_videos, - input_features=input_features, - attention_mask=attention_mask, - input_features_mask=input_features_mask, - position_ids=position_ids, - past_key_values=past_key_values, - mm_token_type_ids=mm_token_type_ids, - inputs_embeds=inputs_embeds, - labels=labels, - use_cache=use_cache, - image_position_ids=image_position_ids, - video_position_ids=video_position_ids, - return_dict=True, - **kwargs, - ) - - hidden_states = outputs.last_hidden_state - # Only compute necessary logits, and do not upcast them to float if we are not computing the loss - slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep - logits = self.lm_head(hidden_states[:, slice_indices, :]) if os.environ.get('UNSLOTH_RETURN_LOGITS', '0') == '1' else EMPTY_LOGITS - loss = None - NOT_RETURN_LOGITS = os.environ.get('UNSLOTH_RETURN_LOGITS', '0') == '0' - RETURN_HIDDEN_STATES = os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES", "0") == "1" - - all_locals = locals() - n_items = None - if 'loss_kwargs' in all_locals: - __kwargs = all_locals['loss_kwargs'] - if type(__kwargs) is dict: - n_items = __kwargs.get("num_items_in_batch", None) - if n_items is None: n_items = __kwargs.get("n_items", None) - if n_items is None and 'kwargs' in all_locals: - __kwargs = all_locals['kwargs'] - if type(__kwargs) is dict: - n_items = __kwargs.get("num_items_in_batch", None) - if n_items is None: n_items = __kwargs.get("n_items", None) - if n_items is None: - all_locals = all_locals.values() - for __kwargs in all_locals: - if type(__kwargs) is dict: - n_items = __kwargs.get("num_items_in_batch", None) - if n_items is None: n_items = __kwargs.get("n_items", None) - break - pass - - if labels is not None: unsloth_fsdp_unshard(self.lm_head) - requires_grad_ = self.lm_head.weight.requires_grad - requires_grad_ = requires_grad_ or self.lm_head.weight.dtype == torch.float32 - - if RETURN_HIDDEN_STATES: - logits = hidden_states[:, slice_indices, :] - elif labels is None: - - - # Set compiler stance to fail on recompiles for inference - global INFERENCE_RUNS - # Skipped while tracing (set_stance raises there, and the counter would guard every step) - # and around a compiled decode step, which eager_on_recompile would otherwise freeze. - if not torch.compiler.is_compiling() and not UNSLOTH_DECODE_COMPILE[0]: - if torch_dynamo_eval_frame is not None: - old_stance = torch_dynamo_eval_frame._stance.stance - else: - old_stance = None - if old_stance is not None and INFERENCE_RUNS == 1: - # Skip guards and return to eager -> we still need guards! - # Already eager (another model's inference, or the user's choice): leave that stance - # object, and so its ownership, as it is. Otherwise switch and claim the new object. - if old_stance != "eager_on_recompile": - previous = getattr(torch_dynamo_eval_frame, "_stance", None) - torch_compiler_set_stance(stance = "eager_on_recompile", skip_guard_eval_unsafe = False) - # The stance object itself, or True where torch keeps no such object. - unsloth_claim_eager_stance(getattr(torch_dynamo_eval_frame, "_stance", True), previous) - if UNSLOTH_ENABLE_LOGGING: - logger_compiler.info( - f"Unsloth: Removing compiler guards after 1 inference run. " \ - f"DYNAMO_STANCE.stance = {torch_dynamo_eval_frame._stance.stance} " \ - f"DYNAMO_STANCE.skip_guard_eval_unsafe = {torch_dynamo_eval_frame._stance.skip_guard_eval_unsafe}" - ) - elif old_stance == "eager_on_recompile": - pass - elif old_stance == "default" and INFERENCE_RUNS > 1: - # Reset compiler stance - torch_compiler_set_stance(stance = "default", skip_guard_eval_unsafe = False) - if UNSLOTH_ENABLE_LOGGING: - logger_compiler.info( - f"Unsloth: Reseting guards. " \ - f"DYNAMO_STANCE.stance = {torch_dynamo_eval_frame._stance.stance} " \ - f"DYNAMO_STANCE.skip_guard_eval_unsafe = {torch_dynamo_eval_frame._stance.skip_guard_eval_unsafe}" - ) - INFERENCE_RUNS = 0 - INFERENCE_RUNS += 1 - - logits = self.lm_head(hidden_states[:, slice_indices, :]) - # Inference returns these logits to the caller, so they must carry the same - # scale and softcap transforms the training branches below apply. Leaving - # them off does not change greedy decoding, since tanh is monotonic, but it - # does change the distribution: sampling, returned logprobs and any scoring - # that reads the logits all see uncapped values. - if () != (): - logits = logits * () - if () != (): - logits = logits / () - if (self.config.get_text_config().final_logit_softcapping) not in (None, (),): - logits = logits / (self.config.get_text_config().final_logit_softcapping) - logits = torch.tanh(logits) - logits = logits * (self.config.get_text_config().final_logit_softcapping) - else: - lm_head_weight = self.lm_head.weight - lm_head_bias = getattr(self.lm_head, "bias", None) - - # ========= NEW fused ========= - _hidden_states = hidden_states[:, slice_indices, :] - torch._dynamo.mark_dynamic(_hidden_states, 1) - torch._dynamo.mark_dynamic(labels, 1) - if attention_mask is not None: - torch._dynamo.mark_dynamic(attention_mask, 1) - # Prefix tuning makes the mask longer than the labels; the stock forward keeps its tail. - _mask = attention_mask - if _mask is not None and _mask.shape[-1] > labels.shape[-1]: - _mask = _mask[..., -labels.shape[-1]:] - loss = unsloth_fused_ce_loss( - trainer = None, - hidden_states = _hidden_states, - lm_head_weight = lm_head_weight, - lm_head_bias = lm_head_bias, - labels = labels, - mask = _mask, - n_items = n_items, - scaling = getattr(self, "accelerator_scaler", None), - target_gb = None, - torch_compile = not UNSLOTH_COMPILE_DISABLE and not lm_head_on_old_cuda_arch(lm_head_weight), - logit_scale_multiply = () if () != () else 0, - logit_scale_divide = () if () != () else 0, - logit_softcapping = (self.config.get_text_config().final_logit_softcapping) if (self.config.get_text_config().final_logit_softcapping) != () else 0, - ) - - - return Gemma4CausalLMOutputWithPast( - loss=loss, - logits=logits, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - image_hidden_states=outputs.image_hidden_states, - audio_hidden_states=outputs.audio_hidden_states, - ) - -class Gemma4ForConditionalGeneration(Gemma4PreTrainedModel, GenerationMixin): - _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} - base_model_prefix = "model" - - def __init__(self, config: Gemma4Config): - super().__init__(config) - self.model = Gemma4Model(config) - self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False) - self.post_init() - - def get_input_embeddings(self): - return self.model.get_input_embeddings() - - def set_input_embeddings(self, value): - self.model.set_input_embeddings(value) - - def get_image_features( - self, - pixel_values: torch.FloatTensor, - image_position_ids: torch.LongTensor | None = None, - **kwargs: Unpack[TransformersKwargs], - ): - r""" - image_position_ids (`torch.LongTensor` of shape `(batch_size, max_patches, 2)`, *optional*): - 2D patch position coordinates from the image processor, with `(-1, -1)` indicating padding. - Passed through to the vision encoder for positional embedding computation. - """ - return self.model.get_image_features(pixel_values, image_position_ids, **kwargs) - - - def forward( - self, - input_ids: torch.LongTensor | None = None, - pixel_values: torch.FloatTensor | None = None, - pixel_values_videos: torch.FloatTensor | None = None, - input_features: torch.FloatTensor | None = None, - attention_mask: torch.Tensor | None = None, - input_features_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - image_position_ids: torch.LongTensor | None = None, - video_position_ids: torch.LongTensor | None = None, - past_key_values: Cache | None = None, - mm_token_type_ids: torch.LongTensor | None = None, - inputs_embeds: torch.FloatTensor | None = None, - labels: torch.LongTensor | None = None, - use_cache: bool | None = None, - logits_to_keep: int | torch.Tensor = 0, - **kwargs: Unpack[TransformersKwargs], - ) -> Gemma4CausalLMOutputWithPast: - return Gemma4ForConditionalGeneration_forward(self, input_ids=input_ids, pixel_values=pixel_values, pixel_values_videos=pixel_values_videos, input_features=input_features, attention_mask=attention_mask, input_features_mask=input_features_mask, position_ids=position_ids, image_position_ids=image_position_ids, video_position_ids=video_position_ids, past_key_values=past_key_values, mm_token_type_ids=mm_token_type_ids, inputs_embeds=inputs_embeds, labels=labels, use_cache=use_cache, logits_to_keep=logits_to_keep, **kwargs) - - def prepare_inputs_for_generation( - self, - input_ids, - past_key_values=None, - inputs_embeds=None, - position_ids=None, - pixel_values=None, - pixel_values_videos=None, - input_features=None, - attention_mask=None, - input_features_mask=None, - token_type_ids=None, - use_cache=True, - logits_to_keep=None, - labels=None, - is_first_iteration=False, - **kwargs, - ): - # Overwritten -- custom `position_ids` and `pixel_values` handling - model_inputs = super().prepare_inputs_for_generation( - input_ids, - past_key_values=past_key_values, - inputs_embeds=inputs_embeds, - attention_mask=attention_mask, - position_ids=position_ids, - use_cache=use_cache, - logits_to_keep=logits_to_keep, - token_type_ids=token_type_ids, - is_first_iteration=is_first_iteration, - **kwargs, - ) - - # If we're in cached decoding stage, multimodal inputs are already cached and can be dropped - if is_first_iteration or not use_cache: - model_inputs["pixel_values"] = pixel_values - model_inputs["pixel_values_videos"] = pixel_values_videos - model_inputs["input_features"] = input_features - model_inputs["input_features_mask"] = input_features_mask - - return model_inputs - - @staticmethod - def create_masks_for_generate( - config: PreTrainedConfig, - inputs_embeds: torch.Tensor, - attention_mask: torch.Tensor | None, - past_key_values: Cache | None, - position_ids: torch.Tensor | None, - mm_token_type_ids: torch.Tensor | None = None, - is_first_iteration: bool | None = False, - **kwargs, - ) -> dict: - if getattr(config.get_text_config(), "use_bidirectional_attention", None) == "vision": - # Larger Gemma 4 models use Gemma 3's bidirectional attention mask for vision inputs - return create_causal_mask_mapping( - config, - inputs_embeds, - attention_mask, - past_key_values, - position_ids, - mm_token_type_ids, - is_first_iteration=is_first_iteration, - **{k: v for k, v in kwargs.items() if k != "pixel_values"}, - ) - else: - # Smaller Gemma models use a conventional casual attention mask - return create_masks_for_generate( - config, inputs_embeds, attention_mask, past_key_values, position_ids, **kwargs - ) - -Gemma4ForConditionalGeneration.__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__ = True diff --git a/unsloth_compiled_cache/unsloth_compiled_module_siglip.py b/unsloth_compiled_cache/unsloth_compiled_module_siglip.py deleted file mode 100644 index 141cbebc9dd4aa66175b9e27e4d0690967d8b187..0000000000000000000000000000000000000000 --- a/unsloth_compiled_cache/unsloth_compiled_module_siglip.py +++ /dev/null @@ -1,506 +0,0 @@ -""" -2026.10.9 -2026.10.9 -5.5.0 -1.15.0 -__UNSLOTH_VERSIONING__ -""" - -# Unsloth auto generated code -# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . - - -import os -import sys -import torch -import importlib.util -import math -if importlib.util.find_spec("unsloth_studio") is None: - UNSLOTH_STUDIO_ENABLED = False -else: - UNSLOTH_STUDIO_ENABLED = os.environ.get("UNSLOTH_STUDIO_DISABLED", "0") == "0" -pass -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -import math - -UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1" -UNSLOTH_ENABLE_CCE = os.environ.get("UNSLOTH_ENABLE_CCE", "1") == "1" -UNSLOTH_COMPILE_DISABLE = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") in ("1", "partial",) -UNSLOTH_COMPILE_LOCATION = os.environ.get("UNSLOTH_COMPILE_LOCATION", "unsloth_compiled_cache") -if UNSLOTH_COMPILE_LOCATION not in sys.path: - sys.path.insert(0, UNSLOTH_COMPILE_LOCATION) - -import logging -logger_compiler = logging.getLogger(__name__) -if UNSLOTH_ENABLE_LOGGING: - logger_compiler.setLevel(logging.DEBUG) - -global INFERENCE_RUNS -INFERENCE_RUNS = 0 - -try: - import torch._dynamo.eval_frame as torch_dynamo_eval_frame - torch_dynamo_eval_frame._stance.stance - torch_compiler_set_stance = torch.compiler.set_stance -except: - torch_dynamo_eval_frame = None - torch_compiler_set_stance = None -pass - -# The stance is process wide, so whether Unsloth owns it lives in unsloth_zoo, shared by -# every generated module: inference in one model's module, training in another's. -from unsloth_zoo.temporary_patches.utils import UNSLOTH_EAGER_STANCE_OWNED, unsloth_claim_eager_stance, unsloth_owns_stance, unsloth_release_stance - -def unsloth_training_stance(): - # eager_on_recompile is global: left on by generate() or eval forwards, every recompile - # in a later training step (new sequence length, model.train() without for_training) - # would silently run the eager body. Training forwards restore the default stance, - # unless the user chose eager_on_recompile themselves. Every set_stance call installs a - # fresh stance object, so ours is still active only if the current one `is` it; a user - # set_stance since then (even to eager_on_recompile again) is theirs and stays. - global INFERENCE_RUNS - INFERENCE_RUNS = 0 - if len(UNSLOTH_EAGER_STANCE_OWNED) != 0 and torch_dynamo_eval_frame is not None: - current = torch_dynamo_eval_frame._stance - # Ownership is only dropped once ours is reset: a temporary stance scope's exit restores - # one of our exact objects, which a later training forward must still recognise. - if current.stance == "eager_on_recompile" and unsloth_owns_stance(current): - # Only the reset object: an outer scope may still restore another one of ours. - # Put back what the caller had before Unsloth switched (force_eager, a backend, ...). - previous = unsloth_release_stance(current) - if previous is not None and hasattr(torch_dynamo_eval_frame, "_set_stance"): - torch_dynamo_eval_frame._set_stance(previous) - else: - torch_compiler_set_stance(stance = "default", skip_guard_eval_unsafe = False) -pass - -from unsloth_zoo import DEVICE_TYPE_TORCH, DEVICE_COUNT -from unsloth_zoo.temporary_patches.utils import UNSLOTH_DECODE_COMPILE - - -from unsloth_zoo.loss_utils import ( - HAS_CUT_CROSS_ENTROPY, - fused_linear_cross_entropy, - unsloth_fused_ce_loss, -) -from unsloth_zoo.fused_losses.forward_adapter import unsloth_fused_lm_head_loss -from unsloth_zoo.fused_losses.forward_install import _can_fuse_loss -from unsloth_zoo.compiler import lm_head_on_old_cuda_arch -from unsloth_zoo.fused_losses.cross_entropy_loss import unsloth_count_aware_cross_entropy, unsloth_loss_count_kwargs, unsloth_fsdp_unshard -from unsloth_zoo.fused_losses.aux_loss import unsloth_ga_scale_aux_loss - -scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention -@torch.compiler.disable(recursive = False) -def disable_compile_scaled_dot_product_attention(*args, **kwargs): - return scaled_dot_product_attention(*args, **kwargs) -pass - - -from transformers.modeling_flash_attention_utils import is_flash_attn_available - -if is_flash_attn_available(): - try: - from transformers.modeling_flash_attention_utils import flash_attn_supports_top_left_mask - except: - flash_attn_supports_top_left_mask = None - try: - from transformers.modeling_flash_attention_utils import _flash_attention_forward - except: - _flash_attention_forward = None - try: - from transformers.modeling_flash_attention_utils import FlashAttentionKwargs - except: - FlashAttentionKwargs = None - try: - from transformers.modeling_flash_attention_utils import flash_attn_varlen_func - except: - flash_attn_varlen_func = None -else: - flash_attn_supports_top_left_mask = None - _flash_attention_forward = None - FlashAttentionKwargs = None - flash_attn_varlen_func = None -pass - - -torch_compile_options = {'epilogue_fusion': True, 'max_autotune': False, 'shape_padding': True, 'trace.enabled': False, 'triton.cudagraphs': False, 'debug': False, 'dce': True, 'memory_planning': True, 'coordinate_descent_tuning': False, 'trace.graph_diagram': False, 'compile_threads': 32, 'group_fusion': True, 'disable_progress': True, 'verbose_progress': False, 'triton.multi_kernel': 0, 'triton.use_block_ptr': False, 'triton.enable_persistent_tma_matmul': False, 'triton.autotune_at_compile_time': False, 'triton.cooperative_reductions': False, 'cuda.compile_opt_level': '-O2', 'cuda.enable_cuda_lto': True, 'combo_kernels': False, 'benchmark_combo_kernel': True, 'combo_kernel_foreach_dynamic_shapes': True} - -from torch.nn import CrossEntropyLoss -from unsloth_zoo.temporary_patches.utils import torch_compile_with_fallback - -@torch_compile_with_fallback(fullgraph = True, dynamic = True, options = torch_compile_options) -def normal_cross_entropy_loss(self, hidden_states, labels): - logits = self.lm_head(hidden_states) - logits = logits.float() - # Shift so that tokens < n predict n - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - # Flatten the tokens - loss_fct = CrossEntropyLoss() - shift_logits = shift_logits.view(-1, self.config.vocab_size) - shift_labels = shift_labels.view(-1) - # Enable model parallelism - shift_labels = shift_labels.to(shift_logits.device) - loss = loss_fct(shift_logits, shift_labels) - return loss, logits -pass - -# We need an empty logits flag to warn people logits will not be returned anymore unless asked ie -# os.environ['UNSLOTH_RETURN_LOGITS'] = '1' -LOGITS_ERROR_STRING = \ - "Unsloth: Logits are empty from 2024.11 onwards. To get raw logits again, please "\ - 'set the environment variable `UNSLOTH_RETURN_LOGITS` to `"1" BEFORE starting to train ie before `trainer.train()`. For example:\n'\ - "```\nimport os\n"\ - "os.environ['UNSLOTH_RETURN_LOGITS'] = '1'\n"\ - "trainer.train()\n```\n"\ - "No need to restart your console - just add `os.environ['UNSLOTH_RETURN_LOGITS'] = '1'` before trainer.train() and re-run the cell!" - -def raise_logits_error(*args, **kwargs): raise NotImplementedError(LOGITS_ERROR_STRING) -def return_none(*args, **kwargs): return None -class EmptyLogits: - def __init__(self): return - def raise_getattr_error(self, attr): - if attr == "to": return return_none - # A catch-all __getattr__ makes hasattr() true for every name, dunders included, so the - # sentinel answers yes to protocol probes it cannot honour. torch.distributed's output cast - # tests `hasattr(x, "__dataclass_fields__")` and then calls `dataclasses.replace(x)` on - # whatever said yes, so with FSDP2 mixed precision every step died in - # `TypeError: replace() should be called on dataclass instances` (unsloth#409, reached - # through `_fsdp_state._cast_output_dtype`). Protocol probes get an honest AttributeError; - # ordinary attribute access still gets the callable that explains UNSLOTH_RETURN_LOGITS. - if len(attr) > 4 and attr.startswith("__") and attr.endswith("__"): - raise AttributeError(f"{type(self).__name__!r} object has no attribute {attr!r}") - return raise_logits_error - __getitem__ = raise_logits_error - __getattr__ = raise_getattr_error - def __repr__(self): return LOGITS_ERROR_STRING - def __str__ (self): return LOGITS_ERROR_STRING - # Stateless pickling so accelerate gather_object works on the sentinel - def __reduce__(self): return (type(self), ()) - # Gathered copies must compare equal in accelerate debug mode - def __eq__(self, other): return type(other).__name__ == "EmptyLogits" - __hash__ = object.__hash__ -pass -EMPTY_LOGITS = EmptyLogits() -functions = dir(torch.Tensor) -for j, function in enumerate(functions): - if function.startswith("__") and function.endswith("__"): - exec(f"def raise_{j}(*args, **kwargs): print('{function}')", globals(), locals()) - try: exec(f"EMPTY_LOGITS.{function} = raise_{j}", globals(), locals()) - except: continue -pass -# The loop above stomps pickle hooks with stubs returning None; restore them. -for function in ("__reduce__", "__reduce_ex__", "__getstate__", "__setstate__"): - try: delattr(EMPTY_LOGITS, function) - except Exception: pass -pass - - -def mask_attention_mask_out(labels = None, attention_mask = None): - # Only a mask over the labels' own positions: MusicGen's attention_mask is the text encoder's, - # while its labels are decoder codebook frames. - if ( - labels is not None and attention_mask is not None - and tuple(attention_mask.shape) == tuple(labels.shape) - ): - attention_mask = attention_mask.to(device = labels.device) - labels[attention_mask == 0] = -100 - return labels -pass - - -from torch import Tensor -import torch -import torch.nn as nn -from torch.nn import functional as F -from unsloth_zoo.temporary_patches.utils import torch_compile_with_fallback -from unsloth_zoo.temporary_patches.utils import torch_compiler_disable_unless_decode -from unsloth_zoo.temporary_patches.common import torch_compile -from typing import Any, List, Optional, Tuple, Union, Dict, Set, Callable -from transformers.models.siglip.modeling_siglip import (Callable, np, torch, nn, init, ACT2FN, ALL_ATTENTION_FUNCTIONS, torch_int, SiglipTextConfig, SiglipVisionConfig) - -@torch_compile_with_fallback(fullgraph = False, dynamic = True, options = torch_compile_options) -def SiglipVisionEmbeddings_forward(self, pixel_values: torch.FloatTensor, interpolate_pos_encoding=False) -> torch.Tensor: - _, _, height, width = pixel_values.shape - target_dtype = self.patch_embedding.weight.dtype - patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype)) # shape = [*, width, grid, grid] - embeddings = patch_embeds.flatten(2).transpose(1, 2) - - if interpolate_pos_encoding: - embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width) - else: - embeddings = embeddings + self.position_embedding(self.position_ids) - return embeddings - -class SiglipVisionEmbeddings(nn.Module): - def __init__(self, config: SiglipVisionConfig): - super().__init__() - self.config = config - self.embed_dim = config.hidden_size - self.image_size = config.image_size - self.patch_size = config.patch_size - - self.patch_embedding = nn.Conv2d( - in_channels=config.num_channels, - out_channels=self.embed_dim, - kernel_size=self.patch_size, - stride=self.patch_size, - padding="valid", - ) - - self.num_patches = (self.image_size // self.patch_size) ** 2 - self.num_positions = self.num_patches - self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) - self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False) - - def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor: - """ - This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution - images. This method is also adapted to support torch.jit tracing and no class embeddings. - - Adapted from: - - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and - - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211 - """ - - num_patches = embeddings.shape[1] - num_positions = self.position_embedding.weight.shape[0] - - # always interpolate when tracing to ensure the exported model works for dynamic input shapes - if not torch.jit.is_tracing() and num_patches == num_positions and height == width: - return self.position_embedding(self.position_ids) - - patch_pos_embed = self.position_embedding.weight.unsqueeze(0) - - dim = embeddings.shape[-1] - - new_height = height // self.patch_size - new_width = width // self.patch_size - - sqrt_num_positions = torch_int(num_positions**0.5) - patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim) - patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2) - - patch_pos_embed = nn.functional.interpolate( - patch_pos_embed, - size=(new_height, new_width), - mode="bicubic", - align_corners=False, - ) - - patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) - return patch_pos_embed - - def forward(self, pixel_values: torch.FloatTensor, interpolate_pos_encoding=False) -> torch.Tensor: - return SiglipVisionEmbeddings_forward(self, pixel_values=pixel_values, interpolate_pos_encoding=interpolate_pos_encoding) - - -@torch_compile_with_fallback(fullgraph = True, dynamic = True, options = torch_compile_options) -def SiglipTextEmbeddings_forward( - self, - input_ids: torch.LongTensor | None = None, - position_ids: torch.LongTensor | None = None, - inputs_embeds: torch.FloatTensor | None = None, -) -> torch.Tensor: - seq_length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2] - max_position_embedding = self.position_embedding.weight.shape[0] - - if seq_length > max_position_embedding: - raise ValueError( - f"Sequence length must be less than max_position_embeddings (got `sequence length`: " - f"{seq_length} and max_position_embeddings: {max_position_embedding}" - ) - - if position_ids is None: - position_ids = self.position_ids[:, :seq_length] - - if inputs_embeds is None: - inputs_embeds = self.token_embedding(input_ids) - - position_embeddings = self.position_embedding(position_ids) - embeddings = inputs_embeds + position_embeddings - - return embeddings - -class SiglipTextEmbeddings(nn.Module): - def __init__(self, config: SiglipTextConfig): - super().__init__() - embed_dim = config.hidden_size - - self.token_embedding = nn.Embedding(config.vocab_size, embed_dim) - self.position_embedding = nn.Embedding(config.max_position_embeddings, embed_dim) - - # position_ids (1, len position emb) is contiguous in memory and exported when serialized - self.register_buffer( - "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False - ) - - def forward( - self, - input_ids: torch.LongTensor | None = None, - position_ids: torch.LongTensor | None = None, - inputs_embeds: torch.FloatTensor | None = None, - ) -> torch.Tensor: - return SiglipTextEmbeddings_forward(self, input_ids=input_ids, position_ids=position_ids, inputs_embeds=inputs_embeds) - - -@torch_compile_with_fallback(fullgraph = True, dynamic = True, options = torch_compile_options) -def eager_attention_forward( - module: nn.Module, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attention_mask: torch.Tensor | None, - scaling: float, - dropout: float = 0.0, - **kwargs, -): - attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling - if attention_mask is not None: - - if isinstance(attention_mask, dict): - - attention_mask = attention_mask.get(getattr(module, 'layer_type', None), None) - - if attention_mask is not None: - - attn_weights = attn_weights + attention_mask - - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype = torch.float32).to(attn_weights.dtype).to(query.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) - - attn_output = torch.matmul(attn_weights, value) - attn_output = attn_output.transpose(1, 2).contiguous() - - return attn_output, attn_weights - - -@torch_compiler_disable_unless_decode -def SiglipAttention_forward( - self, - hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None = None, - **kwargs, -) -> tuple[torch.Tensor, torch.Tensor | None]: - """Input shape: Batch x Time x Channel""" - - batch_size, seq_length, embed_dim = hidden_states.shape - - queries = self.q_proj(hidden_states) - keys = self.k_proj(hidden_states) - values = self.v_proj(hidden_states) - - queries = queries.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) - keys = keys.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) - values = values.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) - - attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( - self.config._attn_implementation, eager_attention_forward - ) - - attn_output, attn_weights = attention_interface( - self, - queries, - keys, - values, - attention_mask, - is_causal=self.is_causal, - scaling=self.scale, - dropout=0.0 if not self.training else self.dropout, - ) - - attn_output = attn_output.reshape(batch_size, seq_length, embed_dim).contiguous() - attn_output = self.out_proj(attn_output) - - return attn_output, attn_weights - -class SiglipAttention(nn.Module): - """Multi-headed attention from 'Attention Is All You Need' paper""" - - def __init__(self, config): - super().__init__() - self.config = config - self.embed_dim = config.hidden_size - self.num_heads = config.num_attention_heads - self.head_dim = self.embed_dim // self.num_heads - if self.head_dim * self.num_heads != self.embed_dim: - raise ValueError( - f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:" - f" {self.num_heads})." - ) - self.scale = self.head_dim**-0.5 - self.dropout = config.attention_dropout - self.is_causal = False - - self.k_proj = nn.Linear(self.embed_dim, self.embed_dim) - self.v_proj = nn.Linear(self.embed_dim, self.embed_dim) - self.q_proj = nn.Linear(self.embed_dim, self.embed_dim) - self.out_proj = nn.Linear(self.embed_dim, self.embed_dim) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None = None, - **kwargs, - ) -> tuple[torch.Tensor, torch.Tensor | None]: - return SiglipAttention_forward(self, hidden_states=hidden_states, attention_mask=attention_mask, **kwargs) - - -@torch_compile_with_fallback(fullgraph = False, dynamic = True, options = torch_compile_options) -def SiglipMLP_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = self.fc1(hidden_states) - hidden_states = self.activation_fn(hidden_states) - hidden_states = self.fc2(hidden_states) - return hidden_states - -class SiglipMLP(nn.Module): - def __init__(self, config): - super().__init__() - self.config = config - self.activation_fn = ACT2FN[config.hidden_act] - self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) - self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - return SiglipMLP_forward(self, hidden_states=hidden_states) - - -@torch_compile_with_fallback(fullgraph = False, dynamic = True, options = torch_compile_options) -def SiglipMultiheadAttentionPoolingHead_forward(self, hidden_state): - batch_size = hidden_state.shape[0] - probe = self.probe.repeat(batch_size, 1, 1) - - hidden_state = self.attention(probe, hidden_state, hidden_state)[0] - - residual = hidden_state - hidden_state = self.layernorm(hidden_state) - hidden_state = residual + self.mlp(hidden_state) - - return hidden_state[:, 0] - -class SiglipMultiheadAttentionPoolingHead(nn.Module): - """Multihead Attention Pooling.""" - - def __init__(self, config: SiglipVisionConfig): - super().__init__() - - self.probe = nn.Parameter(torch.randn(1, 1, config.hidden_size)) - self.attention = torch.nn.MultiheadAttention(config.hidden_size, config.num_attention_heads, batch_first=True) - self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) - self.mlp = SiglipMLP(config) - - def forward(self, hidden_state): - return SiglipMultiheadAttentionPoolingHead_forward(self, hidden_state=hidden_state)