Skip to content

feat: R3 gym notq router replay - #2915

Merged
yuki-97 merged 13 commits into
NVIDIA-NeMo:mainfrom
zyzhou5:r3-gym-notq-router-replay
Jul 1, 2026
Merged

feat: R3 gym notq router replay#2915
yuki-97 merged 13 commits into
NVIDIA-NeMo:mainfrom
zyzhou5:r3-gym-notq-router-replay

Conversation

@zyzhou5

@zyzhou5 zyzhou5 commented Jun 24, 2026

Copy link
Copy Markdown
Contributor

What does this PR do ?

Adds no-TQ async Gym path support for R3 router replay.

When policy.router_replay.enabled=true, async Gym rollouts now request routed expert indices from Gym, validate that they are present, and convert them into Nemo-RL message logs. The routed expert indices are sliced across prompt and generation tokens for each Gym model call so the existing R3 training path can consume them.

This also updates the async OpenAI-compatible vLLM endpoint to return routed expert indices in chat response messages when requested.

@zyzhou5
zyzhou5 requested review from a team as code owners June 24, 2026 19:11
@copy-pr-bot

copy-pr-bot Bot commented Jun 24, 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 added the CI:L1 Run doctests, unit tests, and functional tests label Jun 24, 2026
@zyzhou5 zyzhou5 changed the title R3 gym notq router replay feat: R3 gym notq router replay Jun 24, 2026
@zyzhou5

zyzhou5 commented Jun 24, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 763fb26

@zyzhou5

zyzhou5 commented Jun 24, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test c289b97

@zyzhou5

zyzhou5 commented Jun 24, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test df70c93

@zyzhou5

zyzhou5 commented Jun 24, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test f594c2d

zyzhou5 added 4 commits June 29, 2026 09:47
Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
@zyzhou5 zyzhou5 added CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) and removed CI:L1 Run doctests, unit tests, and functional tests labels Jun 29, 2026
@zyzhou5
zyzhou5 force-pushed the r3-gym-notq-router-replay branch from 2ff1ffc to c25a0bf Compare June 29, 2026 17:00
@zyzhou5

zyzhou5 commented Jun 29, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test c25a0bf

@yuki-97 yuki-97 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

thanks @zyzhou5 , changes LGTM, some comments on test things.

can we add several nightly tests to guard r3 working correctly? we currently seem only cover grpo TQ path in examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3.yaml, could you add the below?

  1. grpo path: #2849
  2. async grpo path: related #2908
  3. gym path: this PR

also maybe good to rename examples/configs/recipes/llm/grpo-qwen3-30ba3b-8n8g-megatron-cp2-r3.yaml to keep the same name style like the other TQ tests.

Comment thread tests/unit/models/generation/test_vllm_utils.py
zyzhou5 added 5 commits June 30, 2026 09:43
Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
@zyzhou5

zyzhou5 commented Jun 30, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 25c2f00

ZhiyuLi-Nvidia
ZhiyuLi-Nvidia previously approved these changes Jun 30, 2026
Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
@zyzhou5

zyzhou5 commented Jun 30, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 6b8df98

Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
@zyzhou5

zyzhou5 commented Jun 30, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 43d34b5

@yuki-97

yuki-97 commented Jul 1, 2026

Copy link
Copy Markdown
Contributor

/ok to test 58b6af1

@yuki-97
yuki-97 enabled auto-merge (squash) July 1, 2026 05:30

@yuki-97 yuki-97 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

LGTM, thanks @zyzhou5

@yuki-97
yuki-97 merged commit af207e8 into NVIDIA-NeMo:main Jul 1, 2026
78 checks passed
@qiaochuz-nv

Copy link
Copy Markdown
Contributor
  1. test_grpo_qwen3_30ba3b_megatron_cp2_r3_async_gym_e2e
    PR covered: NVIDIA-NeMo/RL feat: R3 gym notq router replay #2915

    Functions / paths covered:

    • R3 router replay through NeMo-Gym async generation
    • examples/nemo_gym/run_grpo_nemo_gym.py
    • Gym/OpenAI-compatible server round-trip preserving routed_experts
    • async Gym rollout path
    • Megatron Qwen3-30B-A3B MoE policy path
    • CP2 route sharding
    • non-colocated vLLM/Gym generation
    • train-stage RouterReplay assignment/action
    • forward verifier matches_expected=true
    • CP token identity verification
    • offline SWE1 dataset staging from HF cache
    • CI-scaled 3-node / 24-GPU smoke for upstream 16-node recipe

    Status: PASS
    Evidence:

    • run_id: 20260707T093631Z-6495538f
    • SLURM job: 5584814
    • SLURM: COMPLETED|0:0
    • elapsed: 00:09:14
    • nodes: eos[0143,0192,0255]
    • train/loss steps: 2
    • median(train/token_mult_prob_error): 1.015147626399994
    • R3 trace records: 24640
    • train replay assignments: 6144
    • train replay forward: 6144
    • train replay backward: 6144
    • forward verifier: 6144
    • CP token identity rows: 193480

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) QA:Verified r0.7.0 Ready for QA

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants