diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 66f4015c760..788b1c1d9c0 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, 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 63745ed48e0..0126e3f6065 100644 --- a/modelopt/torch/distill/plugins/megatron.py +++ b/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( 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..57a841adc48 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,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."""