Skip to content

feat: accept opaque string envelopes for routed_experts#2089

Merged
ananthsub merged 1 commit into
NVIDIA-NeMo:mainfrom
zyzhou5:zezhou/routed-experts-b64-envelope
Jul 21, 2026
Merged

feat: accept opaque string envelopes for routed_experts#2089
ananthsub merged 1 commit into
NVIDIA-NeMo:mainfrom
zyzhou5:zezhou/routed-experts-b64-envelope

Conversation

@zyzhou5

@zyzhou5 zyzhou5 commented Jul 20, 2026

Copy link
Copy Markdown
Contributor

Training frameworks attach per-token routed-expert indices ([tokens, num_moe_layers, topk]) to chat responses for MoE router replay. As nested JSON int lists these are multi-MB payloads that every gym hop re-parses, re-validates, and re-serializes at ~1s of single-threaded CPU each, throttling rollout throughput.

Widen RoutedExperts to Union[str, List[List[List[int]]]] so a single opaque envelope string (e.g. NeMo-RL's nrlre1:::) passes through validation in O(1). Gym never inspects the value; the one non-passthrough site (harbor nemo_gym_llm _extract_routed_experts isinstance-list guard, which would silently drop a string) is widened to match. Nested-list payloads remain fully supported.

@copy-pr-bot

copy-pr-bot Bot commented Jul 20, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@zyzhou5

zyzhou5 commented Jul 20, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test c4d96c4

@ZhiyuLi-Nvidia ZhiyuLi-Nvidia left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you @zyzhou5, LGTM.

Training frameworks attach per-token routed-expert indices
([tokens, num_moe_layers, topk]) to chat responses for MoE router
replay. As nested JSON int lists these are multi-MB payloads
(~11.7M ints at 30k-token contexts) that every gym hop re-parses,
re-validates, and re-serializes at ~1s of single-threaded CPU each,
throttling rollout throughput.

Widen RoutedExperts to Union[str, List[List[List[int]]]] so a single
opaque envelope string (e.g. NeMo-RL's nrlre1:<dtype>:<SxLxK>:<base64>)
passes through validation in O(1). Gym never inspects the value; the
one non-passthrough site (harbor nemo_gym_llm _extract_routed_experts
isinstance-list guard, which would silently drop a string) is widened
to match. Nested-list payloads remain fully supported.

Companion to the NeMo-RL change that emits the envelope; measured
there: gym rollout time 132.2s -> 35.4s per step at 15-28k-token
prompts with replay fidelity unchanged.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
@ananthsub
ananthsub force-pushed the zezhou/routed-experts-b64-envelope branch from c4d96c4 to bb8e661 Compare July 21, 2026 17:16
@ananthsub
ananthsub merged commit c40117a into NVIDIA-NeMo:main Jul 21, 2026
15 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants