Skip to content

fix(sc): gate reference logprobs on KL penalty - #3742

Merged
yfw merged 2 commits into
NVIDIA-NeMo:mainfrom
jinglinglingling:linglinj/v2-zero-kl-reference-logprobs
Sep 1, 2026
Merged

fix(sc): gate reference logprobs on KL penalty#3742
yfw merged 2 commits into
NVIDIA-NeMo:mainfrom
jinglinglingling:linglinj/v2-zero-kl-reference-logprobs

Conversation

@jinglinglingling

@jinglinglingling jinglinglingling commented Aug 21, 2026

Copy link
Copy Markdown
Contributor

Summary

This PR forward-ports the Zero-KL reference-logprob gating introduced in #2443 to NeMo RL v2 (single-controller async).

  • gate v2 SingleController reference-logprob requests on both a positive reference_policy_kl_penalty and the explicit skip flag
  • fetch only the policy/reference logprob data-plane columns produced for each train step
  • reject non-finite/negative KL coefficients and reward-KL configurations that would omit required policy logprobs
  • cover initialization, controller calls, train-field propagation, and invalid configuration combinations

This forward-ports the Zero-KL reference-logprob gating behavior from #2443 to the current v2 single-controller async path.

Test plan

  • affected SingleController and setup suites: 56 passed
  • broader affected CPU suites (controller, setup, checkpointing, split API): 102 passed
  • targeted invalid KL coefficient and data-plane field tests passed
  • pre-commit run --all-files

GPU/runtime qualification was not run locally because this environment does not expose a GPU.

@jinglinglingling
jinglinglingling requested review from a team as code owners August 21, 2026 01:08
@copy-pr-bot

copy-pr-bot Bot commented Aug 21, 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.

@jinglinglingling
jinglinglingling requested review from a team as code owners August 21, 2026 02:46

@jinglinglingling jinglinglingling left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Self-review

No action needed — all high-confidence findings from the self-review were resolved and re-reviewed.

Scope

Reviewed the complete v2 Zero-KL reference-logprob gating change with parallel RL-correctness, bug, test, design, and comment reviewers, followed by an independent devil's-advocate pass.

Findings resolved

  • Fixed the zero-KL producer/consumer mismatch: the trainer now fetches only logprob columns produced for the current step.
  • Added direct train-pump coverage for all four policy/reference logprob gate combinations.
  • Constrained reference_policy_kl_penalty to finite non-negative values.
  • Added setup-time rejection for reward-KL configurations that would skip required policy logprobs.

All fixes landed in e06a1b7342855ebeb7d36139853121477bd27533 and the final adversarial pass returned LGTM with no remaining findings at or above the 80% confidence threshold.

Validation

  • 56 final SingleController/setup tests passed.
  • 102 broader affected CPU tests passed across controller, setup, checkpointing, and split-policy APIs.
  • pre-commit run --all-files passed, including Ruff, formatting, and Pyrefly.
  • PR is conflict-free.

Limitation

GPU/runtime qualification was not run locally because this environment does not expose a GPU. Full NVIDIA-runner validation and reviewer approval remain pending.

Generated by Cursor Agent using the upstream /review-pr-team workflow.

Comment thread nemo_rl/algorithms/single_controller.py
Comment thread nemo_rl/algorithms/single_controller_utils/config.py Outdated
Comment thread tests/unit/single_controller/test_single_controller.py Outdated
Keep v2 reference-logprob requests aligned with reference-model initialization so zero-KL runs do not call an unavailable model.

Signed-off-by: Linglin Jing <linglinj@nvidia.com>
Fetch only data-plane columns produced for each v2 train step and reject KL configurations that would silently omit required logprobs.

Signed-off-by: Linglin Jing <linglinj@nvidia.com>
@jinglinglingling
jinglinglingling force-pushed the linglinj/v2-zero-kl-reference-logprobs branch from e06a1b7 to ff0df3f Compare August 29, 2026 03:10
@yfw

yfw commented Aug 31, 2026

Copy link
Copy Markdown
Contributor

/ok to test ff0df3f

@yfw yfw added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Aug 31, 2026
@yfw
yfw enabled auto-merge (squash) August 31, 2026 15:44
@yfw
yfw merged commit 2aafb9d into NVIDIA-NeMo:main Sep 1, 2026
210 of 215 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