fix(qad): report KD as default eval loss

CHANGELOG.rst

@@ -68,6 +68,9 @@ Changelog
 
 **Backward Breaking Changes**
 
+- ``KDTrainer`` / ``QADTrainer`` evaluation now reports KD as the primary
+  ``eval_loss`` and CE as ``eval_ce_loss``; the previous secondary
+  ``eval_kd_loss`` metric is removed.
 - Reorganize custom CUDA / Triton kernels under ``modelopt.torch.kernels`` into ``common/attention``, ``quantization/{conv,gemm}``, and ``sparsity/attention``. High-level APIs (``mtq.quantize``, ``mtsa.sparsify``, etc.) are unchanged, but **any code importing directly from the kernel subpackages must be updated**: there is no backwards-compatibility shim; the old import paths will raise ``ImportError`` / ``ModuleNotFoundError``. Migration table:
 
   - ``from modelopt.torch.kernels import IS_AVAILABLE, attention, attention_calibrate, register_triton_attention`` → ``from modelopt.torch.kernels.common.attention import ...``

modelopt/torch/distill/plugins/huggingface.py

@@ -143,7 +143,7 @@ def __init__(
             temperature=distill_args.temperature, reduction="none"
         )
         self._teacher_prepared = False
-        self._eval_kd_loss_totals = None
+        self._eval_ce_loss_totals = None
 
         if self.use_liger_kernel:
             self._liger_temperature = distill_args.temperature
@@ -188,7 +188,7 @@ def _ds_gather(self, params):
             yield
 
     def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
-        """Train on KD loss and evaluate on CE loss with KD as a metric."""
+        """Train and evaluate on KD loss, with eval CE tracked as a metric."""
         self._ensure_teacher_prepared()
         kd_inputs = {k: v for k, v in inputs.items() if k != "labels"}
         labels = inputs.get("labels")
@@ -199,16 +199,12 @@ def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
             with student_context():
                 outputs = model(**kd_inputs)
         else:
-            loss, outputs = super().compute_loss(model, inputs, return_outputs=True, **kwargs)
-
-        kd_loss = self._compute_kd_loss(outputs, labels, kd_inputs, **kwargs)
-        if is_training:
-            loss = kd_loss
-        else:
+            ce_loss, outputs = super().compute_loss(model, inputs, return_outputs=True, **kwargs)
             batch_size = find_batch_size(inputs)
-            self._record_eval_kd_loss(kd_loss, batch_size)
+            self._record_eval_ce_loss(ce_loss, batch_size)
 
-        return (loss, outputs) if return_outputs else loss
+        kd_loss = self._compute_kd_loss(outputs, labels, kd_inputs, **kwargs)
+        return (kd_loss, outputs) if return_outputs else kd_loss
 
     def _compute_kd_loss(self, outputs, labels, inputs, **kwargs):
         """Run teacher forward and compute KD loss.
@@ -315,29 +311,29 @@ def evaluation_loop(
         ignore_keys=None,
         metric_key_prefix="eval",
     ):
-        """Add KD loss as a secondary evaluation metric."""
-        self._eval_kd_loss_totals = None
+        """Add CE loss as a secondary evaluation metric."""
+        self._eval_ce_loss_totals = None
         output = super().evaluation_loop(
             dataloader,
             description,
             prediction_loss_only=prediction_loss_only,
             ignore_keys=ignore_keys,
             metric_key_prefix=metric_key_prefix,
         )
-        if self._eval_kd_loss_totals is not None:
-            output.metrics[f"{metric_key_prefix}_kd_loss"] = self._get_eval_kd_loss()
+        if self._eval_ce_loss_totals is not None:
+            output.metrics[f"{metric_key_prefix}_ce_loss"] = self._get_eval_ce_loss()
         return output
 
-    def _record_eval_kd_loss(self, loss, batch_size):
+    def _record_eval_ce_loss(self, loss, batch_size):
         count = loss.new_tensor(float(batch_size or 1))
         totals = torch.stack([loss.detach() * count, count])
-        self._eval_kd_loss_totals = (
+        self._eval_ce_loss_totals = (
             totals
-            if self._eval_kd_loss_totals is None
-            else self._eval_kd_loss_totals + totals.to(self._eval_kd_loss_totals.device)
+            if self._eval_ce_loss_totals is None
+            else self._eval_ce_loss_totals + totals.to(self._eval_ce_loss_totals.device)
         )
 
-    def _get_eval_kd_loss(self):
-        totals = self.accelerator.gather_for_metrics(self._eval_kd_loss_totals)
+    def _get_eval_ce_loss(self):
+        totals = self.accelerator.gather_for_metrics(self._eval_ce_loss_totals)
         totals = totals.reshape(-1, 2).sum(dim=0)
         return (totals[0] / totals[1].clamp(min=1)).item()

tests/unit/torch/distill/plugins/test_huggingface_kd.py

@@ -131,17 +131,18 @@ def test_training_loss_is_kd_and_skips_ce(tmp_path):
     assert loss.item() == pytest.approx(expected_kd_loss.item())
 
 
-def test_eval_loss_is_ce_and_kd_is_secondary_metric(tmp_path):
+def test_eval_loss_is_kd_and_ce_is_secondary_metric(tmp_path):
     student, teacher = _make_models()
     batch = _make_batch()
     expected_ce_loss = student(**batch).loss.detach()
+    expected_kd_loss = _manual_kd_loss(student, teacher, batch)
     trainer = _make_trainer(tmp_path, student, teacher)
 
     metrics = trainer.evaluate()
 
-    assert metrics["eval_loss"] == pytest.approx(expected_ce_loss.item())
-    assert "eval_kd_loss" in metrics
-    assert metrics["eval_kd_loss"] != pytest.approx(metrics["eval_loss"])
+    assert metrics["eval_loss"] == pytest.approx(expected_kd_loss.item())
+    assert metrics["eval_ce_loss"] == pytest.approx(expected_ce_loss.item())
+    assert "eval_kd_loss" not in metrics
 
 
 def test_standard_kd_loss_without_labels_uses_mean(tmp_path):