Metrics GRPO

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.10"
+__version__ = "2025.2.11"
 
 __all__ = [
     "SUPPORTS_BFLOAT16",

unsloth/models/rl.py

@@ -30,6 +30,7 @@ from .rl_replacements import (
     RL_FUNCTIONS,
     RL_PRE_ITEMS,
     RL_CONFIG_CHANGES,
+    RL_METRICS_CHANGES,
 )
 selective_log_softmax = RL_REPLACEMENTS["selective_log_softmax"]
 
@@ -310,10 +311,20 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
         RLTrainer_post += neftune_check
     pass
 
+    # Edit optional metrics
+    other_metrics_processor = ""
+    if trainer_file in RL_METRICS_CHANGES:
+        process_extra_args = RL_METRICS_CHANGES[trainer_file]
+        for process_extra_arg in process_extra_args:
+            other_metrics_processor += process_extra_arg(call_args, extra_args)
+    pass
+
     # Add statistics as well!
     extra_args += \
+        "other_metrics = []\n"\
+        f"{other_metrics_processor}\n"\
         "from unsloth_zoo.logging_utils import PatchRLStatistics\n"\
-        f"PatchRLStatistics('{trainer_file}')\n"
+        f"PatchRLStatistics('{trainer_file}', other_metrics)\n"
 
     # Patch optional args
     if trainer_file in RL_EXTRA_ARGS:

unsloth/models/rl_replacements.py

@@ -17,6 +17,7 @@ __all__ = [
     "RL_FUNCTIONS",
     "RL_PRE_ITEMS",
     "RL_CONFIG_CHANGES",
+    "RL_METRICS_CHANGES",
 ]
 
 import re
@@ -24,10 +25,11 @@ import torch
 import inspect
 from collections import defaultdict
 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)
+RL_EXTRA_ARGS      = defaultdict(list)
+RL_FUNCTIONS       = defaultdict(list)
+RL_PRE_ITEMS       = defaultdict(list)
+RL_CONFIG_CHANGES  = defaultdict(list)
+RL_METRICS_CHANGES = dict()
 
 torch_compile_options = {
     "epilogue_fusion"   : True,
@@ -260,3 +262,19 @@ def grpo_trainer_fix_batch_size(RLTrainer_source, RLConfig_source):
     return check_batch_size
 pass
 RL_CONFIG_CHANGES["grpo_trainer"].append(grpo_trainer_fix_batch_size)
+
+
+# Add other reward function names
+def grpo_trainer_metrics(RLTrainer_source, RLConfig_source):
+    if "reward_funcs" not in RLTrainer_source: return ""
+
+    log_metrics = \
+    "if not isinstance(reward_funcs, list): _reward_funcs = [reward_funcs]\n"\
+    "for reward_func in _reward_funcs:\n"\
+    "    try:\n"\
+    "        reward_func_name = reward_func.__name__\n"\
+    "        other_metrics.append(f'rewards/{reward_func_name}')\n"\
+    "    except: pass\n"
+    return log_metrics
+pass
+RL_METRICS_CHANGES["grpo_trainer"].append(grpo_trainer_metrics)