diff --git a/pyproject.toml b/pyproject.toml index 2a6e31dcae..59a7c44737 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,7 @@ triton = [ "triton @ https://github.com/woct0rdho/triton-windows/releases/download/v3.1.0-windows.post5/triton-3.1.0-cp312-cp312-win_amd64.whl ; python_version=='3.12' and platform_system == 'Windows'", ] huggingface = [ - "unsloth_zoo>=2025.2.2", + "unsloth_zoo>=2025.2.5", "packaging", "tyro", "transformers>=4.46.1,!=4.47.0", @@ -50,7 +50,7 @@ huggingface = [ "wheel>=0.42.0", "numpy", "accelerate>=0.34.1", - "trl>=0.7.9,!=0.9.0,!=0.9.1,!=0.9.2,!=0.9.3,<0.15.0", + "trl>=0.7.9,!=0.9.0,!=0.9.1,!=0.9.2,!=0.9.3", "peft>=0.7.1,!=0.11.0", "protobuf<4.0.0", "huggingface_hub", @@ -176,26 +176,26 @@ cu124onlytorch251 = [ "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post1-cp312-cp312-win_amd64.whl ; python_version=='3.12' and platform_system == 'Windows'", ] cu118onlytorch260 = [ - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post2-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and platform_system == 'Linux'", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post2-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and platform_system == 'Linux'", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post2-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and platform_system == 'Linux'", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post2-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and platform_system == 'Linux'", + "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post3-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and platform_system == 'Linux'", + "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post3-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and platform_system == 'Linux'", + "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post3-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and platform_system == 'Linux'", + "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.29.post3-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and platform_system == 'Linux'", ] cu124onlytorch260 = [ - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post2-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and platform_system == 'Linux'", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post2-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and platform_system == 'Linux'", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post2-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and platform_system == 'Linux'", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post2-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and platform_system == 'Linux'", - "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post2-cp39-cp39-win_amd64.whl ; python_version=='3.9' and platform_system == 'Windows'", + "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post3-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and platform_system == 'Linux'", + "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post3-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and platform_system == 'Linux'", + "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post3-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and platform_system == 'Linux'", + "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post3-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and platform_system == 'Linux'", + "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post3-cp39-cp39-win_amd64.whl ; python_version=='3.9' and platform_system == 'Windows'", "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post3-cp310-cp310-win_amd64.whl ; python_version=='3.10' and platform_system == 'Windows'", "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post3-cp311-cp311-win_amd64.whl ; python_version=='3.11' and platform_system == 'Windows'", "xformers @ https://download.pytorch.org/whl/cu124/xformers-0.0.29.post3-cp312-cp312-win_amd64.whl ; python_version=='3.12' and platform_system == 'Windows'", ] cu126onlytorch260 = [ - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post2-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and platform_system == 'Linux'", - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post2-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and platform_system == 'Linux'", - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post2-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and platform_system == 'Linux'", - "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post2-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and platform_system == 'Linux'", + "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post3-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and platform_system == 'Linux'", + "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post3-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and platform_system == 'Linux'", + "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post3-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and platform_system == 'Linux'", + "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.29.post3-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and platform_system == 'Linux'", ] cu118 = [ "unsloth[huggingface]", @@ -344,7 +344,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3", ] colab-new = [ - "unsloth_zoo>=2025.2.2", + "unsloth_zoo>=2025.2.5", "packaging", "tyro", "transformers>=4.46.1,!=4.47.0", @@ -362,7 +362,7 @@ colab-new = [ ] colab-no-deps = [ "accelerate>=0.34.1", - "trl>=0.7.9,!=0.9.0,!=0.9.1,!=0.9.2,!=0.9.3,<0.15.0", + "trl>=0.7.9,!=0.9.0,!=0.9.1,!=0.9.2,!=0.9.3", "peft>=0.7.1", "xformers", "bitsandbytes>=0.46.1", diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 8d0eadb968..df925d746b 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "2025.2.9" +__version__ = "2025.2.10" __all__ = [ "SUPPORTS_BFLOAT16", diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 3d601b0af1..2875ff64a5 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -21,12 +21,25 @@ from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union import inspect import os import re +import torch from unsloth_zoo.compiler import create_new_function from unsloth_zoo.logging_utils import PatchRLStatistics +from unsloth_zoo.rl_replacements import RL_REPLACEMENTS from .rl_replacements import ( RL_EXTRA_ARGS, RL_FUNCTIONS, + RL_PRE_ITEMS, + RL_CONFIG_CHANGES, ) +selective_log_softmax = RL_REPLACEMENTS["selective_log_softmax"] + +torch_compile_options = { + "epilogue_fusion" : True, + "max_autotune" : True, + "shape_padding" : True, + "trace.enabled" : False, + "triton.cudagraphs" : False, +} def PatchRL(FastLanguageModel): @@ -81,6 +94,17 @@ from dataclasses import dataclass, field from packaging.version import Version import torch from contextlib import nullcontext +from torch.nn import functional as F +torch_compile_options = {{ + "epilogue_fusion" : True, + "max_autotune" : True, + "shape_padding" : True, + "trace.enabled" : False, + "triton.cudagraphs" : False, +}} + +{selective_log_softmax_code} +{RL_pre} @dataclass class Unsloth{RLConfig_name}({RLConfig_name}): @@ -142,8 +166,13 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): if RLTrainer.__name__.startswith("Unsloth"): return if RLConfig .__name__.startswith("Unsloth"): return + # Get old source + old_RLTrainer_source = inspect.getsource(RLTrainer) + old_RLConfig_source = inspect.getsource(RLConfig) + all_imports = dir(trainer) - imports = [x for x in all_imports if not x.startswith("_")] + # Fix _deprecate_arguments not getting imported so stop __ but not _ + imports = [x for x in all_imports if not x.startswith("__")] # Get default arguments EMPTY = inspect.Parameter.empty @@ -358,6 +387,13 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): extra_args += num_proc_check pass + # Edit config with anything extra + if trainer_file in RL_CONFIG_CHANGES: + process_extra_args = RL_CONFIG_CHANGES[trainer_file] + for process_extra_arg in process_extra_args: + extra_args += process_extra_arg(old_RLTrainer_source, old_RLConfig_source) + pass + # Edit report_to and default it to nothing if max_steps is like 60 # Create RLConfig args @@ -377,6 +413,17 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): __RLTrainer_doc__ = eval(f"trl.trainer.{RLTrainer_name}").__doc__ __RLConfig_doc__ = eval(f"trl.trainer.{RLConfig_name}") .__doc__ + # Get all pre-modules + if trainer_file in RL_PRE_ITEMS: + RL_pre = "\n".join(RL_PRE_ITEMS[trainer_file]) + else: + RL_pre = "" + pass + + # Selective log softmax + selective_log_softmax_code = inspect.getsource(selective_log_softmax) + + # Get final source code RLTrainer_source = RLTrainer_replacement.format( RLTrainer_name = RLTrainer_name, __RLTrainer_doc__ = __RLTrainer_doc__, @@ -394,6 +441,9 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): RLTrainer_extras = RLTrainer_extras, RLTrainer_post = RLTrainer_post, + RL_pre = RL_pre, + + selective_log_softmax_code = selective_log_softmax_code, ) # Create new function @@ -402,7 +452,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): RLTrainer_source, f"trl.trainer.{trainer_file}", imports, - overwrite = False, + overwrite = True, ) # Patch Trainer diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 82fd3f8d3c..63fe243595 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -15,14 +15,27 @@ __all__ = [ "RL_EXTRA_ARGS", "RL_FUNCTIONS", + "RL_PRE_ITEMS", + "RL_CONFIG_CHANGES", ] import re +import torch import inspect from collections import defaultdict -RL_EXTRA_ARGS = defaultdict(list) -RL_FUNCTIONS = defaultdict(list) +from unsloth_zoo.rl_replacements import RL_REPLACEMENTS +RL_EXTRA_ARGS = defaultdict(list) +RL_FUNCTIONS = defaultdict(list) +RL_PRE_ITEMS = defaultdict(list) +RL_CONFIG_CHANGES = defaultdict(list) +torch_compile_options = { + "epilogue_fusion" : True, + "max_autotune" : True, + "shape_padding" : True, + "trace.enabled" : False, + "triton.cudagraphs" : False, +} # Check untrained tokens def sft_trainer_fix_untraiend_tokens(call_args, extra_args): @@ -161,23 +174,89 @@ RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__move_model_to_vllm) def grpo_trainer__get_per_token_logps(function_name, function): if function_name != "_get_per_token_logps": return function - # Edit model to autocast it - # .*? matches first match. .+? matches final match. - original = re.findall( - r"\n([ ]{4,})(logits = model\(.*?\))", - function, - flags = re.MULTILINE | re.DOTALL, - ) - if len(original) != 0: - spaces, original = original[0] - spaces = len(spaces) - replacer = \ - "if not hasattr(self, '_autocast_dtype'):\n" + \ - " "*(spaces + 4) + "self._autocast_dtype = torch.float16 if os.environ.get('ACCELERATE_MIXED_PRECISION', 'fp16') == 'fp16' else torch.bfloat16\n" + \ - " "*(spaces + 0) + "with torch.amp.autocast(device_type = 'cuda', dtype = self._autocast_dtype):\n" + \ - " "*(spaces + 4) + original - function = function.replace(original, replacer) + def _get_per_token_logps(self, model, input_ids, attention_mask, logits_to_keep): + if not hasattr(self, '_autocast_dtype'): + self._autocast_dtype = torch.float16 if os.environ.get('ACCELERATE_MIXED_PRECISION', 'fp16') == 'fp16' else torch.bfloat16 + with torch.amp.autocast(device_type = 'cuda', dtype = self._autocast_dtype): + # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded + logits = model(input_ids=input_ids, attention_mask=attention_mask, logits_to_keep=logits_to_keep + 1).logits + logits = logits[:, :-1, :] # (B, L-1, V), exclude the last logit: it corresponds to the next token pred + + input_ids = input_ids[:, -logits_to_keep:] + # For transformers<=4.48, logits_to_keep argument isn't supported, so here we drop logits ourselves. + # See https://github.com/huggingface/trl/issues/2770 + logits = logits[:, -logits_to_keep:] + return logits + # return selective_log_softmax(logits, input_ids) # compute logprobs for the input tokens + pass pass + + function = inspect.getsource(_get_per_token_logps) return function pass RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__get_per_token_logps) + +grpo_compute_loss = RL_REPLACEMENTS["grpo_compute_loss"] +RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(grpo_compute_loss)) + +# Edit _get_per_token_logps to handle mixed precision +def grpo_trainer_compute_loss(function_name, function): + if function_name != "compute_loss": return function + + def compute_loss(self, model, inputs, return_outputs = False, num_items_in_batch = None): + if return_outputs: + raise ValueError("The GRPOTrainer does not support returning outputs") + # Compute the per-token log probabilities for the model + + prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"] + completion_ids, completion_mask = inputs["completion_ids"], inputs["completion_mask"] + input_ids = torch.cat([prompt_ids, completion_ids], dim=1) + # attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) + attention_mask = None + logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens + + per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep) + + # Compute the KL divergence between the model and the reference model + 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 + + # x - x.detach() allows for preserving gradients from x + advantages = inputs["advantages"] + # per_token_loss = torch.exp(per_token_logps - per_token_logps.detach()) * advantages.unsqueeze(1) + # per_token_loss = -(per_token_loss - self.beta * per_token_kl) + # loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean() + input_ids = input_ids[:, -logits_to_keep:] + loss, completion_length, mean_kl = grpo_compute_loss( + ref_per_token_logps, per_token_logps, input_ids, completion_mask, self.beta, advantages, + ) + # Log the metrics + # completion_length = self.accelerator.gather_for_metrics(completion_mask.sum(1)).float().mean().item() + self._metrics["completion_length"].append(completion_length.item()) + + # mean_kl = ((per_token_kl * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean() + # self._metrics["kl"].append(self.accelerator.gather_for_metrics(mean_kl).mean().item()) + self._metrics["kl"].append(mean_kl.item()) + return loss + pass + + function = inspect.getsource(compute_loss) + return function +pass +RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer_compute_loss) + +# https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_trainer.py#L356 +# TRL warns if batch size is not a multiple of num_generations -> fix this. +def grpo_trainer_fix_batch_size(RLTrainer_source, RLConfig_source): + if "divisible by the number of generations" not in RLTrainer_source: return "" + if "num_generations" not in RLConfig_source: return "" + + check_batch_size = \ + "div = per_device_train_batch_size // num_generations\n"\ + "if div * num_generations != per_device_train_batch_size:\n"\ + " print('Unsloth: We know expect `per_device_train_batch_size` to be a multiple of `num_generations`.\\n"\ + "We will change the batch size of ' + str(per_device_train_batch_size) + ' to the `num_generations` of ' + str(num_generations))\n"\ + " per_device_train_batch_size = num_generations\n" + return check_batch_size +pass +RL_CONFIG_CHANGES["grpo_trainer"].append(grpo_trainer_fix_batch_size)