fix(loss): return zero, not NaN, when a batch has no supervised tokens

[kabirvashisht4-glitch]

What does this PR do ?

Makes a batch with no supervised tokens contribute 0.0 instead of NaN in the
three losses that were missing the guard MaskedCrossEntropy already has.

num_label_tokens is the global, DP-reduced count of non-ignored labels
(recipes/llm/train_ft.py L1210-1213). It is 0 when every label in the global
batch is -100 — for example a gradient-accumulation window where truncation
cut the answer off every sample, or a bad shard. The sum-reduced loss is then
0.0, and 0.0 / 0 is NaN. NaN propagates through backward() into every
parameter, so the run keeps going with a destroyed model rather than failing.

masked_ce.py L84-88 already handles it:

if num_label_tokens == 0:
    return loss * 0.0

The other three divided unconditionally:

Loss	Before	After
MaskedCrossEntropy	0.0	0.0 (unchanged)
FusedLinearCrossEntropy	NaN	0.0
ChunkedCrossEntropy	NaN	0.0
TEParallelCrossEntropy	NaN	0.0

So which loss you configured decided whether the step survived. That is
reachable without the user changing anything: _maybe_downgrade_loss_fn
(train_ft.py L159) swaps a fused loss for MaskedCrossEntropy when the model
does not declare logits_to_keep, so one config could behave differently across
two models.

The recipe already treats zero as a real case on the validation path
(train_ft.py L1382, max(total_num_label_tokens, 1e-8)), so the training path
was the inconsistent one.

loss * 0.0 rather than a fresh torch.zeros keeps the autograd graph intact,
so the step contributes zero gradient instead of detaching.

Changelog

• FusedLinearCrossEntropy, ChunkedCrossEntropy and TEParallelCrossEntropy
  return loss * 0.0 when num_label_tokens == 0, matching
  MaskedCrossEntropy.
• New tests/unit_tests/loss/test_zero_label_tokens.py covering all four losses
  plus a normalization regression guard.

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 add or update any necessary documentation?

Behaviour now matches the documented num_label_tokens contract, so no doc
change was needed.

Verification (CPU, no GPU needed):

• The linear_ce / te_parallel_ce tests stub cut_cross_entropy and the
  Transformer Engine kernel — following the pattern already in
  test_linear_ce.py — so all four guards are actually exercised here rather
  than skipped, even though neither optional package is installed locally.
• Reverting only the three source hunks makes exactly those three tests fail;
  test_masked_ce_zero_label_tokens_is_zero and
  test_nonzero_label_tokens_still_normalizes pass either way, which is
  intended — they pin the reference behaviour and the path this must not change.
• pytest tests/unit_tests/loss/ → 299 passed, 36 skipped, no regressions.
• ruff format / ruff check clean.

Additional Information

• Related to Fully-masked batch yields NaN loss in 3 of 4 losses; only MaskedCrossEntropy guards it #3796

I kept this as the minimal guard so it is easy to review. If you would rather
the normalization were factored into one shared helper so the four cannot drift
apart again, say so and I will restructure it that way.

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