Masked CE losses overwrite the caller's labels in place, zeroing a reused batch's loss

Describe the bug

MaskedCrossEntropy, ChunkedCrossEntropy and TEParallelCrossEntropy all apply their
optional mask with an in-place labels.masked_fill_:

logits = logits.view(-1, logits.size(-1))
labels = labels.view(-1)
if mask is not None:
    with torch.no_grad():
        ...
        labels.masked_fill_(mask.view(-1) == 0, self.ignore_index)

labels.view(-1) does not copy — it aliases the caller's storage. So the fill writes
ignore_index (-100) directly into the caller's label tensor. In
TEParallelCrossEntropy the aliasing is worse: labels has already been through
labels.to_local(), and that local shard aliases the DTensor's storage, so the write
lands in the caller's distributed tensor.

Nothing in the signature or (until now) the docstring says the argument is consumed.
calculate_loss, the wrapper most callers go through, documents the opposite:
"The caller's mapping and tensors are not mutated."

Steps/Code to reproduce bug

import torch, torch.nn.functional as F
from nemo_automodel.components.loss.masked_ce import MaskedCrossEntropy

torch.manual_seed(0)
logits = torch.randn(6, 5)
labels = torch.randint(0, 5, (6,))

mask_a = torch.tensor([1, 1, 1, 0, 0, 0])   # first half supervised
mask_b = torch.tensor([0, 0, 0, 1, 1, 1])   # second half supervised

def reference(mask):
    t = labels.clone(); t[mask == 0] = -100
    return F.cross_entropy(logits, t, reduction="sum")

loss_fn = MaskedCrossEntropy(reduction="sum")
shared = labels.clone()
a = loss_fn(logits, shared, mask=mask_a)
b = loss_fn(logits, shared, mask=mask_b)    # same tensor reused

print(a.item(), reference(mask_a).item())   # 4.717452 4.717452
print(b.item(), reference(mask_b).item())   # 0.000000 3.509581   <-- wrong
print(shared.tolist())                      # [-100, -100, -100, -100, -100, -100]

Expected behavior

Each call scores the positions its own mask selects, and the caller's labels tensor is
unchanged afterwards.

Actual behavior

The first call destroys the supervision. The second call sees a tensor that is already
all ignore_index, so every position is ignored and the loss is exactly 0.0 —
which backprops a zero gradient. There is no exception and no warning; the step simply
contributes nothing.

Any second consumer of the same tensor is affected the same way: another loss term over
the same batch, a token-accuracy or perplexity metric computed after the loss, a cached
or memory-mapped dataset that hands back the same underlying tensor on the next epoch.

Why it has gone unnoticed

No recipe in-tree currently passes mask= (calculate_loss forwards only logits,
labels and num_label_tokens), so today this is a latent defect on a public API rather
than an active training-corruption bug. It is reachable by anyone calling these loss
modules directly, which the mask parameter exists to support.

The existing coverage cannot catch it. test_masked_cross_entropy_with_mask builds its
reference after calling the loss:

loss_custom = MaskedCrossEntropy()(logits, targets, mask=mask)
targets_ref = targets.clone()      # already mutated to -100 here
targets_ref[mask == 0] = -100      # re-applies an applied mask: a no-op

so it re-applies a mask that is already applied and passes either way.

Environment overview

• Reproduced on CPU, torch 2.x, current main (4a334b644). No GPU or distributed
  setup required.

Proposed fix

Use the out-of-place masked_fill in all three losses, document that labels is left
intact, and snapshot targets before the call in the existing test so it can observe a
regression. PR to follow.