Fix crash when reward function is a functools.partial or callable instance (#6313)

Author: qgallouedec

tests/test_utils.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import copy
+import functools
 import textwrap
 from io import StringIO
 from unittest.mock import patch
@@ -37,6 +38,7 @@
     entropy_from_logits,
     flush_left,
     generate_model_card,
+    get_callable_name,
     get_peft_config,
     hash_module,
     nanstd,
@@ -279,6 +281,36 @@ def test_create_peft_config_use_peft_true(self):
             assert getattr(peft_config, arg) == value
 
 
+class TestGetCallableName(TrlTestCase):
+    def test_function(self):
+        def accuracy_reward(completions):
+            return [0.0] * len(completions)
+
+        assert get_callable_name(accuracy_reward) == "accuracy_reward"
+
+    def test_partial(self):
+        def reward(completions, threshold):
+            return [0.0] * len(completions)
+
+        assert get_callable_name(functools.partial(reward, threshold=0.5)) == "reward"
+
+    def test_nested_partial(self):
+        def reward(completions, threshold):
+            return [0.0] * len(completions)
+
+        assert get_callable_name(functools.partial(functools.partial(reward), threshold=0.5)) == "reward"
+
+    def test_callable_instance(self):
+        class LengthReward:
+            def __call__(self, completions):
+                return [0.0] * len(completions)
+
+        assert get_callable_name(LengthReward()) == "LengthReward"
+
+    def test_lambda(self):
+        assert get_callable_name(lambda completions: [0.0] * len(completions)) == "<lambda>"
+
+
 class TestNanStd(TrlTestCase):
     def test_nanstd_ignores_nans(self):
         x = torch.tensor([1.0, 2.0, 3.0, float("nan")])

trl/experimental/async_grpo/async_rollout_worker.py

@@ -45,7 +45,7 @@
     parse_response,
 )
 from ...import_utils import is_vllm_available
-from ...trainer.utils import print_prompt_completions_sample
+from ...trainer.utils import get_callable_name, print_prompt_completions_sample
 
 
 logger = get_logger(__name__)
@@ -308,7 +308,7 @@ def __init__(
         self.dataset = dataset
         self._dataset_iter = iter(dataset)
         self.reward_funcs = reward_funcs
-        self.reward_func_names = [f.__name__ for f in reward_funcs]
+        self.reward_func_names = [get_callable_name(f) for f in reward_funcs]
         # `add_response_schema` sets the response template (transformers >= 5.13) or legacy schema for known chat
         # templates, so tool calls can be parsed. Skip if one is already set; warn if it's a migratable legacy schema.
         has_template = getattr(processing_class, "response_template", None) is not None

trl/experimental/online_dpo/online_dpo_trainer.py

@@ -54,7 +54,7 @@
 from ...import_utils import is_vllm_available
 from ...models.utils import prepare_deepspeed, prepare_fsdp, unwrap_model_for_generation
 from ...trainer.base_trainer import _BaseTrainer
-from ...trainer.utils import disable_dropout_in_model, ensure_master_addr_port, get_config_model_id
+from ...trainer.utils import disable_dropout_in_model, ensure_master_addr_port, get_callable_name, get_config_model_id
 from ..utils import DPODataCollatorWithPadding, create_reference_model, empty_cache, prepare_peft_model, truncate_right
 from .online_dpo_config import OnlineDPOConfig
 
@@ -217,7 +217,7 @@ def __init__(
             if isinstance(reward_funcs[i], nn.Module):
                 self.reward_func_names.append(get_config_model_id(reward_funcs[i].config).split("/")[-1])
             else:
-                self.reward_func_names.append(reward_funcs[i].__name__)
+                self.reward_func_names.append(get_callable_name(reward_funcs[i]))
         self.reward_funcs = reward_funcs
 
         # Handle reward processing classes for reward_funcs

trl/experimental/sdpo/sdpo_trainer.py

@@ -51,6 +51,7 @@
     RepeatSampler,
     create_model_from_path,
     disable_dropout_in_model,
+    get_callable_name,
     get_config_model_id,
     identity,
     pad,
@@ -630,7 +631,7 @@ def __init__(
             if isinstance(reward_funcs[i], nn.Module):
                 self.reward_func_names.append(get_config_model_id(reward_funcs[i].config).split("/")[-1])
             else:
-                self.reward_func_names.append(reward_funcs[i].__name__)
+                self.reward_func_names.append(get_callable_name(reward_funcs[i]))
         self.reward_funcs = reward_funcs
 
         if args.reward_weights is not None:

trl/trainer/grpo_trainer.py

@@ -81,6 +81,7 @@
     create_model_from_path,
     disable_dropout_in_model,
     entropy_from_logits,
+    get_callable_name,
     get_config_model_id,
     identity,
     nanmax,
@@ -492,7 +493,7 @@ def __init__(
             if isinstance(reward_funcs[i], nn.Module):  # Use Module over PretrainedModel for compat w/ compiled models
                 self.reward_func_names.append(get_config_model_id(reward_funcs[i].config).split("/")[-1])
             else:
-                self.reward_func_names.append(reward_funcs[i].__name__)
+                self.reward_func_names.append(get_callable_name(reward_funcs[i]))
         self.reward_funcs = reward_funcs
 
         # Reward weights

trl/trainer/rloo_trainer.py

@@ -64,6 +64,7 @@
     create_model_from_path,
     disable_dropout_in_model,
     entropy_from_logits,
+    get_callable_name,
     get_config_model_id,
     identity,
     nanmax,
@@ -407,7 +408,7 @@ def __init__(
             if isinstance(reward_funcs[i], nn.Module):  # Use Module over PretrainedModel for compat w/ compiled models
                 self.reward_func_names.append(get_config_model_id(reward_funcs[i].config).split("/")[-1])
             else:
-                self.reward_func_names.append(reward_funcs[i].__name__)
+                self.reward_func_names.append(get_callable_name(reward_funcs[i]))
         self.reward_funcs = reward_funcs
 
         self._has_async_funcs = any(inspect.iscoroutinefunction(func) for func in self.reward_funcs)

trl/trainer/utils.py

@@ -13,14 +13,15 @@
 # limitations under the License.
 
 import asyncio
+import functools
 import hashlib
 import importlib.resources as pkg_resources
 import os
 import random
 import socket
 import threading
 import types
-from collections.abc import Mapping, Sequence, Sized
+from collections.abc import Callable, Mapping, Sequence, Sized
 from contextlib import contextmanager
 from importlib.metadata import version
 from itertools import accumulate
@@ -182,6 +183,17 @@ def disable_dropout_in_model(model: torch.nn.Module) -> None:
             module.p = 0
 
 
+def get_callable_name(func: Callable) -> str:
+    """
+    Return a display name for a callable, supporting the picklable reward forms: module-level functions,
+    [`functools.partial`](https://docs.python.org/3/library/functools.html#functools.partial) (unwrapped to the wrapped
+    function's name), and callable class instances (which fall back to their class name).
+    """
+    while isinstance(func, functools.partial):
+        func = func.func
+    return getattr(func, "__name__", type(func).__name__)
+
+
 def get_quantization_config(model_args: ModelConfig) -> BitsAndBytesConfig | None:
     if model_args.load_in_4bit:
         quantization_config = BitsAndBytesConfig(