fix: eliminate JSON serialization bottleneck in r3 + NeMo-Gym rollouts - #3292
Merged
terrykong merged 4 commits intoJul 30, 2026
Merged
Conversation
…t lists At long context lengths the router-replay routes ([tokens, layers, topk]) attached to each chat response via .tolist() become multi-MB nested JSON int lists (~40MB at 30k tokens). Every hop on the NeMo Gym HTTP path (vLLM OpenAI server, gym model server, agent, resources server, gym actor) re-parses, pydantic-validates, and re-serializes them at ~1s of single-threaded CPU per hop, which throttles async-GRPO rollout production while generation GPUs sit idle. Encode the routes as a single self-describing base64 string (nrlre1:<dtype>:<SxLxK>:<base64>, preserving the resolved routed_experts_dtype) so intermediate hops handle one opaque object. The decoder accepts the legacy nested-list format for compatibility. Validated on a 16-node swe1 gym smoke (Qwen3-30B-A3B, prompts 15-28k tokens, 5 steps, seed 42): mean rollout time r3off 30.0s / r3on unpatched 132.2s (4.4x) / r3on patched 35.4s (1.18x), with per-step gen_kl_error identical to the unpatched R3 run (0.0005-0.0008). Requires the NeMo Gym RoutedExperts type alias to accept str (one-line change in nemo_gym/openai_utils.py, submitted separately). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
zyzhou5
added
the
CI:Lfast
zyzhou5
changed the title
Contributor
Author
|
/ok to test c7f9bcf |
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
Contributor
Author
|
/ok to test d6be365 |
Contributor
|
Awesome @zyzhou5. Thank you for driving this. Looks good to me except the above nit comments. |
Review suggestion from @ZhiyuLi-Nvidia: tobytes() materializes a full-size intermediate copy of the payload; memoryview on the contiguous array feeds b64encode directly. Byte-identical output. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
…tearray Review suggestion from @ZhiyuLi-Nvidia: decode into a bytearray and let torch.frombuffer share its memory instead of round-tripping through a read-only numpy view plus copy(). Copy count is unchanged in the common case (the bytearray construction is the one unavoidable copy), but this drops the numpy dependency, avoids a second copy when the target dtype differs from the wire dtype, and validates base64 strictly. Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
Contributor
Author
|
/ok to test 79fb96b |
Contributor
Author
|
@terrykong can you review this at your convenience? |
terrykong
approved these changes
Jul 30, 2026
pthombre
pushed a commit
that referenced
this pull request
Aug 12, 2026
Post-#3292 the worker attaches routed_experts as the base64 string envelope; torch.tensor() on it raised per call, degrading every capture to capture_failed coords and poisoning all rollouts (the capture-arm smoke's failure mode). Decode the envelope (wire dtype preserved). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Thombre <pthombre@oci-hsg-cs-001-vscode-02.cm.cluster> (cherry picked from commit 0482a980728e8c67576b1074557052acced6db95)
pthombre
pushed a commit
that referenced
this pull request
Aug 12, 2026
Post-#3292 the worker attaches routed_experts as the base64 string envelope; torch.tensor() on it raised per call, degrading every capture to capture_failed coords and poisoning all rollouts (the capture-arm smoke's failure mode). Decode the envelope (wire dtype preserved). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Thombre <pthombre@oci-hsg-cs-001-vscode-02.cm.cluster> (cherry picked from commit 0482a980728e8c67576b1074557052acced6db95)
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 ?
With router replay enabled on the NeMo-Gym path, per-sample routed experts (
[tokens, layers, topk])are attached to chat responses via
.tolist()— at ~30k-token contexts that is ~40 MB of nested JSONints per sample, re-parsed, pydantic-validated, and re-serialized at every gym HTTP hop (~1 s of
single-threaded CPU each). This throttles async-GRPO rollout production to a fraction of R3-off
throughput while generation GPUs sit idle.
This PR encodes the routes as a single self-describing base64 string
(
nrlre1:<dtype>:<SxLxK>:<base64>, preserving the resolvedrouted_experts_dtype), so intermediatehops handle one opaque object. The decoder still accepts the legacy nested-list format.
Validation
The performance improves significantly with this optimization
Report (Test 4 - Exp3): https://api.wandb.ai/links/nvidia-nemo-fw-public/m4onwlx8