Async validation takes too long in some envs
- Dominant language
- Python
- Stars
- 2k
- Forks
- 562
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
**Describe the bug**
In the async grpo, when the validation part is called, rollout is sent using batches through `run_async_multi_turn_rollout`. But the reward computation in `run_async_multi_turn_rollout` is not designed to efficiently process batches, since the reward computation is not natively async and is blocking. This is particularly problematic in LLM-as-a-Judge environments since they take a very long time to respond.
**Steps/Code to reproduce bug**
Run async grpo validation with any llm-as-a-judge environment.
**Expected behavior**
Concurrent computation of reward when `run_async_multi_turn_rollout` is called, either by spawning one thread per rollout, or making reward computation natively async. For instance running [this line](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/experience/rollouts.py#L438) in a separate thread.
**Additional context**
Contributor guide
Research direction
Start with nemo_rl/experience/rollouts.py, especially run_async_multi_turn_rollout and the code at line 438. Reproduce the delay by running async GRPO validation with an LLM-as-a-Judge environment, then trace how batched rollouts invoke blocking reward computation. Done means reward computation proceeds concurrently during validation without the current batch-related delay.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 42/100