refactor(dynamo): move native token response handling to Gym - #3763
refactor(dynamo): move native token response handling to Gym#3763jthomson04 wants to merge 1 commit into
Conversation
|
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. |
github-actions
Bot
added
the
Documentation
|
/ok to test 05a5360 |
05a5360 to
245575e
Compare
|
/ok to test 245575e |
jthomson04
left a comment
jthomson04
left a comment
There was a problem hiding this comment.
Team review of the Dynamo→Gym response-handling move: 6 agents (RL expert, Gym submodule expert, bug finder, design reviewer, test agent, devil's advocate), 18 findings raised, 7 surviving after an adversarial pass.
The refactor itself is a clear win and worth saying so. Gym's replacement isn't just a relocation of _validate_engine_data — it cross-checks the Dynamo bundle against message-level metadata and against vLLM's own response token IDs, neither of which the wrapper ever did. Concrete payoff: adding routed_experts to the native bundle used to mean editing two repos; now it's one entry in _DYNAMO_ENGINE_DATA_FIELD_MAPPING.
The one item I'd genuinely act on before merge is the test-discrimination gap — the suite currently passes with the deleted derivation restored. Everything else is low-severity or documentation.
pre-commit run --all-files is clean. Two items were deliberately not filed here: the missing-engine_data fallback belongs on NeMo-Gym PR 1784, and the "replace the overloaded engine_data flag with a named derive_required_prefix_token_ids boolean" idea belongs in a tracking issue against Gym.
Three findings were dropped by the adversarial pass and are worth naming so they don't get re-raised later: an assert on the hardcoded vllm_model key (the fix would break grpo_nanov3.yaml, and vllm_model_for_training.yaml pins both names anyway), a comment about OmegaConf list replacement (unreachable), and routed_experts not being stripped (router replay is rejected outright on the Dynamo backend).
Generated by Claude Code
| body, | ||
| tokenizer=_Tokenizer(), | ||
| exclude_tools_when_tool_choice_none=True, | ||
| ) |
There was a problem hiding this comment.
tests/unit/models/generation/test_dynamo_token_wrapper.py:392
2 action items.
TL;DR — the suite can no longer detect a revert of this PR's central change: patch the wrapper to re-derive required_prefix_token_ids the deleted way, keep both new raises, and all 14 tests still pass.
PR-introduced. The behavior change here is "stop deriving the prefix, use Gym's". Every test that supplies required_prefix_token_ids sets it to exactly prompt_token_ids + generation_token_ids — precisely what the deleted _derive_required_prefix_token_ids returned:
preserves_prior_prefix:340—[10, 31, 32, 2]vs the message's[10]+[31, 32, 2]public_qwen3_multiturn_prefix_splice_parity:300— literallyfirst_prompt + generation_token_idsnormalizes_prior_tool_arguments:448—[10,777,2,40]+[31,32,2]
So the tests pin the raises, not the value path.
AI-1 — add a discriminating test
Deliberately not asking you to change preserves_prior_prefix's value: its name should keep describing the realistic agreeing-payload case, and repointing it would delete the only test asserting that shape.
AI-2 — cover the inverse guard
grep -rn "requires a tokenized assistant message" tests/ returns nothing — the raise at token_wrapper.py:267-270 is untested.
Both tests below verified passing locally (14 passed → 16 passed):
| ) | |
| ) | |
| def test_prepare_dynamo_chat_completion_request_prefers_gym_prefix_over_message_ids() -> ( | |
| None | |
| ): | |
| tokenizer = _Tokenizer() | |
| body = { | |
| "model": "dummy-model", | |
| # Gym is the source of truth: this must win over the token IDs the | |
| # assistant message carries, which the wrapper no longer derives from. | |
| "required_prefix_token_ids": [10, 31, 32, 2], | |
| "messages": [ | |
| {"role": "user", "content": "hello"}, | |
| { | |
| "role": "assistant", | |
| "content": "first", | |
| "prompt_token_ids": [10], | |
| "generation_token_ids": [55, 56, 2], | |
| "generation_log_probs": [-0.1, -0.2, -0.3], | |
| }, | |
| {"role": "user", "content": "next"}, | |
| ], | |
| } | |
| prepared = prepare_dynamo_chat_completion_request( | |
| body, | |
| tokenizer=tokenizer, | |
| exclude_tools_when_tool_choice_none=True, | |
| ) | |
| assert prepared["nvext"]["token_data"] == [10, 31, 32, 2, 40, 99] | |
| def test_prepare_dynamo_chat_completion_request_rejects_prefix_without_tokenized_turn() -> ( | |
| None | |
| ): | |
| body = { | |
| "model": "dummy-model", | |
| "required_prefix_token_ids": [10, 31, 32, 2], | |
| "messages": [ | |
| {"role": "user", "content": "hello"}, | |
| {"role": "assistant", "content": "first"}, | |
| {"role": "user", "content": "next"}, | |
| ], | |
| } | |
| with pytest.raises( | |
| ValueError, match="required_prefix_token_ids requires a tokenized assistant" | |
| ): | |
| prepare_dynamo_chat_completion_request( | |
| body, | |
| tokenizer=_Tokenizer(), | |
| exclude_tools_when_tool_choice_none=True, | |
| ) |
| @@ -1 +1 @@ | |||
| Subproject commit c3bac96314a59f28b896f597eb9845d175bb0252 | |||
| Subproject commit fad304de9f915059277d76be5d30dfce43ad2cf6 | |||
There was a problem hiding this comment.
1 action item, 1 follow-up.
TL;DR — merging this pin as-is makes the next, unrelated Gym-bump PR fail CI with Commits have DIVERGED, and the PR body's claim that CI is the end-to-end gate is not true — no CI job exercises the Dynamo↔Gym round trip.
Not restating your repin caveat — these are two things it doesn't cover.
AI-1 — the "end-to-end validation gate" claim is false
The only Dynamo CI entry is L1_Functional_Tests_Dynamo.sh → grpo_dynamo.sh:44 → grpo_math_1B_dynamo.yaml, which has zero nemo_gym references — so setup_nemo_gym_config is never called and the wrapper's Gym path never runs. The only dynamo+gym config is grpo-nanov3-30ba3b-3n8g-megatron-dynamo-swe1.yaml, a 3-node/8-GPU recipe that isn't a PR gate. Unit tests are the only coverage.
Action: soften that line in the PR body to "unit tests only; the managed Dynamo+Gym smoke is run manually", or add a dynamo variant of the gym functional test.
Follow-up — the pin will break someone else's PR
git merge-base --is-ancestor fad304de9 origin/main exits 1, and git branch -r --contains fad304de9 returns only fork/codex/dynamo-native-token-transport. Gym main is linear with (#NNNN) squash subjects, so PR 1784 will land under a different SHA. RL main would then pin a commit that is neither ancestor nor descendant of Gym main, and the next Gym bump hits Case 4 at _submodule_check.yml:143 — Commits have DIVERGED from a common ancestor, failing on an author who had nothing to do with it. The existing fast-forward check can't catch this (it only tests ancestry against the PR's own Gym ref, which is why it's green here). Just worth being deliberate that the repin happens before merge, not after.
| caller token IDs. | ||
| uses the policy tokenizer, replaces rendered multi-turn prefixes with Gym's | ||
| exact `required_prefix_token_ids`, and sends `nvext.token_data`. For the Dynamo | ||
| backend, NeMo-RL configures Gym to request `nvext.engine_data`; this request |
There was a problem hiding this comment.
docs/design-docs/dynamo-integration.md:28
1 action item. You updated this design doc for the ownership move; its sibling guide is now false and wasn't.
docs/guides/dynamo-generation.md:60-61 is a Dynamo-upgrade checklist telling the upgrader to reverify the response keys against the wrapper:
nemo_rl/models/generation/dynamo/token_wrapper.py: the real Dynamo response keysnvext.engine_data.{prompt_token_ids,completion_token_ids,completion_logprobs}; reverify them against real Dynamo output …
After this PR, grep -n engine_data nemo_rl/models/generation/dynamo/token_wrapper.py returns exactly one hit — line 322 — and it's the request-side extra_fields append. The reader moved to Gym's _DYNAMO_ENGINE_DATA_FIELD_MAPPING. A future upgrade following this checklist would check the wrong side and miss a wire break — which matters because line 66 of that same guide already documents a real case where it breaks ("Dynamo 1.3.0's legacy tool jail removes nvext.engine_data").
Action: in docs/guides/dynamo-generation.md, repoint that bullet at Gym's app.py mapping, noting the wrapper only requests the field.
|
|
||
| setup_nemo_gym_config(config, tokenizer=None) | ||
|
|
||
| assert config.env["nemo_gym"] == {"port_range_low": 5000} |
There was a problem hiding this comment.
tests/unit/environments/test_nemo_gym.py:1010
2 action items, 1 nit.
Three branches of the new setup_nemo_gym_config block are uncovered. Both tests below verified passing locally.
AI-1 — the extra_fields non-list raise
grep -rn "must be a list" tests/ returns nothing; nemo_gym.py:1133-1136 is untested.
AI-2 — the from-scratch setdefault chain
Both new gym tests pre-populate env["nemo_gym"], so the 5-deep chain at nemo_gym.py:1124-1132 is never built from empty. That's the production-representative path — a minimal dynamo config has no env.nemo_gym block at all — and it's the exact key path Gym reads, so a typo anywhere in it is currently only caught at runtime.
Nit
Lines 987-988 call setup_nemo_gym_config twice with no explanation. It does genuinely prove idempotency (drop the if "engine_data" not in extra_fields guard and the assert sees a duplicate), but a one-line comment would save the next reader the detour — e.g. # Called twice: the dedup guard must make a re-run a no-op.
One thing worth knowing about the assertion at line 996: ["timing", "engine_data"] proves list preservation only for entries seeded on the NeMo-RL side. This dict is merged last into Gym's global config and OmegaConf replaces lists, so an extra_fields entry declared in a Gym YAML would be dropped instead. No Gym YAML declares one today, so there's nothing to fix — just don't read that assertion as broader than it is.
| assert config.env["nemo_gym"] == {"port_range_low": 5000} | |
| assert config.env["nemo_gym"] == {"port_range_low": 5000} | |
| def test_setup_nemo_gym_config_rejects_non_list_dynamo_extra_fields() -> None: | |
| config = SimpleNamespace( | |
| policy={"generation": {"backend": "dynamo", "vllm_cfg": {}}}, | |
| env={ | |
| "nemo_gym": { | |
| "policy_model": { | |
| "responses_api_models": { | |
| "vllm_model": { | |
| "extra_body": {"nvext": {"extra_fields": "engine_data"}} | |
| } | |
| } | |
| } | |
| } | |
| }, | |
| ) | |
| with pytest.raises(ValueError, match="nvext.extra_fields must be a list"): | |
| setup_nemo_gym_config(config, tokenizer=None) | |
| def test_setup_nemo_gym_config_builds_dynamo_engine_data_path_from_scratch() -> None: | |
| config = SimpleNamespace( | |
| policy={"generation": {"backend": "dynamo", "vllm_cfg": {}}}, | |
| env={}, | |
| ) | |
| setup_nemo_gym_config(config, tokenizer=None) | |
| assert config.env == { | |
| "nemo_gym": { | |
| "policy_model": { | |
| "responses_api_models": { | |
| "vllm_model": { | |
| "extra_body": {"nvext": {"extra_fields": ["engine_data"]}} | |
| } | |
| } | |
| } | |
| } | |
| } |
| if assistant_index is not None and required_prefix_token_ids is None: | ||
| raise ValueError( | ||
| "A tokenized assistant message requires required_prefix_token_ids." | ||
| ) |
There was a problem hiding this comment.
nemo_rl/models/generation/dynamo/token_wrapper.py:263-266
2 action items, both in the one suggestion below. Low severity.
TL;DR — required_prefix_token_ids: [] slips past this guard and silently ships an un-spliced prompt, and when the guard does fire its message is the only diagnostic the operator ever sees, because Gym launders the 400 into a 500 that NeMo-RL then retries as an infra failure.
AI-1 — is None lets [] through
[] is not None, so an empty list takes the "present" branch, replace_prefix_tokens short-circuits on if not model_prefix_token_ids, and the un-spliced re-tokenized prompt goes to Dynamo with no error. Verified: [] → token_data [10, 3, 2, 40, 99] where the spliced value is [10, 31, 32, 2, 40, 99].
Reachability is essentially zero — Gym can only emit [] if both prompt_token_ids and generation_token_ids are empty, which no engine produces — so this is a nit, not a live bug. It's worth doing only because the fix is one token, blocks nothing legitimate (an empty prefix with a tokenized assistant turn is already nonsense), and keeps this guard consistent with _latest_tokenized_assistant_index, which uses the same is not None test on the message side.
AI-2 — make the message name the two knobs
The raise fires exactly when Gym attaches message-level token metadata but derivation is gated off — the and at app.py:609. That's a permanent config error, but it doesn't reach the operator as one:
token_wrapper.py:406-407wraps thisValueErrorasHTTPException(400).- 400 ∉
RETRY_ERROR_CODES— correct, no retry at that hop. - But
server_utils.py:593re-emits it as a 500 carrying the text. failures.py:202—status >= 500→RolloutInfraFailure, so it retries across shards and finally dies asexhausted its data budget.
Steps 3-4 are pre-existing and out of scope; this string is the only part of the chain you control. Keeping it under ~200 chars matters — Gym's retry log truncates at content[:200].
One way to do it (verified: 14 passed, and [] now raises):
| if assistant_index is not None and required_prefix_token_ids is None: | |
| raise ValueError( | |
| "A tokenized assistant message requires required_prefix_token_ids." | |
| ) | |
| if assistant_index is not None and not required_prefix_token_ids: | |
| raise ValueError( | |
| "A tokenized assistant message requires required_prefix_token_ids. " | |
| "NeMo-Gym derives it only when the model server sets " | |
| "return_token_id_information: true and nvext.extra_fields " | |
| "contains 'engine_data'." | |
| ) |
| generation_config["stop_strings"] = None | ||
| generation_config["stop_token_ids"] = None | ||
|
|
||
| if generation_config["backend"] == "dynamo": |
There was a problem hiding this comment.
nemo_rl/environments/nemo_gym.py:1123
2 action items. Nit-level.
Two functions this PR gives a new cross-repo contract to, neither of which documents it. Asking for 1-2 sentences each, not full Google docstrings — the module-level helpers in token_wrapper.py all carry one-liners and I don't think that should change.
AI-1 — setup_nemo_gym_config
It's public and now has a backend-conditional side effect on config.env that a caller can't guess. Its immediate neighbour spinup_nemo_gym_actor has a full docstring, so the local convention supports one here. Worth stating: this appends "engine_data" to the Gym policy model's extra_body.nvext.extra_fields, which is how Gym decides to derive required_prefix_token_ids, and it's idempotent.
AI-2 — prepare_dynamo_chat_completion_request
The one-line docstring is unchanged, but required_prefix_token_ids went from ignored-and-popped to load-bearing with two new raise paths. A sentence saying Gym owns prefix derivation and that the two signals must agree would carry it.
| .setdefault("responses_api_models", {}) | ||
| .setdefault("vllm_model", {}) | ||
| ) | ||
| extra_body = model_cfg.setdefault("extra_body", {}) |
There was a problem hiding this comment.
nemo_rl/environments/nemo_gym.py:1130
1 action item. Nit — no shipped config triggers it.
The block type-guards extra_fields at line 1133 but not the two levels above it, so a None at either level hits setdefault on NoneType. Verified against the real function:
env.nemo_gym…vllm_model = {"extra_body": None}
-> AttributeError: 'NoneType' object has no attribute 'setdefault'
env.nemo_gym…vllm_model = {"extra_body": {"nvext": None}}
-> AttributeError: 'NoneType' object has no attribute 'setdefault'
So extra_fields: null raises your clean named ValueError while extra_body: null crashes with a raw traceback. null is the idiomatic way to spell an empty override block in these configs, so someone hand-writing one in a NeMo-RL env.nemo_gym override is plausible — though nothing in examples/ does it today.
This is also inconsistent with the wrapper's own handling, which explicitly tolerates a None nvext at token_wrapper.py:311-315.
Action: either coerce None → {}, or extend the guard to reject a non-dict extra_body/nvext with the same named-ValueError style you already use one line below.
Signed-off-by: jthomson04 <jwillthomson19@gmail.com>
245575e to
d3a007e
Compare
|
/ok to test d3a007e |
1 participant
Summary
Move native Dynamo response-token interpretation from the NeMo-RL token wrapper into NeMo-Gym.
nvext.engine_dataonly when the generation backend is Dynamo.extra_body,nvext, andextra_fieldsvalues while adding the request field once.nvext.token_datarequest preparation in NeMo-RL.Dependency
This PR depends on NeMo-Gym PR 1784. The Gym submodule is pinned to its current exact head,
fad304de9f915059277d76be5d30dfce43ad2cf6, so CI tests both sides together.Before this PR merges, the submodule must be advanced to the exact Gym
maincommit that contains PR 1784.Compatibility
nvext.token_dataandnvext.extra_fields=["engine_data"].uv.lockis unchanged.Validation
uv run --no-sync pytest tests/unit/models/generation/test_dynamo_token_wrapper.py tests/unit/environments/test_nemo_gym.py::test_setup_nemo_gym_config_requests_dynamo_engine_data tests/unit/environments/test_nemo_gym.py::test_setup_nemo_gym_config_does_not_add_dynamo_fields_for_vllm(16 passed)uv run --no-sync pre-commit run --files docs/design-docs/dynamo-integration.md nemo_rl/environments/nemo_gym.py nemo_rl/models/generation/dynamo/token_wrapper.py tests/unit/environments/test_nemo_gym.py tests/unit/models/generation/test_dynamo_token_wrapper.pyThe NeMo-RL CI run on this PR is the end-to-end validation gate.