ChunkedCrossEntropy returns num_chunks x the loss for reduction='mean' and wrong-shaped output for 'none'

[kabirvashisht4-glitch]

Describe the bug

ChunkedCrossEntropy sums its per-chunk results. That is correct for
reduction="sum" (the default, which takes a separate kernel path) but wrong for
the fallback loop used by reduction="mean" and reduction="none":

• "mean" returns num_chunks x the correct loss — 128x at the default
  chunk_len=32 on a 4096-token sequence.
• "none" element-wise adds the per-chunk vectors instead of concatenating
  them, returning a tensor of chunk_len values instead of one per token.

Neither raises; both silently return a wrong number.

nemo_automodel/components/loss/chunked_ce.py L199-216:

if self.reduction == "sum":
    loss = _ChunkedCrossEntropySum.apply(logits, labels, self.ignore_index, self.chunk_len)
else:
    ...
    seq_len = logits.shape[0]
    num_chunks = (seq_len + self.chunk_len - 1) // self.chunk_len
    loss = 0.0
    for logits_chunk, targets_chunk in zip(logits.chunk(num_chunks, dim=0), labels.chunk(num_chunks, dim=0)):
        loss += compute_loss(logits_chunk, targets_chunk, self.ignore_index, self.reduction)

Each iteration returns that chunk's mean (or its per-token vector), and +=
accumulates them. A mean of means must be re-weighted, not summed; per-token
vectors must be concatenated, not added.

The class docstring acknowledges the split — "Other reductions fall back to the
legacy per-chunk torch.compile-d F.cross_entropy loop" — so the fallback
is intended to exist; it is just incorrect.

Steps/Code to reproduce bug

CPU only, no GPU or checkpoint needed:

import torch, torch.nn.functional as F
from nemo_automodel.components.loss.chunked_ce import ChunkedCrossEntropy

torch.manual_seed(0)
N, V = 4096, 32
logits = torch.randn(N, V)
labels = torch.randint(0, V, (N,))

ref = F.cross_entropy(logits, labels, reduction="mean", ignore_index=-100)
got = ChunkedCrossEntropy(reduction="mean")(logits.clone(), labels.clone())   # chunk_len=32
print(float(got), float(ref), float(got) / float(ref))

out = ChunkedCrossEntropy(reduction="none")(logits.clone(), labels.clone())
print(tuple(out.shape))

503.25  3.9316  128.0
(32,)

The factor tracks the chunk count exactly:

chunk_len	chunks over 4096 tokens	ratio vs F.cross_entropy
32 (default)	128	128.0x
1024	4	4.0x
2048	2	2.0x
4096	1	1.0x (correct)

Note the last row: when chunk_len >= seq_len there is exactly one chunk and the
result is right, so a short smoke test looks fine. The error only appears once
the sequence is long enough to chunk — which is the situation this loss exists
for.

reduction="sum" is unaffected: it goes through _ChunkedCrossEntropySum and
matches F.cross_entropy(..., reduction="sum") to within 1e-4.

Expected behavior

ChunkedCrossEntropy matches F.cross_entropy for every reduction it accepts:
"mean" returns the mean over non-ignored tokens regardless of chunk_len, and
"none" returns one value per token.

Environment overview

• main at 3ddef9b, CPU only.

Additional context

ChunkedCrossEntropy is reachable from YAML as a loss_fn._target_ (it appears
in examples/llm_finetune/glm/glm_5.2_lora.yaml and
examples/vlm_finetune/qwen3_5_moe/qwen3_5_35b.yaml), and reduction is a
constructor arg, so a config setting reduction: mean gets a loss inflated by
the chunk count — and gradients scaled to match, which behaves like a silently
multiplied learning rate.

Happy to send a PR: accumulate the token-weighted sum for "mean" and divide
once at the end, concatenate for "none", with CPU tests comparing all three
reductions against F.cross_entropy across several chunk_len values including
one that does not divide the sequence evenly.