Potential NaN risk in loss_functions.py when mask is all zeros

[zpqiu]

Describe the bug

In nemo_reinforcer/algorithms/loss_functions.py, there's a potential risk of getting NaN values when the mask tensor contains all zeros. This occurs in the following code section:

mult_prob_error = ((torch.exp(lp_error) * mask).sum() / mask.sum()).item()

...

with torch.no_grad():
    probs_ratio = masked_mean(ratios.detach(), mask).item()
    probs_ratio_clamped = masked_mean(ratios_clamped.detach(), mask).item()

If mask contains all zeros, the masked_mean function would likely perform a division by zero (since it's calculating a mean over masked elements), resulting in NaN values.

Steps/Code to reproduce bug

NA

Expected behavior

Add a check to handle the case when the mask is all zeros:

with torch.no_grad():
    if mask.sum() > 0:
        probs_ratio = masked_mean(ratios.detach(), mask).item()
        probs_ratio_clamped = masked_mean(ratios_clamped.detach(), mask).item()
    else:
        probs_ratio = 0.0  # Default value when mask is all zeros
        probs_ratio_clamped = 0.0  # Default value when mask is all zeros

Environment overview (please complete the following information)

• Environment location: Docker
• Method of install: [pip install or from source]. uv pip install -e '.[dev,test]'

[zpqiu]

@terrykong I can take on this issue, but I don't have write permissions.

[terrykong]

assigned! thanks for your help @zpqiu !

[terrykong]

so re: masked mean, i think to match the deepseekmath paper, we need to average over seqdim before batch dim.

i hacked it here in aligner, but we'd need to reflect it here terrykong/NeMo-Aligner@02b02e6

@SahilJain314 just an FYI, this could have a some affect on convergence

[zpqiu]

so re: masked mean, i think to match the deepseekmath paper, we need to average over seqdim before batch dim.

i hacked it here in aligner, but we'd need to reflect it here terrykong/NeMo-Aligner@02b02e6

@SahilJain314 just an FYI, this could have a some affect on convergence

Do I need to change it to the method of first taking the mean along the sequence dim and then taking the mean along the batch dim in my PR?
However, DAPO and Dr.GRPO have already explained that this approach to taking means may have some problems.

[SahilJain314]

We can add the 'original' per-sequence average method as a flag in the loss and keep the current scheme as default. Feel free to do that in a subsequent PR because it's going to require a fair bit of validation/testing.