feat: support top-p top-k in grpo by yuki-97 · Pull Request #2053 · NVIDIA-NeMo/RL

yuki-97 · 2026-03-03T06:23:48Z

Thanks a lot for @zhandaz for the original implementation #1578 and re-implement #1938. Really awesome work!
This PR re-implement this based on the refactored loss structure.

What does this PR do ?

For the logprobs:

generation_logprobs, curr_logprobs and prev_logprobs are computed with filtering.
reference_policy_logprobs is computed without filtering.

This PR:

Apply top-p top-k and apply mask -inf on prev_logprobs at get_logprobs, since we need to call compute_and_apply_seq_logprob_error_masking before train, which needs the updated prev_logprobs to compare with generation_logprobs.
Apply top-p top-k and apply mask -inf on curr_logprobs at prepare_loss_input, since we handle distributed logits to logprobs here.
Save and restore sampling params at use_reference_model to make reference_policy_logprobs computed without filtering.
For TP, we convert from vocab-parallel to batch-sequence-parallel layout via all-to-all communication, apply filtering on the full vocabulary, then convert back. This avoids materializing the full vocabulary on any single rank.

More details can be seen at #1578 and #1938.

Test Result

	convergence	time cost
dtensorv1
dtensorv2
mcore

Summary by CodeRabbit

New Features
- Added configurable sampling parameters (temperature, top-k, top-p) for training with new YAML configuration examples.
- Introduced memory-efficient top-k and top-p filtering for logit sampling.
Bug Fixes
- Improved token-level masking for logprob calculations in training pipelines.
Tests
- Added comprehensive test coverage for sampling parameter validation and distributed logprob computations.
- Expanded nightly test suite with new sampling configuration scenarios.

github-actions · 2026-03-03T06:24:09Z

ℹ️ File Consistency Check

Check based on commit: d2c89a8 (PR #2053 from yukih/topp-topk)

✅ DTensor Policy Worker Synchronization Check

Both DTensor policy worker files were modified in this PR:

nemo_rl/models/policy/workers/dtensor_policy_worker.py
nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py

Please ensure that the changes are consistent between both files where applicable.

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

github-actions · 2026-03-03T07:58:41Z

ℹ️ File Consistency Check

Check based on commit: 75ab1d4 (PR #2053 from yukih/topp-topk)

✅ DTensor Policy Worker Synchronization Check

Both DTensor policy worker files were modified in this PR:

nemo_rl/models/policy/workers/dtensor_policy_worker.py
nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py

Please ensure that the changes are consistent between both files where applicable.

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

github-actions · 2026-03-03T08:57:39Z

ℹ️ File Consistency Check

Check based on commit: 7931cb0 (PR #2053 from yukih/topp-topk)

✅ DTensor Policy Worker Synchronization Check

Both DTensor policy worker files were modified in this PR:

nemo_rl/models/policy/workers/dtensor_policy_worker.py
nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py

Please ensure that the changes are consistent between both files where applicable.

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

github-actions · 2026-03-03T09:57:53Z

ℹ️ File Consistency Check

Check based on commit: 76064bd (PR #2053 from yukih/topp-topk)

✅ DTensor Policy Worker Synchronization Check

Both DTensor policy worker files were modified in this PR:

nemo_rl/models/policy/workers/dtensor_policy_worker.py
nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py

Please ensure that the changes are consistent between both files where applicable.

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

github-actions · 2026-03-03T10:31:07Z

ℹ️ File Consistency Check

Check based on commit: a0423a5 (PR #2053 from yukih/topp-topk)

✅ DTensor Policy Worker Synchronization Check

Both DTensor policy worker files were modified in this PR:

nemo_rl/models/policy/workers/dtensor_policy_worker.py
nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py

Please ensure that the changes are consistent between both files where applicable.

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

github-actions · 2026-03-03T15:00:34Z

ℹ️ File Consistency Check

Check based on commit: 3d5ff4f (PR #2053 from yukih/topp-topk)

✅ DTensor Policy Worker Synchronization Check

Both DTensor policy worker files were modified in this PR:

nemo_rl/models/policy/workers/dtensor_policy_worker.py
nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py

Please ensure that the changes are consistent between both files where applicable.

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

github-actions · 2026-03-03T16:23:21Z

ℹ️ File Consistency Check

Check based on commit: 69966dc (PR #2053 from yukih/topp-topk)

✅ DTensor Policy Worker Synchronization Check

Both DTensor policy worker files were modified in this PR:

nemo_rl/models/policy/workers/dtensor_policy_worker.py
nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py

Please ensure that the changes are consistent between both files where applicable.

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

github-actions · 2026-03-04T05:21:15Z

ℹ️ File Consistency Check

Check based on commit: 6ccc659 (PR #2053 from yukih/topp-topk)

✅ DTensor Policy Worker Synchronization Check

Both DTensor policy worker files were modified in this PR:

nemo_rl/models/policy/workers/dtensor_policy_worker.py
nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py

Please ensure that the changes are consistent between both files where applicable.

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

github-actions · 2026-03-04T08:42:29Z

ℹ️ File Consistency Check

Check based on commit: 95ec444 (PR #2053 from yukih/topp-topk)

✅ DTensor Policy Worker Synchronization Check

Both DTensor policy worker files were modified in this PR:

nemo_rl/models/policy/workers/dtensor_policy_worker.py
nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py

Please ensure that the changes are consistent between both files where applicable.

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

github-actions · 2026-03-09T09:19:32Z

ℹ️ File Consistency Check

Check based on commit: cc84132 (PR #2053 from yukih/topp-topk)

✅ DTensor Policy Worker Synchronization Check

Both DTensor policy worker files were modified in this PR:

nemo_rl/models/policy/workers/dtensor_policy_worker.py
nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py

Please ensure that the changes are consistent between both files where applicable.

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

github-actions · 2026-03-09T14:22:25Z

ℹ️ File Consistency Check

Check based on commit: a0474b3 (PR #2053 from yukih/topp-topk)

✅ DTensor Policy Worker Synchronization Check

Both DTensor policy worker files were modified in this PR:

nemo_rl/models/policy/workers/dtensor_policy_worker.py
nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py

Please ensure that the changes are consistent between both files where applicable.

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

github-actions · 2026-03-10T01:27:44Z

ℹ️ File Consistency Check

Check based on commit: b769172 (PR #2053 from yukih/topp-topk)

✅ DTensor Policy Worker Synchronization Check

Both DTensor policy worker files were modified in this PR:

nemo_rl/models/policy/workers/dtensor_policy_worker.py
nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py

Please ensure that the changes are consistent between both files where applicable.

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

(not cp loss_function.py common.py megatron_policy_worker.py) Signed-off-by: Zhanda <zhandazhu@gmail.com>

Signed-off-by: Yuki Huang <yukih@nvidia.com> Co-authored-by: Zhanda <zhandazhu@gmail.com>

Signed-off-by: Yuki Huang <yukih@nvidia.com>

…_masking Signed-off-by: Yuki Huang <yukih@nvidia.com>

Signed-off-by: Yuki Huang <yukih@nvidia.com>

github-actions · 2026-03-10T03:59:30Z

ℹ️ File Consistency Check

Check based on commit: 0e6e0f7 (PR #2053 from yukih/topp-topk)

✅ DTensor Policy Worker Synchronization Check

Both DTensor policy worker files were modified in this PR:

nemo_rl/models/policy/workers/dtensor_policy_worker.py
nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py

Please ensure that the changes are consistent between both files where applicable.

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

yuki-97 · 2026-03-10T03:59:54Z

/ok to test 0e6e0f7

yuki-97 force-pushed the yukih/topp-topk branch from 75ab1d4 to 7931cb0 Compare March 3, 2026 08:57

yuki-97 force-pushed the yukih/topp-topk branch from 7931cb0 to 76064bd Compare March 3, 2026 09:57

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

yuki-97 had a problem deploying to nemo-ci March 3, 2026 09:58 — with GitHub Actions Failure

yuki-97 force-pushed the yukih/topp-topk branch from 76064bd to a0423a5 Compare March 3, 2026 10:30

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

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

yuki-97 had a problem deploying to nemo-ci March 3, 2026 10:31 — with GitHub Actions Failure

yuki-97 added CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) and removed CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) labels Mar 3, 2026

yuki-97 had a problem deploying to nemo-ci March 3, 2026 15:01 — with GitHub Actions Failure

yuki-97 force-pushed the yukih/topp-topk branch from 3d5ff4f to 69966dc Compare March 3, 2026 16:22

yuki-97 added CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) and removed CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) labels Mar 3, 2026

yuki-97 temporarily deployed to nemo-ci March 3, 2026 16:24 — with GitHub Actions Inactive

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

yuki-97 added the CI:L1 Run doctests, unit tests, and functional tests label Mar 4, 2026

yuki-97 temporarily deployed to nemo-ci March 4, 2026 08:43 — with GitHub Actions Inactive

yuki-97 temporarily deployed to nemo-ci March 4, 2026 08:46 — with GitHub Actions Inactive

yuki-97 force-pushed the yukih/topp-topk branch from 95ec444 to 09341b1 Compare March 4, 2026 16:11

yuki-97 removed the CI:L1 Run doctests, unit tests, and functional tests label Mar 4, 2026

yuki-97 temporarily deployed to nemo-ci March 9, 2026 08:51 — with GitHub Actions Inactive

yuki-97 force-pushed the yukih/topp-topk branch from 3e03356 to cc84132 Compare March 9, 2026 09:19

yuki-97 added CI:L1 Run doctests, unit tests, and functional tests and removed CI:L1 Run doctests, unit tests, and functional tests labels Mar 9, 2026

yuki-97 temporarily deployed to nemo-ci March 9, 2026 09:19 — with GitHub Actions Inactive

zhandaz and others added 20 commits March 9, 2026 20:58

Reimplement top-p and top-k from #1578

ce6c0f3

(not cp loss_function.py common.py megatron_policy_worker.py) Signed-off-by: Zhanda <zhandazhu@gmail.com>

apply previous change in loss_function.py

e3219c1

Signed-off-by: Yuki Huang <yukih@nvidia.com> Co-authored-by: Zhanda <zhandazhu@gmail.com>

split to logits_sampling_utils.py

e087ee1

Signed-off-by: Yuki Huang <yukih@nvidia.com>

move TrainingSamplingParams to setup

961f946

Signed-off-by: Yuki Huang <yukih@nvidia.com>

add topp topk in LogprobsPostProcessor

e4658a9

Signed-off-by: Yuki Huang <yukih@nvidia.com>

update megatron

83899a7

Signed-off-by: Yuki Huang <yukih@nvidia.com>

update need_top_k_or_top_p_filtering

cdaa31a

Signed-off-by: Yuki Huang <yukih@nvidia.com>

fix test

c40e497

Signed-off-by: Yuki Huang <yukih@nvidia.com>

apply filter prev_logprobs before compute_and_apply_seq_logprob_error…

5c2abeb

…_masking Signed-off-by: Yuki Huang <yukih@nvidia.com>

update megatron and dtensorv1, and fix unit test

50f7837

Signed-off-by: Yuki Huang <yukih@nvidia.com>

update nightly test

0b333e3

Signed-off-by: Yuki Huang <yukih@nvidia.com>

make megatron worker respect to make_sequence_length_divisible_by

eeddf32

Signed-off-by: Yuki Huang <yukih@nvidia.com>

fix unit test

3cc9cfa

Signed-off-by: Yuki Huang <yukih@nvidia.com>

limit nightly test threshold

120594a

Signed-off-by: Yuki Huang <yukih@nvidia.com>

remove default value and update top_k check

2283932

Signed-off-by: Yuki Huang <yukih@nvidia.com>

address comments

793088f

Signed-off-by: Yuki Huang <yukih@nvidia.com>

add doc

77e4233

Signed-off-by: Yuki Huang <yukih@nvidia.com>

add functional test

2a7202a

Signed-off-by: Yuki Huang <yukih@nvidia.com>

fix unit test

4b020fc

Signed-off-by: Yuki Huang <yukih@nvidia.com>

update threshold

0e6e0f7

Signed-off-by: Yuki Huang <yukih@nvidia.com>

Conversation

yuki-97 commented Mar 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Test Result

Summary by CodeRabbit

Uh oh!

github-actions bot commented Mar 3, 2026

ℹ️ File Consistency Check

✅ DTensor Policy Worker Synchronization Check

Uh oh!

github-actions bot commented Mar 3, 2026

ℹ️ File Consistency Check

✅ DTensor Policy Worker Synchronization Check

Uh oh!

github-actions bot commented Mar 3, 2026

ℹ️ File Consistency Check

✅ DTensor Policy Worker Synchronization Check

Uh oh!

github-actions bot commented Mar 3, 2026

ℹ️ File Consistency Check

✅ DTensor Policy Worker Synchronization Check

Uh oh!

github-actions bot commented Mar 3, 2026

ℹ️ File Consistency Check

✅ DTensor Policy Worker Synchronization Check

Uh oh!

github-actions bot commented Mar 3, 2026

ℹ️ File Consistency Check

✅ DTensor Policy Worker Synchronization Check

Uh oh!

github-actions bot commented Mar 3, 2026

ℹ️ File Consistency Check

✅ DTensor Policy Worker Synchronization Check

Uh oh!

github-actions bot commented Mar 4, 2026

ℹ️ File Consistency Check

✅ DTensor Policy Worker Synchronization Check

Uh oh!

github-actions bot commented Mar 4, 2026

ℹ️ File Consistency Check

✅ DTensor Policy Worker Synchronization Check

Uh oh!

github-actions bot commented Mar 9, 2026

ℹ️ File Consistency Check

✅ DTensor Policy Worker Synchronization Check

Uh oh!

github-actions bot commented Mar 9, 2026

ℹ️ File Consistency Check

✅ DTensor Policy Worker Synchronization Check

Uh oh!

github-actions bot commented Mar 10, 2026

ℹ️ File Consistency Check

✅ DTensor Policy Worker Synchronization Check

Uh oh!

github-actions bot commented Mar 10, 2026

ℹ️ File Consistency Check

✅ DTensor Policy Worker Synchronization Check

Uh oh!

yuki-97 commented Mar 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

yuki-97 commented Mar 3, 2026 •

edited

Loading