refactor(nemo-gym): centralize the gym actor factory - #3367
Draft
ananthsub wants to merge 1 commit into
Draft
Conversation
The NemoGym actor was constructed in three places (grpo.setup, distillation.setup, single_controller_utils.setup) that had drifted apart, each forgetting a different field: - single_controller never passed use_fastokens or routed_experts_dtype - distillation never passed require_routed_experts - grpo and single_controller built runtime_env inline and set VIRTUAL_ENV to the *interpreter path* rather than the venv root Collapse all three onto spinup_nemo_gym_actor(), split out build_nemo_gym_config() so the config split (NeMo-RL knobs vs Gym's global config) is testable on its own, and use the shared make_actor_runtime_env() helper, which resolves VIRTUAL_ENV correctly. This makes the actor-creation path a single place to extend for multi-actor sharding. Distillation keeps require_routed_experts off: it never configures vLLM to emit routed experts, so requiring them would fail every rollout. The construction assertions that lived in the distillation tests move to the factory's own test module, where they cover all three callers at once. Signed-off-by: Ananth Subramaniam <ansubramania@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. |
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
1 participant
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 ?
Moves the construction of the
NemoGymRay actor into a single factory, so the three places that start one no longer each carry their own copy of the setup logic.nemo_rl/algorithms/grpo.py,nemo_rl/algorithms/distillation.py, andnemo_rl/algorithms/single_controller_utils/setup.pyeach built the actor's config inline. The three copies had already drifted: they split NeMo-RL-only keys out of the gym config differently, and they disagreed about the uv cache and venv directories and about router replay. Adding anything to the actor's setup meant finding and editing all three.This PR adds two functions to
nemo_rl/environments/nemo_gym.py:build_nemo_gym_config(...)assembles the config the actor receives: it separates the keys NeMo-RL owns from the ones it forwards to Gym, resolves the uv cache and venv directories, and applies the router replay settings.spinup_nemo_gym_actor(...)creates the actor with the right runtime environment and scheduling, then spins it up.All three call sites now call the factory with keyword arguments. Behavior is unchanged: this is a refactor, and the resulting actor config is the same one the previous inline code produced.
The detailed assertions about config assembly moved from
tests/unit/algorithms/test_distillation.py(where they only covered the distillation copy) totests/unit/environments/test_nemo_gym_utils.py, where they cover the shared factory.This is the first of eight stacked PRs adding actor sharding to NeMo-Gym. It is a prerequisite rather than a feature: the later PRs need one place to change.
Issues
None.
Usage
No user-facing change. Configs and command lines are unaffected.
Before your PR is "Ready for review"
Pre checks:
Additional Information
workplace_assistantrun on 2 GPUs. Validation accuracy is bit-identical to the pre-change baseline on that fixture, which is the signal that the refactor did not alter config assembly.