Fully-masked batch yields NaN loss in 3 of 4 losses; only MaskedCrossEntropy guards it

[kabirvashisht4-glitch]

Describe the bug

MaskedCrossEntropy guards against a zero num_label_tokens; the other three
losses that normalise the same way do not, so a batch with no supervised tokens
produces NaN instead of 0.0. NaN propagates through backward() into every
parameter, so the run keeps going with a destroyed model rather than failing.

nemo_automodel/components/loss/masked_ce.py L84-88 has the guard:

if num_label_tokens is not None:
    assert self.reduction == "sum", "num_label_tokens is only supported when reduction is 'sum'"
    if num_label_tokens == 0:
        return loss * 0.0
    loss = loss / num_label_tokens

The other three divide unconditionally:

Loss	Line	Guard
masked_ce.py MaskedCrossEntropy	86-88	yes
linear_ce.py FusedLinearCrossEntropy	264	no
chunked_ce.py ChunkedCrossEntropy	219	no
te_parallel_ce.py	188	no

The recipe already treats zero as a real case on the validation path
(recipes/llm/train_ft.py L1382):

val_loss = total_loss / max(total_num_label_tokens, 1e-8)

so the training path is the inconsistent one.

Why the count can be zero

num_label_tokens is the global, DP-reduced number of non-ignored labels
(train_ft.py L1210-1213):

num_label_tokens = torch.tensor(
    sum((batch["labels"] != -100).sum().item() for batch in batches), dtype=torch.long
)
num_label_tokens = self._dp_allreduce(num_label_tokens).item()

It reaches zero when every label in the global batch is -100 — for example a
gradient-accumulation window where truncation cut the answer off every sample
(answer_only_loss_mask with seq_length shorter than the prompts), or a bad
shard. The existing guard in MaskedCrossEntropy and the max(..., 1e-8) on the
validation path suggest this has been hit before.

Consequence: which loss you configure decides whether the step survives.
MaskedCrossEntropy contributes 0.0; the other three poison the model. The
recipe can also swap between them on its own — _maybe_downgrade_loss_fn
(train_ft.py L159) replaces a fused loss with MaskedCrossEntropy when the
model does not declare logits_to_keep — so the same config can behave
differently on two models.

Steps/Code to reproduce bug

CPU only, no GPU or checkpoint needed:

import torch
from nemo_automodel.components.loss.masked_ce import MaskedCrossEntropy
from nemo_automodel.components.loss.chunked_ce import ChunkedCrossEntropy

B, S, V = 2, 8, 32
torch.manual_seed(0)
logits = torch.randn(B, S, V)
labels = torch.full((B, S), -100)                 # every position ignored
num_label_tokens = int((labels != -100).sum())    # 0, as the recipe computes it

print(MaskedCrossEntropy()(logits, labels.clone(), num_label_tokens=num_label_tokens))
print(ChunkedCrossEntropy()(logits.clone(), labels.clone(), num_label_tokens=num_label_tokens))

tensor(0.)
tensor(nan)

The unguarded branch is 0.0 / 0 — the sum-reduced loss over zero supervised
tokens is 0.0, and dividing that by 0 gives NaN rather than 0.0.

Expected behavior

All four losses agree: a batch with no supervised tokens contributes 0.0,
matching MaskedCrossEntropy and the validation path, so an empty
gradient-accumulation window is a no-op step rather than a silent model wipe.

Environment overview

• main at 3ddef9b, CPU only. ChunkedCrossEntropy reproduces directly;
  FusedLinearCrossEntropy needs cut_cross_entropy and te_parallel_ce needs
  Transformer Engine, but both perform the identical unguarded division.

Additional context

Happy to send a PR adding the same num_label_tokens == 0 guard to the three
losses, plus a CPU unit test per loss. Alternatively the normalisation could be
factored into one shared helper so the four cannot drift again — say which you
prefer and I will follow that shape.