fix(loss): stop masked losses from overwriting the caller's labels

[kabirvashisht4-glitch]

What does this PR do ?

Fixes #3917: MaskedCrossEntropy, ChunkedCrossEntropy and TEParallelCrossEntropy
applied their optional mask with an in-place labels.masked_fill_, writing
ignore_index into the caller's tensor.

labels.view(-1) aliases the caller's storage rather than copying it, so the fill is
visible to whoever owns the batch. TEParallelCrossEntropy 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.

The failure is silent and it zeroes a loss. Reusing the tensor for a second mask scores
the intersection of the two masks, not the second one:

	mask_a = first half	mask_b = second half (same tensor)
expected	4.717452	3.509581
before	4.717452	0.000000
after	4.717452	3.509581

A loss of exactly 0.0 backprops a zero gradient with no exception and no warning. The
same applies to any second consumer of the batch: another loss term, a token-accuracy or
perplexity metric computed after the loss, or a cached dataset that returns the same
underlying tensor on the next epoch.

calculate_loss — the wrapper most callers go through — documents the opposite
guarantee: "The caller's mapping and tensors are not mutated."

Scope

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

Why the existing test could not catch it

test_masked_cross_entropy_with_mask built its reference after the loss call:

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

It re-applied an already-applied mask, so it passed with or without the bug. This PR
snapshots targets beforehand, which makes the existing assertion able to fail.

Changelog

• Replace the in-place labels.masked_fill_ with the out-of-place masked_fill in
  masked_ce.py, chunked_ce.py and te_parallel_ce.py.
• Document on all three that labels is not modified in place. chunked_ce.py
  previously documented the opposite ("ignored positions are replaced with
  ignore_index in this tensor"); that line is now correct.
• Snapshot targets before the loss call in test_masked_cross_entropy_with_mask so it
  can observe an in-place write.
• Add, for MaskedCrossEntropy and ChunkedCrossEntropy: a test that the caller's
  labels is unchanged, and a test that two disjoint masks over one tensor each score
  their own positions. Add the corresponding no-mutation test for
  TEParallelCrossEntropy, guarded on HAVE_TE_PARALLEL_CE and CUDA.

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?

Verified locally:

• pytest tests/unit_tests/loss/ -> 313 passed, 37 skipped.
• The four new CPU tests fail on main and pass here. Reverting only the
  nemo_automodel/ changes and keeping the tests reproduces:
  second mask scored 0.000000, expected 6.012943.
• ruff format --check and ruff check clean on every changed file.
• CPU only; the TE test is skipped without Transformer Engine and CUDA.

Additional Information

• The three call sites are the only in-place writes to labels under
  nemo_automodel/components/loss/.

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

[akoumpa]

Hi @kabirvashisht4-glitch can you take a look at the conflicts when you get the chance? thank you