fix(grpo): correct sequence importance ratio metric

[CarlosGomes98]

What does this PR do ?

Fixes an issue in the metric computation of importance ratio

The issue seems to be that the importance weights are reshaped to [B, 1] in order to be applied to all the tokens in each sequence, but are left that way when they are then reused together with the sample mask. That broadcasts to [B, B]. The whole matrix then gets summed and divided by the number of valid samples, which results in an error. e.g. for 2 samples, both valid the value will be 2x as large as it should be.

The training loss is unnafected, but the metric is wrong

Issues

List issues that this PR closes (syntax):

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

[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.

[yuki-97]

good catch and thanks for the fix!

[yuki-97]

/ok to test 45abcca