fix(megatron): honor policy.logprob_chunk_size in the training loss path by kaloyan-inherent · Pull Request #2872 · NVIDIA-NeMo/RL

kaloyan-inherent · 2026-06-18T19:35:04Z

What does this PR do ?

This PR closes #2871 by passing chunk_size for the logprobs calculation.

Issues

Closes #2871

Signed-off-by: Kaloyan <253267049+kaloyan-inherent@users.noreply.github.com>

copy-pr-bot · 2026-06-18T19:35:08Z

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.

terrykong

@ZhiyuLi-Nvidia can you review?

ZhiyuLi-Nvidia · 2026-06-22T21:01:32Z

Thank you @kaloyan-inherent generally look good to me. Could you just add a equivalence test of logprob function with and without chunk_size?

Signed-off-by: Kaloyan <253267049+kaloyan-inherent@users.noreply.github.com>

kaloyan-inherent · 2026-06-22T21:33:54Z

Thank you @kaloyan-inherent generally look good to me. Could you just add a equivalence test of logprob function with and without chunk_size?

thanks for the review @ZhiyuLi-Nvidia! sounds good -- i added the test

Signed-off-by: Kaloyan <253267049+kaloyan-inherent@users.noreply.github.com>

ZhiyuLi-Nvidia · 2026-06-22T22:22:31Z

/ok to test c7d05e8

ZhiyuLi-Nvidia

Thank you @kaloyan-inherent. LGTM!

…ath (#2872) Signed-off-by: Kaloyan <253267049+kaloyan-inherent@users.noreply.github.com> Signed-off-by: Anna Shors <ashors@nvidia.com>

…2872 The rebased Gym PR #2872 tree pins openai==2.44.0 while sglang==0.5.12.post1 pins openai==2.6.1; no environment installs both extras together, so declare the conflict and relock. uv.lock is the final state validated by 'uv lock --check' in the container preflights. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Thombre <pthombre@nvidia.com>

Bump the Gym pin from 84f2b29e to 8e983c57: latest PR #2872 head (beaa94c8, adds the staging-consolidation refactor 313754e2, swe-agents orjson reads, and a main merge) with the strict-strip fix 84f2b29e cherry-picked on top, since the PR still lacks it and dropping it would reintroduce the runaway-generation regression validated in E5. The refactor removed nemo_gym.token_id_capture.staging.attribution; resolve_terminal is now re-exported from the staging package with an unchanged call signature, so only the import moves. Gym's pyproject.toml is byte-identical across the bump, so uv.lock is untouched. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Thombre <pthombre@nvidia.com>

Bump the live Gym worktree from Stack #2777 top 577ded84 to the rebased-on-main PR #2872 head bc17ce17dc9744e16dae94d0fff4749b4811838c to rerun the fullperf capture experiment against the rebased stack. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Gym PR #2872 (rebased on main) pins openai==2.44.0 while sglang==0.5.12.post1 pins openai==2.6.1. Declare nemo-gym and sglang extras as conflicting -- no environment installs both together -- so uv forks openai (2.6.1 for sglang, 2.44.0 for gym) instead of failing resolution. Relocked with uv 0.11.28 to match the perf container. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Bump the Gym pin from 84f2b29e to 8e983c57: latest PR #2872 head (beaa94c8, adds the staging-consolidation refactor 313754e2, swe-agents orjson reads, and a main merge) with the strict-strip fix 84f2b29e cherry-picked on top, since the PR still lacks it and dropping it would reintroduce the runaway-generation regression validated in E5. The refactor removed nemo_gym.token_id_capture.staging.attribution; resolve_terminal is now re-exported from the staging package with an unchanged call signature, so only the import moves. Gym's pyproject.toml is byte-identical across the bump, so uv.lock is untouched. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Thombre <pthombre@nvidia.com>

Bump the live Gym worktree from Stack #2777 top 577ded84 to the rebased-on-main PR #2872 head bc17ce17dc9744e16dae94d0fff4749b4811838c to rerun the fullperf capture experiment against the rebased stack. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Thombre <pthombre@nvidia.com>

Gym PR #2872 (rebased on main) pins openai==2.44.0 while sglang==0.5.12.post1 pins openai==2.6.1. Declare nemo-gym and sglang extras as conflicting -- no environment installs both together -- so uv forks openai (2.6.1 for sglang, 2.44.0 for gym) instead of failing resolution. Relocked with uv 0.11.28 to match the perf container. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Thombre <pthombre@nvidia.com>

Bump the Gym pin from 84f2b29e to 8e983c57: latest PR #2872 head (beaa94c8, adds the staging-consolidation refactor 313754e2, swe-agents orjson reads, and a main merge) with the strict-strip fix 84f2b29e cherry-picked on top, since the PR still lacks it and dropping it would reintroduce the runaway-generation regression validated in E5. The refactor removed nemo_gym.token_id_capture.staging.attribution; resolve_terminal is now re-exported from the staging package with an unchanged call signature, so only the import moves. Gym's pyproject.toml is byte-identical across the bump, so uv.lock is untouched. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Thombre <pthombre@nvidia.com>

The previous pin d8a11dba existed only in the local submodule clone (never pushed to NVIDIA-NeMo/Gym), which broke every CI job at submodule checkout ('not our ref') and the submodule fast-forward check. Re-pin to 7dea0e786, the live head of NVIDIA-NeMo/Gym#2872, and relock for its pyproject delta (uv 0.11.28, lockfile revision 3). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Thombre <pthombre@nvidia.com>

Pin 7b842eab: verify_and_linearize replaced with the metadata-only implementation this branch's finalizer now calls (StagedCallBaseSnapshot, ExtrasCommitment, public route codec + span decision table). The SHA is pushed to NVIDIA-NeMo/Gym#2872, so CI submodule checkout resolves it. Signed-off-by: Pranav Thombre <pthombre@nvidia.com>

logprob chunking: fix path to not cast fp32

b4d12f4

Signed-off-by: Kaloyan <253267049+kaloyan-inherent@users.noreply.github.com>

kaloyan-inherent requested review from a team as code owners June 18, 2026 19:35

kaloyan-inherent mentioned this pull request Jun 18, 2026

[bug] policy.logprob_chunk_size is ignored by the training loss path which leads to increased memory #2871

Closed

github-actions Bot added the community-request label Jun 18, 2026

svcnvidia-nemo-ci added the waiting-on-maintainers Waiting on maintainers to respond label Jun 20, 2026

terrykong reviewed Jun 22, 2026

View reviewed changes

terrykong requested a review from ZhiyuLi-Nvidia June 22, 2026 16:54

svcnvidia-nemo-ci removed the waiting-on-maintainers Waiting on maintainers to respond label Jun 22, 2026

ZhiyuLi-Nvidia reviewed Jun 22, 2026

View reviewed changes

Comment thread tests/unit/distributed/test_distributed_logprob.py

svcnvidia-nemo-ci added the waiting-on-customer Waiting on the original author to respond label Jun 22, 2026

pr review: address review by adding equivalence test

4ec5a73

Signed-off-by: Kaloyan <253267049+kaloyan-inherent@users.noreply.github.com>

kaloyan-inherent requested a review from ZhiyuLi-Nvidia June 22, 2026 21:33

ZhiyuLi-Nvidia reviewed Jun 22, 2026

View reviewed changes

Comment thread nemo_rl/distributed/model_utils.py Outdated

ZhiyuLi-Nvidia reviewed Jun 22, 2026

View reviewed changes

Comment thread tests/unit/distributed/test_distributed_logprob.py

suffix: add fp32 for clarity

c7d05e8

Signed-off-by: Kaloyan <253267049+kaloyan-inherent@users.noreply.github.com>

kaloyan-inherent requested a review from ZhiyuLi-Nvidia June 22, 2026 21:53

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

copy-pr-bot Bot temporarily deployed to public June 22, 2026 22:22 Inactive

copy-pr-bot Bot temporarily deployed to public June 22, 2026 22:23 Inactive

copy-pr-bot Bot temporarily deployed to test June 22, 2026 22:25 Inactive

copy-pr-bot Bot temporarily deployed to public June 22, 2026 22:26 Inactive

svcnvidia-nemo-ci removed the waiting-on-customer Waiting on the original author to respond label Jun 23, 2026

terrykong enabled auto-merge (squash) June 23, 2026 15:59

ZhiyuLi-Nvidia approved these changes Jun 23, 2026

View reviewed changes

terrykong merged commit f67e9e6 into NVIDIA-NeMo:main Jun 23, 2026
76 of 77 checks passed

Navigation Menu

Uh oh!

fix(megatron): honor policy.logprob_chunk_size in the training loss path - #2872

fix(megatron): honor policy.logprob_chunk_size in the training loss path#2872
terrykong merged 3 commits into
NVIDIA-NeMo:mainfrom
kaloyan-inherent:kally/fix-logprob-chunk-loss-path

kaloyan-inherent commented Jun 18, 2026

Uh oh!

copy-pr-bot Bot commented Jun 18, 2026

Uh oh!

terrykong left a comment

Uh oh!

ZhiyuLi-Nvidia commented Jun 22, 2026

Uh oh!

Uh oh!

kaloyan-inherent commented Jun 22, 2026

Uh oh!

Uh oh!

Uh oh!

ZhiyuLi-Nvidia commented Jun 22, 2026

Uh oh!

ZhiyuLi-Nvidia left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

kaloyan-inherent commented Jun 18, 2026

What does this PR do ?

Issues

Uh oh!

copy-pr-bot Bot commented Jun 18, 2026

Uh oh!

terrykong left a comment

Choose a reason for hiding this comment

Uh oh!

ZhiyuLi-Nvidia commented Jun 22, 2026

Uh oh!

Uh oh!

kaloyan-inherent commented Jun 22, 2026

Uh oh!

Uh oh!

Uh oh!

ZhiyuLi-Nvidia commented Jun 22, 2026

Uh oh!

ZhiyuLi-Nvidia left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants