fix(distill): flatten MFTLoss labels along with the logits - #2645

What does this PR do?

Type of change: Bug fix

MFTLoss.forward flattens both logit tensors to (batch * positions, vocab) and its docstring only assumes the class dimension is last, so leading dimensions are clearly meant to be allowed:

soft_log_probs = soft_log_probs.view(-1, soft_log_probs.size(-1))   # (new B, C)
target_logits  = target_logits.view(-1, target_logits.size(-1))     # (new B, C)
soft_targets = self._prepare_corrected_distributions(target_logits, labels, ...)

The labels were passed through untouched, and _prepare_corrected_distributions rejects anything that is not 1-D:

ValueError: Logits must be a 2D tensor and labels must be a 1D tensor.

So the shapes a language model actually produces — (batch, seq_len, vocab) logits against (batch, seq_len) labels — cannot be used:

That is the setting Minifinetuning (arXiv:2506.15702) is for, so in practice a caller had to flatten the labels themselves and nothing documented that.

Flattening them alongside the logits fixes it, and the docstring now says what shape the labels are expected in.

Testing

test_mft_loss_accepts_sequence_shaped_logits drives MFTLoss at (2, 8, 50) / (2, 8) and asserts the result equals the pre-flattened (16, 50) / (16,) call. It fails on main with the ValueError above and passes here.

Worth noting why this was not caught: the existing test_distillation_model_mft drives MFTLoss through a vision model, whose logits are already (batch, classes) and labels already (batch,), so the flattening never does anything there.

tests/unit/torch/distill is 30 passed, 1 skipped locally. The skip is plugins/test_huggingface_kd.py, which needs transformers and is not installed in my environment; it exercises LogitsDistillationLoss, which this change does not touch.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ — already-1-D labels reshape to themselves, so existing callers are unaffected
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ✅
• Did you get Claude approval on this PR?: N/A

Additional Information

Unrelated to my open ONNX PRs (#2553, #2554, #2567, #2575, #2583) — no shared files.

One thing I noticed next door and did not touch, in case it is of interest: LogitsDistillationLoss defaults to reduction="mean", while MFTLoss in the same file defaults to "batchmean". PyTorch warns on every call that "mean" is not the KL divergence value and that it will be changed to behave as "batchmean" in a future major release, so the default path is currently a factor of the vocabulary size away from the other two reductions and will shift silently when that lands. Changing a training default is your call rather than mine, so I have left it alone — happy to open a separate issue if useful.

Summary by CodeRabbit

• Bug Fixes
  • Fixed MFTLoss to handle sequence-shaped labels alongside sequence-shaped logits, producing results consistent with equivalent flattened inputs.
• Documentation
  • Clarified that labels should provide one value per logit position.