Speedup the training process when the reward_function is cpu-intensive

[peterjc123]

The generation and rewarding process are all running under one async context, which means it utilizes only one CPU. It can be very slow, often 2-4x slower than other frameworks that are based on ray tasks. For a different perspective, by swapping the reward function with a placeholder function, the rollout time reduces by 50%.

[peterjc123]

With a deeper look, it seems to be that request.post is the actual blocker. So it is solved by setting --use-distributed-post.

[zimingyou01]

How did you solve your problem? I think setting --use-distributed-post can only make requests to sglang router in parallel but all the samples' reward computation remains in the same asyncio event loop. So this setting can not accelarate rollout if reward calculation is cpu-heavy, right? @peterjc123

[peterjc123]

@zimingyou01 Yes, you are right. Turns out that my reward function is actually not cpu intensive, and the real problem is the limited size of concurrent connections in a single process, so --use-distributed-post solves my problem. Maybe you could try wrapping the reward calls with ProcessPoolExecutor, I have tried that but it doesn't help in my case.