feat: overlap NeMo Gym init with vLLM init by yfw · Pull Request #2741 · NVIDIA-NeMo/RL

yfw · 2026-06-08T22:57:25Z

What does this PR do ?

Pulls in and squashes several commits from the super / ultra branches related to gym and vllm startup. Includes a Gym bump to latest main. Pulled in the commits together and squashed because several early commits were later partially un-done by the "Pipeclean GenRMModel flow and local_vllm_model use" MR (f9cce2ae33776de372a7a7801de5b933c495fe45).

Claude generated description below:

Re-architect NeMo-Gym startup so the gym is created inside grpo.setup() and its spin-up overlaps vLLM/policy initialization: vLLM does a deferred model load (reserve ports first), the gym actor starts via _spinup(), and vLLM + policy + gym initialize concurrently via init_tasks. This converges main's NeMo-Gym init to the ultra-v3 final shape.

Changes:

grpo.setup(): build the gym in-loop and return its actor (10- → 11-tuple)
vLLM: defer_model_load + load_and_start() to overlap with gym spin-up
nemo_gym: split init → _spinup(); add ray_num_gpus_per_node / ray_namespace
run_grpo_nemo_gym.py: consume the gym actor from setup()
bump Gym submodule to latest main

Provenance — ultra-v3-posttraining gym-init series (public GitHub, NVIDIA-NeMo/RL):
Reflected in this branch (net):
- Support for ray spinup within gym:
c722c91
- overlap NeMo Gym init with vLLM init:
eb401ba
In the series but UNDONE by the Pipeclean GenRMModel commit, so intentionally NOT in this branch:
- add genrm rlhf to gym:
5bbf159
- Fix genrm for async gym:
3f391da
- Fix for ray state API usage (list_nodes with limit):
5c6d5e4
- fix genrm reasoning off:
48fe510
- Pass explicit PG to gym:
46e5ec5

Final state defined by:

Pipeclean GenRMModel flow and local_vllm_model use (f9cce2ae33776de372a7a7801de5b933c495fe45)

Update (2026-06-09): Separate the vllm deferred load changes to a separate PR: #2750

Issues

List issues that this PR closes (syntax):

Usage

You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

...

Point 3rdparty/Gym-workspace/Gym at the current tip of NVIDIA-NeMo/Gym main in preparation for converging NeMo-Gym init to the ultra-v3 final state. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com> Co-Authored-By: Felipe Vieira Frujeri <ffrujeri@nvidia.com> Co-Authored-By: Peter Jin <pjin@nvidia.com> Co-Authored-By: Yashaswi Karnati <ykarnati@nvidia.com>

Port the deferred-load mechanism from the ultra-v3 lineage (eb401ba), adapted to main's refactored vLLM workers: - VllmGeneration: defer_model_load flag; when set, workers only reserve ports and dp_openai_server_base_urls is populated immediately via _collect_reserved_urls(); load_and_start() performs the heavy load later. - BaseVllmGenerationWorker: split __init__ into _init_config()/_apply_vllm_patches() /_load_model(); patches now applied for all workers (incl. non-owners). - VllmAsyncGenerationWorker: reserve a listening socket in __init__ and hand it to uvicorn via server.serve(sockets=) (no close/rebind race); load_model() and get_reserved_url() added. - Add test_vllm_deferred_model_load. Enables overlapping vLLM model loading with NeMo Gym spinup. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com> Co-Authored-By: Felipe Vieira Frujeri <ffrujeri@nvidia.com> Co-Authored-By: Peter Jin <pjin@nvidia.com> Co-Authored-By: Yashaswi Karnati <ykarnati@nvidia.com>

Converge NemoGym to the ultra-v3 final shape (c722c91, post-f9cce2ae cleanup): - Move the head-server/rollout-helper startup out of __init__ into an explicit _spinup() so the actor can be created cheaply and spun up once vLLM server URLs are ready (overlapping with vLLM model load). - Add NemoGymConfig.ray_num_gpus_per_node and ray_namespace; surface them into the gym global config when set. (No ray_gpu_nodes/ray_gpu_pgs — removed by the pipeclean cleanup.) - Remove health_check; callers now use _spinup as the blocking spinup wait. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com> Co-Authored-By: Felipe Vieira Frujeri <ffrujeri@nvidia.com> Co-Authored-By: Peter Jin <pjin@nvidia.com> Co-Authored-By: Yashaswi Karnati <ykarnati@nvidia.com>

Converge grpo.setup() to the ultra-v3 gym-init shape (c722c91 + eb401ba, post-f9cce2ae cleanup), so NeMo Gym spinup overlaps vLLM model loading and so future ultra cherry-picks land cleanly: - When NeMo Gym is enabled, reserve vLLM ports via deferred model load, then run init_vllm/init_policy/init_nemo_gym concurrently via ThreadPoolExecutor (sequential when colocated). init_nemo_gym builds NemoGymConfig with the reserved URLs + ray_num_gpus_per_node/ray_namespace and calls actor._spinup(). - setup() now returns the gym actor (3rd element); run_grpo_nemo_gym.py consumes it instead of building/health-checking the gym itself. - Non-gym vLLM/SGLang/Megatron paths keep main's existing initialize_generation_with_policy flow. Cluster creation is unchanged (gym scheduling via soft NodeAffinity); GPU-node reservation is a follow-up. - invalid_tool_call_patterns/thinking_tags are preserved as NemoGymConfig fields. Excludes the tabled gym-venv-caching (05eb916) so it can be applied separately. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com> Co-Authored-By: Felipe Vieira Frujeri <ffrujeri@nvidia.com> Co-Authored-By: Peter Jin <pjin@nvidia.com> Co-Authored-By: Yashaswi Karnati <ykarnati@nvidia.com>

Apply ruff format to the deferred-load changes and remove the unused `ray` import in run_grpo_nemo_gym.py (the health_check call it served was removed now that the gym actor is spun up inside setup()). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com> Co-Authored-By: Felipe Vieira Frujeri <ffrujeri@nvidia.com> Co-Authored-By: Peter Jin <pjin@nvidia.com> Co-Authored-By: Yashaswi Karnati <ykarnati@nvidia.com>

copy-pr-bot · 2026-06-08T22:57:29Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

github-actions · 2026-06-08T22:58:32Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 97b1fe8 (PR #2741 from yifu/ultra_gym)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

- setup() returns an 11-tuple (the NeMo-Gym actor is at index 2). Update the three non-gym callers (run_grpo, run_grpo_sliding_puzzle, run_vlm_grpo) to unpack the extra element so they don't fail with a ValueError, and fix the setup() docstring, which listed a stale tuple. - Drop the dead ray_namespace / ray_num_gpus_per_node gym config fields: the pinned Gym (9b52dc8) has zero references to them, so they were set into the gym global config but silently ignored. Keep ray_head_node_address, which the Gym does consume. The _apply_vllm_patches() non-owner write race is handled separately by #2739, which lands before this PR. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

github-actions · 2026-06-08T23:54:22Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 165ac4f (PR #2741 from yifu/ultra_gym)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

yfw · 2026-06-08T23:56:17Z

/ok to test 165ac4f

- setup() returns an 11-tuple (the NeMo-Gym actor is at index 2). Update the three non-gym callers (run_grpo, run_grpo_sliding_puzzle, run_vlm_grpo) to unpack the extra element so they don't fail with a ValueError, and fix the setup() docstring, which listed a stale tuple. - Drop the dead ray_namespace / ray_num_gpus_per_node gym config fields: the pinned Gym (9b52dc8) has zero references to them, so they were set into the gym global config but silently ignored. Keep ray_head_node_address, which the Gym does consume. The _apply_vllm_patches() non-owner write race is handled separately by #2739, which lands before this PR. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

yfw · 2026-06-09T00:36:36Z

/ok to test 86ebabb

yfw · 2026-06-09T18:05:22Z

/ok to test 86ebabb

github-actions · 2026-06-10T15:20:38Z

✅ Submodule Fast-Forward Check Results

Check based on commit: cf139fe (PR #2741 from yifu/ultra_gym)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

yfw · 2026-06-10T15:43:40Z

/ok to test d00a5e1

github-actions · 2026-06-10T15:44:43Z

✅ Submodule Fast-Forward Check Results

Check based on commit: d00a5e1 (PR #2741 from yifu/ultra_gym)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

yfw · 2026-06-10T15:59:44Z

/ok to test d00a5e1

Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Felipe Vieira Frujeri <ffrujeri@nvidia.com> Co-authored-by: Peter Jin <pjin@nvidia.com> Co-authored-by: Yashaswi Karnati <ykarnati@nvidia.com>

Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Felipe Vieira Frujeri <ffrujeri@nvidia.com> Co-authored-by: Peter Jin <pjin@nvidia.com> Co-authored-by: Yashaswi Karnati <ykarnati@nvidia.com> Signed-off-by: Anna Shors <ashors@nvidia.com>

yfw and others added 5 commits June 7, 2026 11:52

yfw requested review from a team as code owners June 8, 2026 22:57

yfw requested review from ananthsub and ffrujeri June 8, 2026 22:58

yfw requested a review from yashaswikarnati June 8, 2026 23:01

yfw commented Jun 8, 2026

View reviewed changes

Comment thread nemo_rl/algorithms/grpo.py

Comment thread nemo_rl/models/generation/vllm/vllm_worker.py

Comment thread nemo_rl/environments/nemo_gym.py Outdated

yfw requested a review from a team as a code owner June 8, 2026 23:53

yfw added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Jun 8, 2026

copy-pr-bot Bot temporarily deployed to public June 8, 2026 23:56 Inactive

copy-pr-bot Bot temporarily deployed to public June 8, 2026 23:57 Inactive

copy-pr-bot Bot temporarily deployed to public June 9, 2026 00:00 Inactive

yfw added 3 commits June 8, 2026 17:26

chore(grpo): sort imports to satisfy ruff isort pre-commit hook

8432d24

Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

Merge remote-tracking branch 'origin/main' into yifu/ultra_gym

309d7ed

Merge remote-tracking branch 'origin/main' into yifu/ultra_gym

86ebabb

copy-pr-bot Bot temporarily deployed to public June 9, 2026 00:37 Inactive

copy-pr-bot Bot temporarily deployed to test June 9, 2026 00:39 Inactive

copy-pr-bot Bot temporarily deployed to public June 9, 2026 00:41 Inactive

yfw mentioned this pull request Jun 9, 2026

feat(vllm): add deferred model load for overlapped NeMo Gym init #2750

Merged

4 tasks

yfw changed the title ~~feat: init NeMo-Gym in setup() with overlapped deferred vLLM load~~ feat: overlap NeMo Gym init with vLLM init Jun 9, 2026

yfw added CI:L1 Run doctests, unit tests, and functional tests and removed CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) labels Jun 9, 2026

Merge branch 'main' into yifu/ultra_gym

cf139fe

Bump gym

d00a5e1

Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

copy-pr-bot Bot had a problem deploying to public June 10, 2026 15:44 Failure

copy-pr-bot Bot temporarily deployed to test June 10, 2026 15:47 Inactive

copy-pr-bot Bot had a problem deploying to public June 10, 2026 16:03 Failure

copy-pr-bot Bot temporarily deployed to public June 10, 2026 16:51 Inactive

copy-pr-bot Bot temporarily deployed to public June 10, 2026 16:55 Inactive

copy-pr-bot Bot temporarily deployed to test June 10, 2026 17:06 Inactive

ananthsub approved these changes Jun 10, 2026

View reviewed changes

yfw enabled auto-merge (squash) June 10, 2026 17:56

yfw merged commit 6a9146f into main Jun 10, 2026
114 of 121 checks passed

yfw deleted the yifu/ultra_gym branch June 10, 2026 20:33

yuki-97 mentioned this pull request Jun 11, 2026

nemo_gym 'ActorHandle' object has no attribute 'health_check' #2784

Closed

mxinO mentioned this pull request Jun 12, 2026

feat: overlap NeMo Gym init with vLLM init in distillation #2793

Merged

4 tasks

Navigation Menu

Uh oh!

feat: overlap NeMo Gym init with vLLM init - #2741

feat: overlap NeMo Gym init with vLLM init#2741
yfw merged 11 commits into
mainfrom
yifu/ultra_gym

yfw commented Jun 8, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Jun 8, 2026

Uh oh!

github-actions Bot commented Jun 8, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented Jun 8, 2026

Uh oh!

yfw commented Jun 8, 2026

Uh oh!

yfw commented Jun 9, 2026

Uh oh!

yfw commented Jun 9, 2026

Uh oh!

github-actions Bot commented Jun 10, 2026

Uh oh!

yfw commented Jun 10, 2026

Uh oh!

github-actions Bot commented Jun 10, 2026

Uh oh!

yfw commented Jun 10, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

yfw commented Jun 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Issues

Usage

Before your PR is "Ready for review"

Additional Information

Uh oh!

copy-pr-bot Bot commented Jun 8, 2026

Uh oh!

github-actions Bot commented Jun 8, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented Jun 8, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

yfw commented Jun 8, 2026

Uh oh!

yfw commented Jun 9, 2026

Uh oh!

yfw commented Jun 9, 2026

Uh oh!

github-actions Bot commented Jun 10, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

yfw commented Jun 10, 2026

Uh oh!

github-actions Bot commented Jun 10, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

yfw commented Jun 10, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

yfw commented Jun 8, 2026 •

edited

Loading