GRPO optimized

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

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)