Fix pickle error in prime reward manager when prime in main_ppo.py

[cgpeter96]

fixed #1293

[wuxibin89]

ThreadPoolExecutor is not a good idea since compute_score is compute intensive, python GIL prevents it to use multi cores.

[cgpeter96]

ThreadPoolExecutor is not a good idea since compute_score is compute intensive, python GIL prevents it to use multi cores.

Understood, I will update the code shortly.

[patrik-bartak]

I actually ran into the same issue as you, and I think your solution is unnecessarily complex. I also got the issue that the function is not pickleable - the issue is that the wrapped_fn if defined within a function, and not top level.

A simple solution that works for me is replacing the wrapped function with a partial over the raw function.

...
    def wrapped_fn(*args, **kwargs):
        return raw_fn(*args, **kwargs, **reward_kwargs)

    return wrapped_fn

->

...
    return partial(raw_fn, **reward_kwargs)

No other changes are needed. Do you agree?

I don't see what the point of wrapping the function is, other than passing in the reward_kwargs, which the partial already does. I think someone added it without testing it with the prime reward manager.

[patrik-bartak]

Also, the function get_custom_reward_fn is duplicated across main_dapo, main_eval, main_ppo. If you want I can make a PR to share the implementation of this function.

[cgpeter96]

Also, the function get_custom_reward_fn is duplicated across main_dapo, main_eval, main_ppo. If you want I can make a PR to share the implementation of this function.

Of course, simple implementation is best :)

I actually ran into the same issue as you, and I think your solution is unnecessarily complex. I also got the issue that the function is not pickleable - the issue is that the wrapped_fn if defined within a function, and not top level.

A simple solution that works for me is replacing the wrapped function with a partial over the raw function.

...
    def wrapped_fn(*args, **kwargs):
        return raw_fn(*args, **kwargs, **reward_kwargs)

    return wrapped_fn

->

...
    return partial(raw_fn, **reward_kwargs)

No other changes are needed. Do you agree?

I don't see what the point of wrapping the function is, other than passing in the reward_kwargs, which the partial already does. I think someone added it without testing it with the prime reward manager.

YES! You are right. The simple solution is best

[cgpeter96]

Also, the function get_custom_reward_fn is duplicated across main_dapo, main_eval, main_ppo. If you want I can make a PR to share the implementation of this function.

Yes!

[cgpeter96]

@wuxibin89 Could you given me some feedback 😊. This pr will help someone who wanna use prime reward manager in ppo/grpo training but failed.

[CLAassistant]

All committers have signed the CLA.

[vadimkantorov]

Hi @wuxibin89 @vermouth1992 ! We also stumbled in this issue with a custom reward function defined in a sidekick python file: [Error] Task failed: Can't get local object 'get_custom_reward_fn.<locals>.wrapped_fn', completion: ...

We tried the partial(...) solution but then hit another bug: [Error] Task failed: Can't pickle <function compute_score at 0x7f23310b4ea0>: it's not the same object as custom_module.compute_score, completion: ...

We're calling verl trainer as follows:

python -m verl.trainer.main_ppo \
    custom_reward_function.path=/mnt/fs/verl/custom_function.py \
    custom_reward_function.name=compute_score \
    reward_model.reward_manager=prime \
    ...

Would you have any suggestions how to fix this?

[zhourunlong]

Hi @wuxibin89 @vermouth1992 ! We also stumbled in this issue with a custom reward function defined in a sidekick python file: [Error] Task failed: Can't get local object 'get_custom_reward_fn.<locals>.wrapped_fn', completion: ...

We tried the partial(...) solution but then hit another bug: [Error] Task failed: Can't pickle <function compute_score at 0x7f23310b4ea0>: it's not the same object as custom_module.compute_score, completion: ...

We're calling verl trainer as follows:

python -m verl.trainer.main_ppo \
    custom_reward_function.path=/mnt/fs/verl/custom_function.py \
    custom_reward_function.name=compute_score \
    reward_model.reward_manager=prime \
    ...

Would you have any suggestions how to fix this?

Meeting exactly the same issue.

[george1459]

I ran into this issue as well (even with #2239). Here is my fix for the issue which worked for me.
Let me know if you guys are interested in having this as a PR.
@vadimkantorov @zhourunlong perhaps you guys can give this a try if interested.

[warmsnow-sh]

This code seems to be inconsistent with the latest version. Are there any latest solutions to this problem at present?

[warmsnow-sh]

I ran into this issue as well (even with #2239). Here is my fix for the issue which worked for me. Let me know if you guys are interested in having this as a PR. @vadimkantorov @zhourunlong perhaps you guys can give this a try if interested.

I tried your plan, but it reported an error
(TaskRunner pid=1138472) [Error] Task failed: 'list' object has no attribute 'get',
I'm still checking if there are any other potential issues with my code. I wonder why this problem(prime manager) hasn't been fixed for so long