[yuki-97]

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:

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

This PR:

1. 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.
2. Apply top-p top-k and apply mask -inf on curr_logprobs at prepare_loss_input, since we handle distributed logits to logprobs here.
3. Save and restore sampling params at use_reference_model to make reference_policy_logprobs computed without filtering.
4. 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]

ℹ️ 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]

ℹ️ 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]

ℹ️ 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]

ℹ️ 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]

ℹ️ 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]

ℹ️ 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]

ℹ️ 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]

ℹ️ 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]

ℹ️ 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]

ℹ️ 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]

/ok to test 0e6e0f7