Repository navigation
Run the AsyncGRPO train test against a transformers-backed vLLM server - #8137
qgallouedec wants to merge 21 commits into
Conversation
The trainer's real rollout worker and weight sync now run in test_init_minimal and test_train, against a local HTTP server that answers the vLLM endpoints they call and samples with transformers.
NCCL refuses two ranks on one device, so the fixture hands the trainer's weight stream to the server in place of the NCCL broadcast. The server now samples from the current policy instead of the initial weights.
# Conflicts: # tests/experimental/test_async_grpo_trainer.py
With the server following the policy, lr=0.1 collapses it to one completion within about 20 steps. Every group then has equal rewards and is dropped for zero std, so no step completes and the test waits forever.
Their trainers sync weights through vLLM's client. Installing vLLM in the main experimental job would pin torch to vLLM's version (2.13 instead of 2.15) for every experimental test.
The workers reach the server over HTTP with aiohttp and import nothing from vLLM. The weight sync, which does, keeps its own check in WeightTransferClient.
The fixture gives the trainer the server in place of WeightTransferClient, which needs vLLM and one device per NCCL rank. The train tests then run without vLLM, and the gates that only covered the worker's vLLM check go.
This reverts commit 201ab10.
# Conflicts: # tests/experimental/test_async_distillation_trainer.py
The fixture starts a server per model id. The server reports a 512-token context, answers the top-k candidate-set endpoint, serves vision-language checkpoints, replays scripted completions, and joins its request threads on close. test_train no longer needs its token_budget.
# Conflicts: # tests/experimental/test_async_grpo_trainer.py
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 86ef884. Configure here.
Codex Review Summary: ✅ CompletedThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
| return Handler | ||
|
|
||
|
|
||
| @pytest.fixture |
There was a problem hiding this comment.
should this fixture define a scope here ?
There was a problem hiding this comment.
No, on purpose. Each test trains the served model and may give the server a script, so a shared server would carry both into the next test. It also uses monkeypatch, which is function-scoped, and the model loads in a few seconds.
…erClient, and name the fixture make_vllm_server
…at train without a server, instead of applying them to every test
Part of #7841, prepares #8146.
The AsyncGRPO train tests stub the rollout worker, so rewards, tools and sampling never run in them. This adds
TransformersVLLMServer, a local HTTP server that answers the vLLM endpoints the worker calls and samples with transformers on the trainer's GPU. Thevllm_serverfixture also hands it to the trainer in place ofWeightTransferClient, which needs vLLM and one GPU per NCCL rank, so each sync loads the weights into the server.test_trainnow runs the real worker against it, and the rollout workers drop their vLLM check, since they only use aiohttp.vllm servetakes 157 s for the same model on an H100.Note
Low Risk
Primarily test and dependency changes plus relaxing an import guard on rollout workers; production weight sync still uses vLLM when configured.
Overview
Async
AsyncRolloutWorkerconstruction no longer requires a full vLLM install—only aiohttp—matching the worker’s HTTP-only client role;aiohttp>=3.13.3is added to dev dependencies for that reason.Tests gain
TransformersVLLMServer, a local HTTP shim that implements the rollout endpoints (including sampling masks / inference routes) and samples with transformers on the trainer device, plusInProcessWeightTransferClientto copy synced weights into that server instead of NCCLWeightTransferClient.test_traindrives the real rollout worker against this server (withvaried_rewardso groups stay scorable) and asserts weights reach the served model; the stub rollout/weight-transfer fixtures are opt-in viausefixturesrather than autouse.@require_vllmis dropped from rollout-worker wiring tests that only instantiate the worker.Reviewed by Cursor Bugbot for commit 1a688df. Bugbot is set up for automated code reviews on this repo. Configure here.