Skip to content

feat: support R3 async RL without TQ - #2908

Merged
terrykong merged 1 commit into
NVIDIA-NeMo:mainfrom
zyzhou5:r3-async-no-tq
Jun 25, 2026
Merged

feat: support R3 async RL without TQ#2908
terrykong merged 1 commit into
NVIDIA-NeMo:mainfrom
zyzhou5:r3-async-no-tq

Conversation

@zyzhou5

@zyzhou5 zyzhou5 commented Jun 23, 2026

Copy link
Copy Markdown
Contributor

What does this PR do ?

  • Preserve R3/router-replay routed_experts when building policy input batches in GRPO training paths.
  • Enable the async no-TQ path to consume rollout-recorded routes for prev-logprob and training.
  • Add a clear guard for the unsupported async+TQ combination.

@zyzhou5
zyzhou5 requested review from a team as code owners June 23, 2026 23:42
@copy-pr-bot

copy-pr-bot Bot commented Jun 23, 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 Jun 23, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test c4305cd

@zyzhou5

zyzhou5 commented Jun 24, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 9f2afe8

@zyzhou5 zyzhou5 added the CI:L1 Run doctests, unit tests, and functional tests label Jun 24, 2026
@zyzhou5

zyzhou5 commented Jun 24, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 9f2afe8

@zyzhou5

zyzhou5 commented Jun 24, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 40fa5d3

ZhiyuLi-Nvidia
ZhiyuLi-Nvidia previously approved these changes Jun 24, 2026

@ZhiyuLi-Nvidia ZhiyuLi-Nvidia 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.

Awesome change. LGTM!

@zyzhou5

zyzhou5 commented Jun 24, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 7a1c5ce

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

zyzhou5 commented Jun 24, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 3b617bf

@terrykong
terrykong merged commit 3cb22af into NVIDIA-NeMo:main Jun 25, 2026
92 checks passed
ashors1 pushed a commit that referenced this pull request Jun 27, 2026
Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
Signed-off-by: Anna Shors <ashors@nvidia.com>
@qiaochuz-nv

Copy link
Copy Markdown
Contributor

Test coverage added:

1. test_grpo_router_replay_async_dp_off_daily_pr
   PR covered: NVIDIA-NeMo/RL #2908

   Functions / paths covered:
   - async GRPO R3 path without TransferQueue
   - data_plane.enabled=false
   - _build_async_grpo_train_data
     - verifies routed_experts survives async producer/consumer batching
     - verifies train/logprob data carries R3 route tensors
   - Megatron Qwen3-30B-A3B MoE policy path
   - TP4 / PP2 / CP2 / EP4 topology
   - non-colocated vLLM generation
   - RouterReplay replay in:
     - prev-logprob
     - train
   - CP2 token identity / zigzag route-sharding correctness
   - forward verifier matches_expected=true

   Status: PASS
   Evidence:
   - run_id: 20260706T185536Z-0496ec2b
   - SLURM job: 5579660
   - SLURM: COMPLETED|0:0
   - elapsed: 00:08:52
   - nodes: eos[0516,0518,0520]
   - train/loss steps: 2
   - R3 trace records: 21632
   - replay assignments:
     - prev-logprob: 3072
     - train: 3072
   - replay forward:
     - prev-logprob: 3072
     - train: 3072
   - replay backward train: 3072
   - forward verifier:
     - prev-logprob: 3072
     - train: 3072
   - CP token identity rows: 47168

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

Labels

CI:L1 Run doctests, unit tests, and functional tests QA:Verified Ready for QA

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants