[BREAKING][rollout] feat: Added asynchronous reward model calculation in agent loop

[echo-rain]

What does this PR do?

This PR will be based on PR#3055, and will further support asynchronous calculation of reward models based on the agent loop which only supports asynchronous reward function calculation.

Checklist Before Starting

• Search for similar PRs. Paste at least one query link here: ...
• Format the PR title as [{modules}] {type}: {description} (This will be checked by the CI)
  • {modules} include fsdp, megatron, sglang, vllm, rollout, trainer, ci, training_utils, recipe, hardware, deployment, ray, worker, single_controller, misc, perf, model, algo, env, tool, ckpt, doc, data
  • If this PR involves multiple modules, separate them with , like [megatron, fsdp, doc]
  • {type} is in feat, fix, refactor, chore, test
  • If this PR breaks any API (CLI arguments, config, function signature, etc.), add [BREAKING] to the beginning of the title.
  • Example: [BREAKING][fsdp, megatron] feat: dynamic batching

Test

For changes that can not be tested by CI (e.g., algorithm implementation, new model support), validate by experiment(s) and show results like training curve plots, evaluation results, etc.

API and Usage Example

If you want to use this feature, you need to add the following configuration to the startup script configuration item

    reward_model.enable_resource_pool=True 
    reward_model.n_gpus_per_node=1 
    reward_model.nnodes=1 

Design & Code Changes

Demonstrate the high-level design if this PR is complex, and list the specific changes.

Checklist Before Submitting

Important

Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review.

• Read the Contribute Guide.
• Apply pre-commit checks: pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always
• Add / Update the documentation.
• Add unit or end-to-end test(s) to the CI workflow to cover all the code. If not feasible, explain why: ...
• Once your PR is ready for CI, send a message in the ci-request channel in the verl Slack workspace. (If not accessible, please try the Feishu group (飞书群).)

[CLAassistant]

All committers have signed the CLA.

[gemini-code-assist]

Code Review

This pull request introduces asynchronous reward model calculation within the agent loop. It achieves this by passing the reward model worker group (rm_wg) down to the AgentLoopWorker and RewardManagerWorker. A new reward_wrapper method is introduced to handle the asynchronous scoring. The changes also include support for a dedicated resource pool for the reward model, improving resource management. The overall implementation is consistent and correctly enables the new asynchronous workflow. My review includes one high-severity suggestion to improve configuration validation by using ValueError instead of assert, which can be disabled in production environments.

[gemini-code-assist]

Using assert for validating user configuration can be risky, as assertions can be disabled with Python's -O (optimize) flag. This could lead to silent failures or unexpected behavior in production environments if the configuration is invalid. It's safer to raise a ValueError for configuration errors to ensure the validation is always performed.

Suggested change

	assert config.reward_model.n_gpus_per_node > 0, "config.reward_model.n_gpus_per_node must be greater than 0"
	assert config.reward_model.nnodes > 0, "config.reward_model.nnodes must be greater than 0"
	if config.reward_model.n_gpus_per_node <= 0:
	raise ValueError("config.reward_model.n_gpus_per_node must be greater than 0")
	if config.reward_model.nnodes <= 0:
	raise ValueError("config.reward_model.nnodes must be greater than 0")

[wuxibin89]

There's a big difference between reward model as service and reward model as worker group.

Reward model as service: request sent to OpenAI compatible http server with chat completion api, e.g vllm server, sglang server. These http servers can accept individual request and batch them automatically. Hence in agent loop, once a sequence generation finished, we can send it to reward service immediately.

However, reward model as worker group lack the ability of automatic request batching, so if we send individual request in agent loop, these requests are processed serially, which is very inefficient.

[echo-rain]

There's a big difference between reward model as service and reward model as worker group.

Reward model as service: request sent to OpenAI compatible http server with chat completion api, e.g vllm server, sglang server. These http servers can accept individual request and batch them automatically. Hence in agent loop, once a sequence generation finished, we can send it to reward service immediately.

However, reward model as worker group lack the ability of automatic request batching, so if we send individual request in agent loop, these requests are processed serially, which is very inefficient.

You are right. I will add an executor in the next commit for the real batch execution, and only register the request when processing a single request.

[echo-rain]

There's a big difference between reward model as service and reward model as worker group.

Reward model as service: request sent to OpenAI compatible http server with chat completion api, e.g vllm server, sglang server. These http servers can accept individual request and batch them automatically. Hence in agent loop, once a sequence generation finished, we can send it to reward service immediately.

However, reward model as worker group lack the ability of automatic request batching, so if we send individual request in agent loop, these requests are processed serially, which is very inefficient.

The latest commit has added dynamic batching for reward model inference. While verifying performance data, we found that using a separate resource pool for the reward model in the on-policy scenario resulted in low computing resource utilization.

To address the issue of computing resource utilization, a potential solution would be to have rollout and reward pools use mutually exclusive resource pools, with actors sharing the computing resources of both. However, the current resource pool isolation mechanism does not appear to be easily modified to achieve this solution, so it is not implemented in this PR.