diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 3d601b0af1..a216f4f387 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -26,8 +26,17 @@ from unsloth_zoo.logging_utils import PatchRLStatistics from .rl_replacements import ( RL_EXTRA_ARGS, RL_FUNCTIONS, + RL_PRE_ITEMS, ) +torch_compile_options = { + "epilogue_fusion" : True, + "max_autotune" : True, + "shape_padding" : True, + "trace.enabled" : False, + "triton.cudagraphs" : False, +} + def PatchRL(FastLanguageModel): from trl.models.utils import unwrap_model_for_generation @@ -74,6 +83,23 @@ def PatchRL(FastLanguageModel): pass +# https://github.com/huggingface/trl/blob/main/trl/trainer/utils.py#L1674 +@torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options,) +def _selective_log_softmax(logits, index): + logits = logits.to(torch.float32) + selected_logits = torch.gather(logits, dim=-1, index=index.unsqueeze(-1)).squeeze(-1) + # loop to reduce peak mem consumption + # logsumexp_values = torch.stack([torch.logsumexp(lg, dim=-1) for lg in logits]) + logsumexp_values = torch.logsumexp(logits, dim = -1) + per_token_logps = selected_logits - logsumexp_values # log_softmax(x_i) = x_i - logsumexp(x) + return per_token_logps +pass + +def selective_log_softmax(logits, index): + return _selective_log_softmax(logits, index) +pass + + RLTrainer_replacement = ''' import os from typing import * @@ -81,6 +107,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}): @@ -377,6 +414,19 @@ 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 RLTrainer_name in RL_PRE_ITEMS: + RL_pre = "\n".join(RL_PRE_ITEMS) + else: + RL_pre = "" + pass + + # Selective log softmax + selective_log_softmax_code = \ + inspect.getsource(_selective_log_softmax) + "\n" + \ + inspect.getsource(selective_log_softmax) + "\n" + + # Get final source code RLTrainer_source = RLTrainer_replacement.format( RLTrainer_name = RLTrainer_name, __RLTrainer_doc__ = __RLTrainer_doc__, @@ -394,6 +444,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 +455,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..39db053550 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -15,6 +15,7 @@ __all__ = [ "RL_EXTRA_ARGS", "RL_FUNCTIONS", + "RL_PRE_ITEMS", ] import re @@ -22,7 +23,15 @@ import inspect from collections import defaultdict RL_EXTRA_ARGS = defaultdict(list) RL_FUNCTIONS = defaultdict(list) +RL_PRE_ITEMS = 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 +170,109 @@ 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) + + +# Custom compiled GRPO loss - creates 3 Triton kernels +@torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options,) +def _grpo_compute_loss(old_logits, new_logits, input_ids, mask, beta): + old_logits = old_logits.to(torch.float32) + new_logits = new_logits.to(torch.float32) + input_ids = input_ids.unsqueeze(-1) + + # x_i - logsumexp(x_i) + old_x = torch.gather(old_logits, dim = -1, index = input_ids).squeeze(-1) + new_x = torch.gather(new_logits, dim = -1, index = input_ids).squeeze(-1) + old = old_x - torch.logsumexp(old_logits, dim = -1) + new = new_x - torch.logsumexp(new_logits, dim = -1) + + kl_i = torch.exp(old - new) - (old - new) - 1.0 + loss_i = torch.exp(new - new.detach()) * advantages.unsqueeze(1) + loss_i = -(loss_i - beta * kl_i) + + mask = mask.to(torch.float32) + n_mask = mask.sum(1) + loss_per_reward = (loss_i * mask).sum(1) / n_mask + loss = loss_per_reward.mean() + + # Get metrics as well which are folded + with torch.inference_mode(): + completion_length = n_mask.mean() + mean_kl_per_reward = (kl_i * mask).sum(1) / n_mask + mean_kl = mean_kl_per_reward.mean() + pass + return loss, completion_length, mean_kl +pass +def grpo_compute_loss(old_logits, new_logits, input_ids, mask, beta): + loss, completion_length, mean_kl = _grpo_compute_loss(old_logits, new_logits, input_ids, mask, beta) + return loss, completion_length.item(), mean_kl.item() +pass +RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource((_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() + + loss, completion_length, mean_kl = grpo_compute_loss( + ref_per_token_logps, per_token_logps, input_ids, completion_mask, self.beta, + ) + # Log the metrics + # completion_length = self.accelerator.gather_for_metrics(completion_mask.sum(1)).float().mean().item() + self._metrics["completion_length"].append(completion_length) + + # 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) + return loss + pass + + function = inspect.getsource(compute_loss) + return function +pass +RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer_compute_loss)