fix(distillation): use _spinup instead of removed health_check for NeMo-Gym - #2786
Merged
Conversation
Distillation can now reuse the NeMo Gym rollout path that GRPO already uses, while keeping the new runner and smoke coverage scoped to the distillation integration. Constraint: NeMo Gym requires the vLLM async HTTP generation path. Rejected: Mirror GRPO's full NeMo Gym test matrix | distillation only adds orchestration around the shared rollout and environment layers Confidence: medium Scope-risk: moderate Tested: tests/unit/algorithms/test_distillation.py::test_distillation_train_uses_nemo_gym_rollout_when_enabled in container Tested: ruff check and ruff format --check on touched Python files Tested: tests/functional/distillation_nemo_gym.sh in container via Slurm job 210602 Signed-off-by: Meng Xin <mxin@nvidia.com>
The NeMo Gym distillation example now uses a distinct default teacher and the current DTensor v2 path used by base distillation. The functional smoke keeps CI cheap by overriding to the small Qwen3 base/instruct pair while still exercising NeMo Gym rollout, teacher logits, and student training with nonzero KL. Constraint: Split PR must stay scoped to NeMo Gym distillation without quant changes Rejected: Same teacher/student smoke | it can produce zero KL and miss the distillation path Rejected: DTensor v1 | smoke jobs stalled during teacher state-dict loading Confidence: high Scope-risk: narrow Tested: uv run --no-sync ruff check examples/nemo_gym/run_distillation_nemo_gym.py nemo_rl/algorithms/distillation.py tests/unit/algorithms/test_distillation.py Tested: uv run --no-sync ruff format --check examples/nemo_gym/run_distillation_nemo_gym.py nemo_rl/algorithms/distillation.py tests/unit/algorithms/test_distillation.py Tested: bash -n tests/functional/distillation_nemo_gym.sh && bash -n tests/functional/L1_Functional_Tests_Gym.sh Tested: Slurm unit job 210733, 12 passed Tested: Slurm functional smoke job 210729, train/loss=0.8154925107955933, generation=2.2170605659484863s, mean_gen_tokens=128.0 Not-tested: Full unmodified 1.7B-teacher example run; functional smoke intentionally overrides to the cheap small-model pair Signed-off-by: Meng Xin <mxin@nvidia.com>
NeMo Gym is now available for on-policy distillation, so the public docs should not read as GRPO-only. The usage examples call out required dataset path overrides because the checked-in config intentionally keeps placeholder paths. Constraint: Split PR stays scoped to NeMo Gym distillation docs and examples Rejected: Present the example config as runnable without overrides | the config contains placeholder dataset paths Confidence: high Scope-risk: narrow Tested: git diff --check -- README.md docs/about/algorithms/on-policy-distillation.md docs/design-docs/nemo-gym-integration.md Tested: uv run --group docs python -m sphinx.cmd.build -b dummy docs docs/_build/dummy Signed-off-by: Meng Xin <mxin@nvidia.com>
The top-level README should stay focused on the existing high-level workflow list. Keep the NeMo Gym distillation usage details in the algorithm and integration docs instead. Constraint: User requested no top-level README update Confidence: high Scope-risk: narrow Tested: git diff --check -- README.md docs/about/algorithms/on-policy-distillation.md docs/design-docs/nemo-gym-integration.md Signed-off-by: Meng Xin <mxin@nvidia.com>
…Mo-Gym run_distillation_nemo_gym.py mirrored the pre-#2741 GRPO gym-init flow and called nemo_gym.health_check.remote() to block until the gym is ready. #2741 ("overlap NeMo Gym init with vLLM init") removed the NemoGym actor's health_check() method in favor of _spinup(), so the call raised AttributeError: 'ActorHandle' object has no attribute 'health_check' and broke gb200_L1_Functional_Tests_Gym. create_env() still creates the actor correctly; only the blocking spin-up call needed renaming to match the new API (grpo.setup() uses the same ray.get(actor._spinup.remote())). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
…spinup Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
Contributor
Author
|
/ok to test af236c8 |
terrykong
reviewed
Jun 11, 2026
terrykong
left a comment
terrykong
left a comment
Collaborator
There was a problem hiding this comment.
lgtm, but can you change the distillation nemo gym test to "fast"
so we can catch this. i think this was missed in the original pr
Mark the distillation NeMo-Gym functional test as `fast` in L1_Functional_Tests_Gym.sh so it also runs under CI:Lfast (FAST=1), not only the full L1 tier. Catches regressions like the health_check->_spinup breakage earlier, on more PRs. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
yfw
added
CI:Lfast
Contributor
Author
|
/ok to test 7616c64 |
terrykong
approved these changes
Jun 11, 2026
ashors1
pushed a commit
that referenced
this pull request
Jun 27, 2026
…Mo-Gym (#2786) Signed-off-by: Meng Xin <mxin@nvidia.com> Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com> Co-authored-by: Meng Xin <mxin@nvidia.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.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 ?
Fix distillation startup. Use _spinup instead of health_check
Issues
Closes #2784
Usage
# Add a code snippet demonstrating how to use thisBefore your PR is "Ready for review"
Pre checks:
Additional Information