fix(loss): count num_valid_samples from the sample mask, not the batch dim

[tianyi-zhang-02]

What does this PR do ?

Makes num_valid_samples mean the same thing in every loss: the number of samples the microbatch actually contributes, taken from the sample mask.

Issues

None to close.

Usage

No config change. The metric is internal.

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 (unit only)
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

Seven losses in loss_functions.py report num_valid_samples. Four take it from the sample mask. Three take it from the batch dimension:

loss	before
MseValueLossFn	int(values.shape[0])
DistillationLossFn	data["input_ids"].shape[0]
CrossTokenizerDistillationLossFn	data["input_ids"].shape[0]

That number is a gate, not just a log line. All three DTensor workers do:

if num_valid_samples > 0:
    mb_losses.append(loss.item())
    all_mb_metrics.append(loss_metrics)

So a microbatch whose sample_mask is entirely zero reports the raw batch size, passes the gate, and contributes its zero loss to the step's reported mean. ClippedPGLossFn skips that microbatch; the critic and the distillation losses do not.

A fully-masked microbatch is not hypothetical on these three paths in particular — overlong filtering and env-flagged masking zero loss_multiplier on the distillation side, and the critic's seq_logprob_error_threshold filter zeroes sample_mask before the value loss runs. dpo.py and rm.py also use the value as a weighted-average denominator.

DistillationLossFn keeps a batch-dimension fallback rather than requiring the mask: its masking branch is conditional today, and with no mask every sample is valid by definition. (#3496 turns that else into a raise; this is written so it is correct either way, and I'll rebase whichever lands second.)

Tests are CPU-only — they build the loss inputs directly rather than going through prepare_loss_input, so no GPU and no distributed init. Mutation-tested: reverting either site to the batch dimension, and breaking the no-mask fallback, each turn them red.

The cross-tokenizer loss gets the same one-line change but no new test. Every CPU test for it calls the private helpers (_compute_gold, _compute_p_kl); none reaches __call__, where the metrics dict is built. Flagging that rather than claiming coverage I don't have. :)

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