fix(grpo): mask excluded values before logprob exponentiation

[bzantium]

What does this PR do ?

Prevent masked tokens and excluded samples from producing NaN gradients in ClippedPGLossFn.

The loss currently computes probability ratios and KL terms before applying the loss mask. In float32, a log-probability difference of 200 overflows exp; multiplying the result by a zero mask still yields NaN. Masked -inf or NaN inputs can also contaminate reductions and diagnostics.

Zero the excluded log-probabilities and advantages before those operations, using the existing combined token/sample mask. Values at valid positions are unchanged. This covers the policy ratio, importance-sampling correction, reference KL penalty, and log-probability diagnostics without clamping valid-token ratios.

The regression test compares a clean batch against the same batch with extreme values in a masked prompt position and an excluded sample. It checks finite loss, gradients, and metrics, plus identical loss and input gradients. Its 18 cases cover token/sequence loss reduction, sequence importance ratios, on-policy/off-policy ratios, and finite overflow/-inf/NaN inputs.

Issues

No linked issue.

Usage

No configuration changes are needed.

python -m pytest tests/unit/algorithms/test_loss_functions.py -k clipped_pg_loss

Before your PR is "Ready for review"

• Read and followed the contributor guidelines.
• Added regression tests.
• Ran the full unit and functional test suites locally; see the focused validation below.
• No user-facing documentation changes are needed for this numerical bug fix.

Additional Information

Validated against upstream main d633032b2a8017c69ddfff9183deb42cab6b36f4:

• Focused pytest run: 35 passed, 19 skipped, 49 deselected, including all 18 new regression cases. The skipped tests require CUDA, which was unavailable in this environment.

  python -m pytest tests/unit/algorithms/test_loss_functions.py \
    -k clipped_pg_loss --confcutdir=tests/unit/algorithms --no-cov -q

  --confcutdir excludes the parent suite's automatic Ray-cluster fixtures for this local CPU run; the actual upstream loss implementation and test module are imported without stubs.
• An isolated reproduction using the upstream loss method and KL/reduction helpers failed all 18 cases before the fix and passed all 18 afterward.
• Ruff lint, Ruff format check, and git diff --check passed for the changed files.
• The full GPU unit and functional suites have not been run locally.

[copy-pr-bot]

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.