feat: overlap NeMo Gym init with vLLM init in distillation - #2793
Merged
Conversation
NeMo-Gym distillation needs the Gym actor to share the deferred vLLM server URLs produced during setup. Moving actor creation into distillation setup keeps the runner from constructing a second environment path and matches the setup-owned initialization used by GRPO-style NeMo-Gym flows. Constraint: origin/main already contains the deferred vLLM load API Rejected: Keep actor construction in the runner | duplicates setup responsibility and misses the overlapped vLLM/Gym initialization path Confidence: high Scope-risk: narrow Tested: uv run ruff format --check on changed files Tested: uv run --group dev pre-commit run --files examples/nemo_gym/run_distillation_nemo_gym.py examples/run_distillation.py nemo_rl/algorithms/distillation.py tests/unit/algorithms/test_distillation.py Tested: uv run --all-groups focused distillation NeMo-Gym unit checks Tested: Slurm qwen3-0.6B NeMo-Gym distillation functional smoke job 212974 Not-tested: full unit suite due local pytest/Ray fixture setup cost Signed-off-by: Meng Xin <mxin@nvidia.com>
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
mxinO
changed the title
…-fix Signed-off-by: Meng Xin <mxin@nvidia.com> # Conflicts: # examples/nemo_gym/run_distillation_nemo_gym.py
mxinO
changed the title
mxinO
marked this pull request as ready for review
June 12, 2026 05:33
|
Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
Contributor
Author
|
/ok to test |
yfw
reviewed
Jun 12, 2026
Distillation now mirrors the GRPO NeMo-Gym startup config added with the release-image venv caching work. The Gym actor reads uv cache and venv locations from its global config, so relying only on actor runtime environment inheritance can still rebuild server venvs during spinup. Constraint: NeMo-Gym consumes uv_cache_dir and uv_venv_dir from initial_global_config_dict. Rejected: Rely on make_actor_runtime_env environment inheritance | the Gym server config path does not read these values from environment variables. Confidence: high Scope-risk: narrow Tested: uv run --no-sync ruff format --check nemo_rl/algorithms/distillation.py tests/unit/algorithms/test_distillation.py Tested: uv run --no-sync ruff check nemo_rl/algorithms/distillation.py tests/unit/algorithms/test_distillation.py Tested: uv run --no-sync python -m py_compile nemo_rl/algorithms/distillation.py tests/unit/algorithms/test_distillation.py Not-tested: pytest in this local worktree; torch is not installed in the local uv environment. Signed-off-by: Meng Xin <mxin@nvidia.com>
The distillation startup path now mirrors GRPO by forwarding image-baked uv cache and venv paths into NeMo-Gym. Cover both default insertion and explicit user override preservation so future edits do not accidentally replace setdefault semantics. Constraint: Reviewer requested parity with GRPO's NeMo-Gym release-image venv config path. Confidence: high Scope-risk: narrow Tested: uv run --no-sync python -m py_compile nemo_rl/algorithms/distillation.py tests/unit/algorithms/test_distillation.py Tested: uv run --no-sync ruff format --check nemo_rl/algorithms/distillation.py tests/unit/algorithms/test_distillation.py Tested: uv run --no-sync ruff check nemo_rl/algorithms/distillation.py tests/unit/algorithms/test_distillation.py Not-tested: targeted pytest; local uv environment lacks torch imported by tests/unit/conftest.py. Signed-off-by: Meng Xin <mxin@nvidia.com>
…-fix Signed-off-by: Meng Xin <mxin@nvidia.com>
Contributor
Author
|
/ok to test |
yfw
approved these changes
Jun 13, 2026
yuki-97
approved these changes
Jun 13, 2026
yuki-97
added
CI:Lfast
Contributor
|
L1_Functional_Tests_Megatron_3 also fails in main L1_Functional_Tests_Megatron_3 and unrelated to this PR, so change to FYI @mxinO we can use |
yuki-97
enabled auto-merge (squash)
June 13, 2026 07:37
5 tasks
ashors1
pushed a commit
that referenced
this pull request
Jun 27, 2026
Signed-off-by: Meng Xin <mxin@nvidia.com> Signed-off-by: Anna Shors <ashors@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do ?
align nemo-gym distillation with #2741
This PR ports the core NeMo-Gym startup ownership model from #2741 to distillation: the Gym actor is created inside
distillation.setup(), spun up via_spinup(), and returned to the runner instead of being constructed inrun_distillation_nemo_gym.py.The current implementation overlaps deferred vLLM model load with NeMo-Gym startup, which fixes the removed
health_checkpath and avoids runner-side Gym construction.One remaining parity gap with #2741 is startup overlap with policy initialization. GRPO overlaps Gym startup with vLLM/policy initialization. For distillation, teacher policy setup is separate and should likely stay sequential/offloaded first, but student policy initialization can follow the GRPO pattern:
vLLM -> student_policysequentially in one task while Gym spins up concurrentlyThis is not required for functional correctness, but it would make distillation match the intended #2741 startup design more closely and reduce setup time.
Issues
List issues that this PR closes (syntax):
Usage
# Add a code snippet demonstrating how to use thisBefore your PR is "Ready for review"
Pre checks:
Additional Information