From 4c1a80103da4fda18c5fc45a4bcd608dd097e3c0 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Wed, 28 Oct 2026 00:40:51 -0700 Subject: [PATCH 1/3] Compute the Top-K KD loss's full-vocab normalizers without fp32 copies of the logits TopLogitsKLLoss kept about five float32 copies of the student logits alive at the backward peak, which runs out of memory at long sequences with large vocabularies. The normalizers now run over row chunks, and a custom autograd function saves only the logits and writes their gradient chunk by chunk. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- CHANGELOG.rst | 1 + modelopt/torch/distill/plugins/megatron.py | 93 ++++++++++++++++--- .../distill/plugins/test_distill_megatron.py | 26 +++++- 3 files changed, 101 insertions(+), 19 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 66f4015c760..d9de64cbf67 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -75,6 +75,7 @@ Changelog - Fix ONNX quantization with the default ``concat_elimination`` pass crashing in calibration with ``KeyError: '_ReduceMin'`` when a ``Concat`` input is a constant, such as the class token of a ViT. - Fix an illegal memory access in static NVFP4 fake quantization for tensors with more than 2^31 elements, such as stacked MoE expert weights. The kernel is also about 14x faster. - Fix Megatron Top-K/Top-P distillation gradients being multiplied by the tensor-parallel size when tensor parallelism is enabled; reported loss values are unchanged. +- Fix Megatron Top-K/Top-P distillation running out of memory at long sequence lengths with large vocabularies: since 0.48 its full-vocabulary normalizers kept several float32 copies of the logits alive. It now adds about one copy of the logits in their own dtype, with identical loss and gradients. - Fix the Puzzletron memory objective pricing the KV cache with ``hidden_size // num_attention_heads`` instead of the model's ``head_dim``, which misestimated it for models such as Mistral-Small, GPT-OSS and Gemma-3 and for every candidate with a pruned hidden size. Recompute subblock statistics produced with earlier versions. - Fix ONNX INT4 quantization ending the Python process when the INT4 packing extension fails to build, for example without a C++ compiler. It now falls back to the Python implementation. - Fix ``fsdp2_aware_weight_update`` masking an error raised during its unshard (e.g. a CUDA OOM) with a spurious ``UnboundLocalError``. diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py index 63745ed48e0..adfb41bfde7 100644 --- a/modelopt/torch/distill/plugins/megatron.py +++ b/modelopt/torch/distill/plugins/megatron.py @@ -403,13 +403,81 @@ def _tp_logsumexp(self, logits: Tensor) -> Tensor: return logits_max + torch.log(denom) +# fp32 elements per row chunk in the Top-K loss's full-vocab passes (1 GiB). +_TOP_LOGITS_CHUNK_NUMEL = 1 << 28 + + +def _row_chunks(rows: Tensor) -> list[slice]: + step = max(1, _TOP_LOGITS_CHUNK_NUMEL // rows.size(-1)) + return [slice(i, i + step) for i in range(0, rows.size(0), step)] + + +def _scaled_logsumexp( + logits: Tensor, temperature: float, tp_group, top_k: int = 0 +) -> tuple[Tensor, Tensor, Tensor]: + """fp32 ``logsumexp(logits / temperature)`` over the TP-sharded vocab dim, by row chunks. + + Also returns the local Top-K values and indices of ``logits / temperature`` (empty if ``top_k=0``). + """ + rows = logits.detach().reshape(-1, logits.size(-1)) + # Rounding is monotonic, so this is also the max of the scaled fp32 logits. + lse = rows.amax(dim=-1, keepdim=True).float() / temperature + if tp_group is not None: + torch.distributed.all_reduce(lse, op=torch.distributed.ReduceOp.MAX, group=tp_group) + sum_exp = torch.empty_like(lse) + top_vals = lse.new_empty(rows.size(0), top_k) + top_idx = torch.empty(rows.size(0), top_k, dtype=torch.long, device=rows.device) + for c in _row_chunks(rows): + scaled = rows[c].to(torch.float32, copy=True).div_(temperature) + if top_k: + top_vals[c], top_idx[c] = torch.topk(scaled, top_k, dim=-1) + sum_exp[c] = scaled.sub_(lse[c]).exp_().sum(dim=-1, keepdim=True) + if tp_group is not None: + torch.distributed.all_reduce(sum_exp, group=tp_group) + shape = (*logits.shape[:-1], -1) + return (lse + sum_exp.log()).view(shape), top_vals.view(shape), top_idx.view(shape) + + +class _ScaledLogSumExpAndGather(torch.autograd.Function): + """``_scaled_logsumexp(logits)`` and ``(logits / temperature).gather(-1, index)`` in fp32. + + Saves only ``logits`` and writes their gradient chunk by chunk, instead of keeping fp32 copies + of the full logits alive for backward. + """ + + @staticmethod + def forward(ctx, logits, index, temperature, tp_group): + lse, _, _ = _scaled_logsumexp(logits, temperature, tp_group) + ctx.save_for_backward(logits, index, lse) + ctx.temperature, ctx.tp_group = temperature, tp_group + return lse, logits.gather(-1, index).float() / temperature + + @staticmethod + def backward(ctx, grad_lse, grad_vals): + logits, index, lse = ctx.saved_tensors + if ctx.tp_group is not None: + # As in the autograd all-reduce of the normalizer, every TP rank's loss reaches each shard. + grad_lse = grad_lse.clone() + torch.distributed.all_reduce(grad_lse, group=ctx.tp_group) + rows = logits.reshape(-1, logits.size(-1)) + grad_lse, lse = grad_lse.reshape(-1, 1), lse.reshape(-1, 1) + index, grad_vals = index.reshape(-1, index.size(-1)), grad_vals.reshape(-1, index.size(-1)) + grad = torch.empty_like(rows) + for c in _row_chunks(rows): + # softmax(x / T) * grad_lse, plus grad_vals scattered to the gathered entries, over T. + g = rows[c].to(torch.float32, copy=True).div_(ctx.temperature).sub_(lse[c]).exp_() + g.mul_(grad_lse[c]).scatter_add_(-1, index[c], grad_vals[c]) + grad[c] = g.div_(ctx.temperature) + return grad.view(logits.shape), None, None, None + + class TopLogitsKLLoss(LogitsKLLoss): """Calculates KL-Divergence loss restricted to the Teacher's Top-K (and optionally Top-P) entries. Calculates using the global Top-K entries without gathering full logits. NOTE: Will gather Top-K logits per rank, so mind the value of K for communication. The full-vocab - normalizers still allocate fp32 copies of the local logit shards (the teacher's is freed right - away; the student's is retained for backward), and add two TP all-reduces per distribution. + normalizers run over fp32 row chunks of the local logit shards rather than fp32 copies of them, + and add two TP all-reduces per distribution. Both distributions are normalized over the *full* vocabulary (not re-normalized over the Top-K), matching the offline cached-logits KD loss in Megatron-LM. A "ghost" token holding the @@ -471,21 +539,16 @@ def forward(self, predictions: Tensor, targets: Tensor) -> Tensor: # We take K from each rank and then find the global Top-K of all those. local_top_k = min(self.top_k, targets.size(-1)) - # Teacher: full-vocab normalizer and local Top-K, then free its fp32 copy before the student's. - output_teacher = targets.float() / self._temperature - teacher_lse = self._tp_logsumexp(output_teacher) - top_teacher_vals, top_idx = torch.topk(output_teacher, local_top_k, dim=-1) - del output_teacher - - # Student: the full-vocab normalizer is inherent to the ghost-token formulation. - output_student = predictions.float() / self._temperature - student_lse = self._tp_logsumexp(output_student) - top_student_vals = torch.gather(output_student, dim=-1, index=top_idx) - del output_student + # Full-vocab normalizers (inherent to the ghost-token formulation) and the local Top-K. + tp_group = parallel_state.get_tensor_model_parallel_group() if tp_size > 1 else None + teacher_lse, top_teacher_vals, top_idx = _scaled_logsumexp( + targets, self._temperature, tp_group, local_top_k + ) + student_lse, top_student_vals = _ScaledLogSumExpAndGather.apply( + predictions, top_idx, self._temperature, tp_group + ) if tp_size > 1: - tp_group = parallel_state.get_tensor_model_parallel_group() - # Gather all candidates into shape [N_rows, local_k * tp_size] # Use all_gather from torch.distributed.nn.functional to preserve gradients all_teacher_vals = dist_nn.functional.all_gather( diff --git a/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py b/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py index 18cc2807a8e..6ce24c43f82 100644 --- a/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py +++ b/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py @@ -17,6 +17,7 @@ import warnings from functools import partial from types import SimpleNamespace +from unittest.mock import patch import pytest import torch @@ -28,6 +29,7 @@ from megatron.core import parallel_state import modelopt.torch.distill as mtd +import modelopt.torch.distill.plugins.megatron as megatron_plugin from modelopt.torch.distill.plugins.megatron import ( DistillationConfig, LogitsAndIntermediatesLossBalancer, @@ -479,9 +481,13 @@ def _test_topk_logits_kl_loss_numerics_ghost_token_reference( student_shard = student.chunk(size, dim=-1)[rank].detach().contiguous().requires_grad_() teacher_shard = teacher.chunk(size, dim=-1)[rank].contiguous() k = 4 - loss = TopLogitsKLLoss(cfg, top_k=k, top_p=top_p, reverse=reverse, temperature=temperature)( - student_shard, teacher_shard - )[0] + # Chunks of 2 rows, so the loss's chunked full-vocab passes take several chunks. + with patch.object(megatron_plugin, "_TOP_LOGITS_CHUNK_NUMEL", 2 * student_shard.size(-1)): + loss_fn = TopLogitsKLLoss( + cfg, top_k=k, top_p=top_p, reverse=reverse, temperature=temperature + ) + loss = loss_fn(student_shard, teacher_shard)[0] + loss.mean().backward() q_full = F.log_softmax(teacher / temperature, dim=-1) p_full = F.log_softmax(student / temperature, dim=-1) @@ -504,7 +510,6 @@ def _test_topk_logits_kl_loss_numerics_ghost_token_reference( p, q = q, p ref = (keep * q.exp() * (q - p)).sum(-1).transpose(0, 1) torch.testing.assert_close(loss, ref, atol=1e-6, rtol=1e-5) - loss.mean().backward() ref.mean().backward() torch.testing.assert_close( student_shard.grad, student.grad.chunk(size, dim=-1)[rank], atol=1e-6, rtol=1e-5 @@ -526,6 +531,19 @@ def test_topk_logits_kl_loss_numerics_ghost_token_reference( ) +def test_topk_logits_kl_loss_memory(): + """Forward and backward add about one gradient the size of the logits, not fp32 copies of them.""" + cfg = SimpleNamespace(tensor_model_parallel_size=1) + student = torch.randn(1024, 2, 32768, device="cuda", dtype=torch.bfloat16, requires_grad=True) + teacher = torch.randn_like(student, requires_grad=False) + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + base = torch.cuda.memory_allocated() + with patch.object(megatron_plugin, "_TOP_LOGITS_CHUNK_NUMEL", 64 * student.size(-1)): + TopLogitsKLLoss(cfg, top_k=64)(student, teacher)[0].sum().backward() + assert torch.cuda.max_memory_allocated() - base < 1.5 * student.nbytes + + @pytest.mark.parametrize("temperature", [0.5, 2.0, 3.7]) def test_logits_kl_losses_temperature_scaling(temperature): """Dense and Top-K losses match a plain ``log_softmax(x / T)`` reference at T != 1.""" From c6d9e538d9c8f7d6dc9cd3f6308d07bceaacf420 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Wed, 28 Oct 2026 01:45:51 -0700 Subject: [PATCH 2/3] Keep the Top-K KD loss exact for logits with a large common offset Work relative to the row max: the normalizer and the Top-K values are both shifted by it, so it cancels out of the log-probs, and backward rebuilds the softmax from the shifted logits instead of exp(x - lse) with a rounded lse. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- modelopt/torch/distill/plugins/megatron.py | 49 ++++++++++--------- .../distill/plugins/test_distill_megatron.py | 16 ++++++ 2 files changed, 42 insertions(+), 23 deletions(-) diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py index adfb41bfde7..0d84b57479b 100644 --- a/modelopt/torch/distill/plugins/megatron.py +++ b/modelopt/torch/distill/plugins/megatron.py @@ -414,59 +414,61 @@ def _row_chunks(rows: Tensor) -> list[slice]: def _scaled_logsumexp( logits: Tensor, temperature: float, tp_group, top_k: int = 0 -) -> tuple[Tensor, Tensor, Tensor]: - """fp32 ``logsumexp(logits / temperature)`` over the TP-sharded vocab dim, by row chunks. +) -> tuple[Tensor, Tensor, Tensor, Tensor]: + """fp32 row max ``m`` and ``logsumexp(x - m)`` of ``x = logits / temperature``, by row chunks. - Also returns the local Top-K values and indices of ``logits / temperature`` (empty if ``top_k=0``). + Both are over the TP-sharded vocab dim. Also returns the local Top-K values of ``x - m`` and their + indices (empty if ``top_k=0``). Working relative to ``m`` keeps the results exact for logits with + a large common offset. """ rows = logits.detach().reshape(-1, logits.size(-1)) # Rounding is monotonic, so this is also the max of the scaled fp32 logits. - lse = rows.amax(dim=-1, keepdim=True).float() / temperature + row_max = rows.amax(dim=-1, keepdim=True).float() / temperature if tp_group is not None: - torch.distributed.all_reduce(lse, op=torch.distributed.ReduceOp.MAX, group=tp_group) - sum_exp = torch.empty_like(lse) - top_vals = lse.new_empty(rows.size(0), top_k) + torch.distributed.all_reduce(row_max, op=torch.distributed.ReduceOp.MAX, group=tp_group) + sum_exp = torch.empty_like(row_max) + top_vals = row_max.new_empty(rows.size(0), top_k) top_idx = torch.empty(rows.size(0), top_k, dtype=torch.long, device=rows.device) for c in _row_chunks(rows): - scaled = rows[c].to(torch.float32, copy=True).div_(temperature) + shifted = rows[c].to(torch.float32, copy=True).div_(temperature).sub_(row_max[c]) if top_k: - top_vals[c], top_idx[c] = torch.topk(scaled, top_k, dim=-1) - sum_exp[c] = scaled.sub_(lse[c]).exp_().sum(dim=-1, keepdim=True) + top_vals[c], top_idx[c] = torch.topk(shifted, top_k, dim=-1) + sum_exp[c] = shifted.exp_().sum(dim=-1, keepdim=True) if tp_group is not None: torch.distributed.all_reduce(sum_exp, group=tp_group) shape = (*logits.shape[:-1], -1) - return (lse + sum_exp.log()).view(shape), top_vals.view(shape), top_idx.view(shape) + return tuple(t.view(shape) for t in (row_max, sum_exp.log(), top_vals, top_idx)) class _ScaledLogSumExpAndGather(torch.autograd.Function): - """``_scaled_logsumexp(logits)`` and ``(logits / temperature).gather(-1, index)`` in fp32. + """``logsumexp(x - m)`` and ``(x - m).gather(-1, index)`` of ``_scaled_logsumexp`` in fp32. Saves only ``logits`` and writes their gradient chunk by chunk, instead of keeping fp32 copies - of the full logits alive for backward. + of the full logits alive for backward. ``m`` cancels out of the log-probs, so it gets no gradient. """ @staticmethod def forward(ctx, logits, index, temperature, tp_group): - lse, _, _ = _scaled_logsumexp(logits, temperature, tp_group) - ctx.save_for_backward(logits, index, lse) + row_max, lse, _, _ = _scaled_logsumexp(logits, temperature, tp_group) + ctx.save_for_backward(logits, index, row_max, lse) ctx.temperature, ctx.tp_group = temperature, tp_group - return lse, logits.gather(-1, index).float() / temperature + return lse, logits.gather(-1, index).float() / temperature - row_max @staticmethod def backward(ctx, grad_lse, grad_vals): - logits, index, lse = ctx.saved_tensors + logits, index, row_max, lse = ctx.saved_tensors if ctx.tp_group is not None: # As in the autograd all-reduce of the normalizer, every TP rank's loss reaches each shard. grad_lse = grad_lse.clone() torch.distributed.all_reduce(grad_lse, group=ctx.tp_group) rows = logits.reshape(-1, logits.size(-1)) - grad_lse, lse = grad_lse.reshape(-1, 1), lse.reshape(-1, 1) + row_max, lse, grad_lse = row_max.reshape(-1, 1), lse.reshape(-1, 1), grad_lse.reshape(-1, 1) index, grad_vals = index.reshape(-1, index.size(-1)), grad_vals.reshape(-1, index.size(-1)) grad = torch.empty_like(rows) for c in _row_chunks(rows): - # softmax(x / T) * grad_lse, plus grad_vals scattered to the gathered entries, over T. - g = rows[c].to(torch.float32, copy=True).div_(ctx.temperature).sub_(lse[c]).exp_() - g.mul_(grad_lse[c]).scatter_add_(-1, index[c], grad_vals[c]) + # softmax(x) * grad_lse, plus grad_vals scattered to the gathered entries, over T. + g = rows[c].to(torch.float32, copy=True).div_(ctx.temperature).sub_(row_max[c]) + g.sub_(lse[c]).exp_().mul_(grad_lse[c]).scatter_add_(-1, index[c], grad_vals[c]) grad[c] = g.div_(ctx.temperature) return grad.view(logits.shape), None, None, None @@ -539,9 +541,10 @@ def forward(self, predictions: Tensor, targets: Tensor) -> Tensor: # We take K from each rank and then find the global Top-K of all those. local_top_k = min(self.top_k, targets.size(-1)) - # Full-vocab normalizers (inherent to the ghost-token formulation) and the local Top-K. + # Full-vocab normalizers (inherent to the ghost-token formulation) and the local Top-K, both + # relative to each distribution's row max, which cancels out of the log-probs. tp_group = parallel_state.get_tensor_model_parallel_group() if tp_size > 1 else None - teacher_lse, top_teacher_vals, top_idx = _scaled_logsumexp( + _, teacher_lse, top_teacher_vals, top_idx = _scaled_logsumexp( targets, self._temperature, tp_group, local_top_k ) student_lse, top_student_vals = _ScaledLogSumExpAndGather.apply( diff --git a/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py b/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py index 6ce24c43f82..375e0777a43 100644 --- a/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py +++ b/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py @@ -544,6 +544,22 @@ def test_topk_logits_kl_loss_memory(): assert torch.cuda.max_memory_allocated() - base < 1.5 * student.nbytes +def test_topk_logits_kl_loss_offset_invariant(): + """A common logit offset far larger than the logits changes neither the loss nor the gradient.""" + cfg = SimpleNamespace(tensor_model_parallel_size=1) + torch.manual_seed(SEED) + # Multiples of 1/8 below 16 stay exact in fp32 after adding 2**20. + student = torch.randint(-64, 64, (4, 3, 16)) / 8 + teacher = torch.randint(-64, 64, (4, 3, 16)) / 8 + results = [] + for offset in (0.0, 2.0**20): + shifted = (student + offset).requires_grad_() + loss = TopLogitsKLLoss(cfg, top_k=4)(shifted, teacher + offset)[0] + loss.sum().backward() + results.append((loss.detach(), shifted.grad)) + torch.testing.assert_close(results[1], results[0]) + + @pytest.mark.parametrize("temperature", [0.5, 2.0, 3.7]) def test_logits_kl_losses_temperature_scaling(temperature): """Dense and Top-K losses match a plain ``log_softmax(x / T)`` reference at T != 1.""" From 685afb47df7acbbf4c70e667ec1b8f4f60d337f4 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Wed, 28 Oct 2026 09:59:47 -0700 Subject: [PATCH 3/3] Keep the Top-K KD loss working for empty batch or sequence inputs Reshape the normalizer results with explicit widths, since -1 cannot be inferred for zero-element tensors, and add forward/backward coverage. Reword the changelog entry: results now differ from before for logits with a large common offset. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- CHANGELOG.rst | 2 +- modelopt/torch/distill/plugins/megatron.py | 6 ++++-- .../torch/distill/plugins/test_distill_megatron.py | 11 +++++++++++ 3 files changed, 16 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index d9de64cbf67..788b1c1d9c0 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -75,7 +75,7 @@ Changelog - Fix ONNX quantization with the default ``concat_elimination`` pass crashing in calibration with ``KeyError: '_ReduceMin'`` when a ``Concat`` input is a constant, such as the class token of a ViT. - Fix an illegal memory access in static NVFP4 fake quantization for tensors with more than 2^31 elements, such as stacked MoE expert weights. The kernel is also about 14x faster. - Fix Megatron Top-K/Top-P distillation gradients being multiplied by the tensor-parallel size when tensor parallelism is enabled; reported loss values are unchanged. -- Fix Megatron Top-K/Top-P distillation running out of memory at long sequence lengths with large vocabularies: since 0.48 its full-vocabulary normalizers kept several float32 copies of the logits alive. It now adds about one copy of the logits in their own dtype, with identical loss and gradients. +- Fix Megatron Top-K/Top-P distillation running out of memory at long sequence lengths with large vocabularies: since 0.48 its full-vocabulary normalizers kept several float32 copies of the logits alive. It now adds about one copy of the logits in their own dtype, and stays accurate for logits with a large common offset. - Fix the Puzzletron memory objective pricing the KV cache with ``hidden_size // num_attention_heads`` instead of the model's ``head_dim``, which misestimated it for models such as Mistral-Small, GPT-OSS and Gemma-3 and for every candidate with a pruned hidden size. Recompute subblock statistics produced with earlier versions. - Fix ONNX INT4 quantization ending the Python process when the INT4 packing extension fails to build, for example without a C++ compiler. It now falls back to the Python implementation. - Fix ``fsdp2_aware_weight_update`` masking an error raised during its unshard (e.g. a CUDA OOM) with a spurious ``UnboundLocalError``. diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py index 0d84b57479b..0126e3f6065 100644 --- a/modelopt/torch/distill/plugins/megatron.py +++ b/modelopt/torch/distill/plugins/megatron.py @@ -436,8 +436,10 @@ def _scaled_logsumexp( sum_exp[c] = shifted.exp_().sum(dim=-1, keepdim=True) if tp_group is not None: torch.distributed.all_reduce(sum_exp, group=tp_group) - shape = (*logits.shape[:-1], -1) - return tuple(t.view(shape) for t in (row_max, sum_exp.log(), top_vals, top_idx)) + # Explicit widths: a -1 cannot be inferred for empty batch or sequence dims. + return tuple( + t.view(*logits.shape[:-1], t.size(-1)) for t in (row_max, sum_exp.log(), top_vals, top_idx) + ) class _ScaledLogSumExpAndGather(torch.autograd.Function): diff --git a/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py b/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py index 375e0777a43..57a841adc48 100644 --- a/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py +++ b/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py @@ -560,6 +560,17 @@ def test_topk_logits_kl_loss_offset_invariant(): torch.testing.assert_close(results[1], results[0]) +@pytest.mark.parametrize("shape", [(0, 3, 16), (4, 0, 16)]) +def test_topk_logits_kl_loss_empty_inputs(shape): + """Empty batch or sequence inputs give an empty loss and gradient.""" + cfg = SimpleNamespace(tensor_model_parallel_size=1) + student = torch.randn(shape, requires_grad=True) + loss = TopLogitsKLLoss(cfg, top_k=4)(student, torch.randn(shape))[0] + loss.sum().backward() + assert loss.shape == (shape[1], shape[0]) + assert student.grad.shape == shape + + @pytest.mark.parametrize("temperature", [0.5, 2.0, 3.7]) def test_logits_kl_losses_temperature_scaling(temperature): """Dense and Top-K losses match a plain ``log_softmax(x / T)`` reference at T != 1."""