Skip to content

feat(grpo): optional early stop at a validation metric threshold - #3404

Merged
yuki-97 merged 6 commits into
NVIDIA-NeMo:mainfrom
michal2409:split/stop-at-accuracy
Aug 3, 2026
Merged

feat(grpo): optional early stop at a validation metric threshold#3404
yuki-97 merged 6 commits into
NVIDIA-NeMo:mainfrom
michal2409:split/stop-at-accuracy

Conversation

@michal2409

@michal2409 michal2409 commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

Train-to-accuracy runs (benchmarking, recipe qualification, ablations) need training to stop as soon as the quality bar is met, rather than at a step count nobody can know in advance. This PR adds an optional early stop on a validation metric, via two config keys defaulted in the exemplar YAMLs:

  • grpo.stop_at_validation_metric (str | null) — which reported validation metric to compare; null (the default) disables early stopping. Any reported metric works (for example a grouped pass@k); a metric that validation does not report fails loudly listing the available keys.
  • grpo.stop_at_validation_threshold (float | null) — end training once the chosen metric reaches this value; setup() asserts it is set whenever the metric is.

The stop takes effect at the end of the stopping step: with checkpointing enabled the step is saved first (forced like a last step, carrying its validation metrics for get_best_checkpoint_path()), and every exit flushes pending checkpoint finalization the same way the existing timeout/max-steps early returns do.

Covered by unit tests parametrized over all three trainers (grpo_train, async_grpo_train, grpo_train_sync): periodic threshold stop, initial-validation stop, configurable metric, loud failure on an unreported metric, and checkpoint-on-stop.

@michal2409
michal2409 requested review from a team as code owners July 29, 2026 11:06
@copy-pr-bot

copy-pr-bot Bot commented Jul 29, 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.

Adds two required grpo keys, defaulted in the exemplar YAMLs:

- stop_at_validation_threshold (float | null): end training once the
  chosen validation metric reaches this value; null (the default)
  disables early stopping.
- stop_at_validation_metric (str, default accuracy): which reported
  validation metric to compare; a metric validation does not report
  fails loudly listing the available keys.

The stop takes effect at the end of the stopping step: with
checkpointing enabled the step is saved first (forced like a last
step, carrying its validation metrics), and every exit flushes pending
checkpoint finalization the same way the timeout/max-steps early
returns do.

Signed-off-by: Michal Futrega <mfutrega@nvidia.com>
@michal2409
michal2409 force-pushed the split/stop-at-accuracy branch from 92ef483 to c149546 Compare July 30, 2026 13:50
@michal2409 michal2409 changed the title feat(grpo): optional early stop at a validation accuracy threshold feat(grpo): optional early stop at a validation metric threshold Jul 30, 2026

@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 @michal2409 , LGTM, just two minor nits.

Comment thread tests/unit/algorithms/test_grpo.py Outdated
Comment thread nemo_rl/algorithms/grpo.py Outdated
yuki-97

This comment was marked as duplicate.

…at-accuracy

Signed-off-by: Michal Futrega <mfutrega@nvidia.com>

# Conflicts:
#	nemo_rl/algorithms/grpo_sync.py
Per review: a null metric disables the early stop (the threshold is no
longer the switch), and setup() asserts the threshold is set whenever
the metric is. Also parametrize the stop tests over grpo_train_sync via
mock_sync_grpo_infrastructure and add an initial-validation stop test
for all three trainers.

Signed-off-by: Michal Futrega <mfutrega@nvidia.com>
yuki-97
yuki-97 previously approved these changes Jul 31, 2026

@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 @michal2409 , LGTM

@yuki-97 yuki-97 added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Jul 31, 2026
@yuki-97

yuki-97 commented Jul 31, 2026

Copy link
Copy Markdown
Contributor

/ok to test aecf0af

…at-accuracy

Signed-off-by: Michal Futrega <mfutrega@nvidia.com>
Their contract test validates the full MasterConfig, so the new
required grpo keys must be present (same trap as the research template
config).

Signed-off-by: Michal Futrega <mfutrega@nvidia.com>
@michal2409

Copy link
Copy Markdown
Contributor Author

/ok to test 0e14faa

max(train/critic/loss) lands at 6.68-7.00 in CI for unrelated PRs
(NVIDIA-NeMo#3401, NVIDIA-NeMo#3404, NVIDIA-NeMo#3423) since the vLLM 0.25.1 bump; the same critic-side
drift is already tracked in NVIDIA-NeMo#3412. Placeholder bump, like the grad_norm
bound raised in NVIDIA-NeMo#3280.

Signed-off-by: Michal Futrega <mfutrega@nvidia.com>
@michal2409

Copy link
Copy Markdown
Contributor Author

/ok to test 0e14faa

michal2409 added a commit to michal2409/RL that referenced this pull request Aug 2, 2026
max(train/critic/loss) lands at 6.68-7.00 in CI for unrelated PRs
(NVIDIA-NeMo#3401, NVIDIA-NeMo#3404, NVIDIA-NeMo#3423) since the vLLM 0.25.1 bump; the same critic-side
drift is already tracked in NVIDIA-NeMo#3412. Placeholder bump, like the grad_norm
bound raised in NVIDIA-NeMo#3280.

Signed-off-by: Michal Futrega <mfutrega@nvidia.com>
@michal2409

Copy link
Copy Markdown
Contributor Author

Same pre-existing L1 PPO flake as on #3401 (critic loss vs the 6.0 bound since the vLLM 0.25.1 bump): raised the bound to 8.0 with a note pointing at #3412.

@michal2409

Copy link
Copy Markdown
Contributor Author

/ok to test 4ab2fe5

@yuki-97

yuki-97 commented Aug 3, 2026

Copy link
Copy Markdown
Contributor

thanks for the bump for L1 PPO threshold, tracked at #3412 (comment) and bump critic/loss (max) bound from 6.0 to 8.0 in this PR to unblock main.

@yuki-97
yuki-97 merged commit 4dbfa20 into NVIDIA-NeMo:main Aug 3, 2026
81 checks passed
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)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants