From cf03d1209ba33e29d45a8e2f7e815ea4a6dea7e9 Mon Sep 17 00:00:00 2001 From: realAsma Date: Tue, 16 Jun 2026 18:58:02 +0000 Subject: [PATCH 1/3] fix(qad): report kd as default eval loss Signed-off-by: realAsma --- modelopt/torch/distill/plugins/huggingface.py | 27 ++++++++++--------- .../distill/plugins/test_huggingface_kd.py | 9 ++++--- 2 files changed, 19 insertions(+), 17 deletions(-) diff --git a/modelopt/torch/distill/plugins/huggingface.py b/modelopt/torch/distill/plugins/huggingface.py index 32a60468921..097195345db 100644 --- a/modelopt/torch/distill/plugins/huggingface.py +++ b/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") @@ -206,7 +206,8 @@ def compute_loss(self, model, inputs, return_outputs=False, **kwargs): loss = kd_loss else: batch_size = find_batch_size(inputs) - self._record_eval_kd_loss(kd_loss, batch_size) + self._record_eval_ce_loss(loss, batch_size) + loss = kd_loss return (loss, outputs) if return_outputs else loss @@ -315,8 +316,8 @@ 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, @@ -324,20 +325,20 @@ def evaluation_loop( 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() diff --git a/tests/unit/torch/distill/plugins/test_huggingface_kd.py b/tests/unit/torch/distill/plugins/test_huggingface_kd.py index eb527171207..0c27f9c4ca5 100644 --- a/tests/unit/torch/distill/plugins/test_huggingface_kd.py +++ b/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): From 467f9af7d602ef1b8500fe6b4bec9e1a14453e22 Mon Sep 17 00:00:00 2001 From: realAsma Date: Tue, 16 Jun 2026 18:59:32 +0000 Subject: [PATCH 2/3] docs: note qad eval metric change Signed-off-by: realAsma --- CHANGELOG.rst | 3 +++ 1 file changed, 3 insertions(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index cab0e5a3942..56380b01b55 100755 --- a/CHANGELOG.rst +++ b/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 ...`` From d0c3d1d71fc554b6ab31beb0368f61f4e88132f7 Mon Sep 17 00:00:00 2001 From: realAsma Date: Wed, 17 Jun 2026 02:51:29 +0000 Subject: [PATCH 3/3] fix(qad): simplify eval ce loss tracking Signed-off-by: realAsma --- modelopt/torch/distill/plugins/huggingface.py | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/modelopt/torch/distill/plugins/huggingface.py b/modelopt/torch/distill/plugins/huggingface.py index 097195345db..a6ecdd0d7c0 100644 --- a/modelopt/torch/distill/plugins/huggingface.py +++ b/modelopt/torch/distill/plugins/huggingface.py @@ -199,17 +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_ce_loss(loss, batch_size) - loss = kd_loss + 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.