diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index df925d746b..2a5b71d399 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.10" +__version__ = "2025.2.11" __all__ = [ "SUPPORTS_BFLOAT16", diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 2875ff64a5..7b363d8fc1 100644 --- a/unsloth/models/rl.py +++ b/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: diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 63fe243595..1e13068213 100644 --- a/unsloth/models/rl_replacements.py +++ b/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)