fix(sc): retain rollout metrics through training by yfw · Pull Request #3880 · NVIDIA-NeMo/RL

yfw · 2026-08-27T19:51:37Z

Summary

Ports the rollout-metric aggregation behavior from #2175 to the SingleController async/TransferQueue path.

retain per-group rollout metrics when prompt groups are committed to TQReplayBuffer
carry metrics only from prompt groups selected by the sampler for the current training step
aggregate selected metrics with the shared aggregate_rollout_metrics semantic rules and log them with the corresponding train step
honor the existing W&B opt-in before retaining NeMo Gym full-result tables, keeping disabled payloads out of replay checkpoints

No new configuration or migration changes are required.

Validation

focused SingleController, replay-buffer, sampler, rollout-manager, setup, and aggregation tests: 151 passed
ruff check and ruff format --check: passed
pyrefly check on changed production files: passed

Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

copy-pr-bot · 2026-08-27T19:51:41Z

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.

yfw

Rollout metrics reaching the SingleController train logs is a real gap: on main the SC path computed record.rollout_metrics and discarded it, with no other route to those numbers. The implementation matches its nearest analogs (grpo.py:4845-4849, ppo.py:2432) almost line for line and reuses the shared aggregate_rollout_metrics rather than reimplementing its min/max/sum/histogram semantics, which is the right call. The re-attach in _finalize_selection is genuinely necessary — concat keeps only the first meta's extra_info — and the pop targets a fresh concat-produced meta, so it cannot corrupt buffered state. Multi-chunk accumulation is correctly scoped, and resume from a pre-PR checkpoint degrades gracefully.

One inline comment is a real bug on an adjacent line this PR activates; the other two are low-priority test pins.

Minor, no action required: the TODO at single_controller.py:1401-1403 still lists rollout_metrics as outstanding — this PR implements it. The histogram log item in that same TODO is a different thing (the explicit log_histogram calls that grpo.py:2562-2582 makes for baseline/advantage) and is still outstanding.

Generated by Claude Code

Honor the existing W&B opt-in before retaining NeMo Gym full-result tables in SingleController rollout metrics. Add coverage for default-off behavior, setup propagation, checkpoint compatibility, and removal before DP sharding. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

yfw · 2026-08-28T07:45:27Z

/ok to test 14a80c8

Keep the constructor-bypassing NeMo Gym redispatch fixture aligned with the production default so rollout metric computation can exercise the retry path. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

yfw · 2026-08-28T18:22:11Z

/ok to test 3be567b

Keep the partial checkpoint setup config aligned with the production logger shape so full-result table gating can be evaluated during resume tests. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

yfw · 2026-08-29T00:20:24Z

/ok to test 103fef5

Resolve the SingleController setup import overlap while preserving rollout metric logging and warm-start checkpoint validation. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

yfw · 2026-08-29T01:03:37Z

/ok to test 38ee49e

Keep the partial PPO setup config aligned with the production logger shape so rollout table gating can be evaluated across critic setup tests. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

Resolve SingleController rollout metric retention against native TQ checkpoint recovery while preserving both behaviors. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

Preserve rollout metric retention and optional full-result table logging while incorporating reward penalties, reference-logprob gating, and the latest SingleController changes. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

yfw · 2026-09-01T17:12:29Z

/ok to test 8491d65

Treat generation worker IDs as string metric labels so retained per-worker token counts remain valid native TQ replay metadata. Cover both metric production and checkpoint round trips. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

mehraakash

The full-result-table opt-in is useful, but the branch is no longer safe to merge as-is.

The rollout-metric retention and aggregation portion has already landed on main through #3883, together with speculative-decoding metrics and replay-manifest handling. This PR now conflicts with current main, and taking the PR versions during conflict resolution could regress those newer changes.

Please rebase and reduce this PR to the remaining unique change:

propagate log_nemo_gym_full_result_tables into the SingleController rollout manager;
construct full-result tables only when enabled;
retain focused default-off, opt-in, setup-propagation, and checkpoint-serialization tests.

The current SingleController functional failure is a real checkpointing regression caused by hashing rollout_metrics, so I don’t think this is ready until that path is fixed and green.

Keep the full-result-table opt-in on top of main rollout metrics and replay-manifest handling. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

yfw · 2026-09-03T06:10:13Z

/ok to test 782dff6

fix(sc): retain rollout metrics through training

a57599f

Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

yfw requested review from a team as code owners August 27, 2026 19:51

yfw commented Aug 27, 2026

View reviewed changes

Comment thread nemo_rl/algorithms/async_utils/replay_buffer.py Outdated

Comment thread tests/unit/single_controller/test_sampler_interface.py Outdated

Comment thread tests/unit/single_controller/test_single_controller_actor.py

yfw requested a review from a team as a code owner August 27, 2026 22:59

yfw added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Aug 28, 2026

test(sc): initialize full-result logging gate

3be567b

Keep the constructor-bypassing NeMo Gym redispatch fixture aligned with the production default so rollout metric computation can exercise the retry path. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

test(sc): initialize checkpoint logger fixture

103fef5

Keep the partial checkpoint setup config aligned with the production logger shape so full-result table gating can be evaluated during resume tests. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

Merge origin/main into yifu/v2_rollout_metrics

38ee49e

Resolve the SingleController setup import overlap while preserving rollout metric logging and warm-start checkpoint validation. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

yfw added 3 commits August 29, 2026 16:23

test(sc): initialize PPO logger fixture

df2d93f

Keep the partial PPO setup config aligned with the production logger shape so rollout table gating can be evaluated across critic setup tests. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

Merge origin/main into yifu/v2_rollout_metrics

fc4c88a

Resolve SingleController rollout metric retention against native TQ checkpoint recovery while preserving both behaviors. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

Merge origin/main into yifu/v2_rollout_metrics

8491d65

Preserve rollout metric retention and optional full-result table logging while incorporating reward penalties, reference-logprob gating, and the latest SingleController changes. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

fix(sc): checkpoint worker rollout metrics

ad996d1

Treat generation worker IDs as string metric labels so retained per-worker token counts remain valid native TQ replay metadata. Cover both metric production and checkpoint round trips. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

mehraakash requested changes Sep 2, 2026

View reviewed changes

Comment thread nemo_rl/algorithms/async_utils/replay_buffer.py

Comment thread nemo_rl/experience/rollout_manager.py

yfw added 2 commits September 2, 2026 21:26

Merge origin/main into yifu/v2_rollout_metrics

b49bd8f

Keep the full-result-table opt-in on top of main rollout metrics and replay-manifest handling. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

test(sc): checkpoint opt-in rollout result tables

782dff6

Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

mehraakash approved these changes Sep 4, 2026

View reviewed changes

yfw merged commit 9b75ab7 into main Sep 4, 2026
86 checks passed

yfw deleted the yifu/v2_rollout_metrics branch September 4, 2026 18:59

Navigation Menu

Uh oh!

fix(sc): retain rollout metrics through training - #3880

fix(sc): retain rollout metrics through training#3880
yfw merged 11 commits into
mainfrom
yifu/v2_rollout_metrics

yfw commented Aug 27, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Aug 27, 2026

Uh oh!

yfw left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yfw commented Aug 28, 2026

Uh oh!

yfw commented Aug 28, 2026

Uh oh!

yfw commented Aug 29, 2026

Uh oh!

yfw commented Aug 29, 2026

Uh oh!

yfw commented Sep 1, 2026

Uh oh!

mehraakash left a comment

Uh oh!

Uh oh!

Uh oh!

yfw commented Sep 3, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

yfw commented Aug 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Validation

Uh oh!

copy-pr-bot Bot commented Aug 27, 2026

Uh oh!

yfw left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yfw commented Aug 28, 2026

Uh oh!

yfw commented Aug 28, 2026

Uh oh!

yfw commented Aug 29, 2026

Uh oh!

yfw commented Aug 29, 2026

Uh oh!

yfw commented Sep 1, 2026

Uh oh!

mehraakash left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

yfw commented Sep 3, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

yfw commented Aug 27, 2026 •

edited

Loading