Fix the Top-K KD loss running out of memory at long sequences with large vocabularies

CHANGELOG.rst

@@ -75,6 +75,7 @@ Changelog
 - Fix ONNX quantization with the default ``concat_elimination`` pass crashing in calibration with ``KeyError: '<tensor>_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, 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``.

modelopt/torch/distill/plugins/megatron.py

@@ -403,13 +403,85 @@ 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, Tensor]:
+    """fp32 row max ``m`` and ``logsumexp(x - m)`` of ``x = logits / temperature``, by row chunks.
+
+    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.
+    row_max = rows.amax(dim=-1, keepdim=True).float() / temperature
+    if tp_group is not None:
+        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):
+        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(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)
+    # 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):
+    """``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. ``m`` cancels out of the log-probs, so it gets no gradient.
+    """
+
+    @staticmethod
+    def forward(ctx, logits, index, temperature, tp_group):
+        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 - row_max
+
+    @staticmethod
+    def backward(ctx, grad_lse, grad_vals):
+        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))
+        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) * 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
+
+
 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 +543,17 @@ 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, 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(
+            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(

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,46 @@ 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
+
+
+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("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."""