From b4112bdf82978c4318e74990c50d71d489a948d3 Mon Sep 17 00:00:00 2001 From: Asha Anoosheh Date: Thu, 17 Sep 2026 19:01:48 +0200 Subject: [PATCH 1/3] Bring offline KD upgrades such as Ghost Token and Top-P to Megatron KD plugin Signed-off-by: Asha Anoosheh --- CHANGELOG.rst | 4 + examples/megatron_bridge/distill.py | 42 +++- modelopt/torch/distill/plugins/megatron.py | 214 ++++++++++++------ .../distill/plugins/test_distill_megatron.py | 190 +++++++++++++++- 4 files changed, 376 insertions(+), 74 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 46168b7fe0f..8b78b03a004 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -15,6 +15,7 @@ Changelog *Megatron Framework (M-LM / M-Bridge)* +- Add optional Top-P (nucleus) truncation to the Megatron ``TopKLogitsKLLoss`` via ``logit_kl_top_p`` and ``logit_kl_top_p_min_k`` in ``DistillationConfig``: after the global Top-K selection, only the smallest prefix whose cumulative teacher probability reaches ``top_p`` (with a floor of ``min_k`` entries) contributes to the KL, mirroring ``--logits-save-top-p`` / ``--logits-save-top-p-min-k`` in Megatron-LM's logits saver. - Add an end-to-end W4A4 NVFP4 PTQ and QAD tutorial for Qwen3.6-35B-A3B also covering evaluation and vLLM throughput benchmarking. See `examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/README.md `_ for details. *Misc* @@ -23,6 +24,8 @@ Changelog **Backward Breaking Changes** +- ``modelopt.torch.distill.plugins.megatron.TopKLogitsKLLoss`` (``logit_kl_topk`` in ``DistillationConfig``) now normalizes both distributions over the full vocabulary instead of re-normalizing over the Top-K entries, and by default appends a "ghost" token holding the probability mass outside the Top-K to both student and teacher (matching Megatron-LM's offline cached-logits KD loss). Loss values change for existing ``logit_kl_topk`` runs; set ``logit_kl_ghost_token: false`` to drop the ghost token. +- ``LogitsAndIntermediatesLossBalancer`` (Megatron distillation plugin) no longer rescales the distillation loss to the magnitude of the LM loss. The total is now the fixed convex combination ``(1 - alpha) * lm_loss + alpha * kd_loss`` with ``DistillationConfig.kd_loss_alpha`` (default ``0.9``, in [0, 1]), matching Megatron-LM's offline cached-logits KD. ``skip_lm_loss`` is now derived from ``kd_loss_alpha`` (skipped iff ``1.0``), so the LM loss is computed by default where it was previously skipped. ``examples/megatron_bridge/distill.py`` gains ``--kd_loss_alpha``. - Layerwise calibration now uses prior-layer QDQ activations by default (``layerwise.get_qdq_activations_from_prev_layer=True``). Set it to ``False`` to preserve full-precision activations for subsequent layers (the default behavior for @@ -32,6 +35,7 @@ Changelog **Deprecations** +- ``DistillationConfig.kd_loss_scale`` and ``DistillationConfig.skip_lm_loss`` (Megatron distillation plugin) are deprecated. ``kd_loss_scale`` is ignored with a ``DeprecationWarning``; a user-provided ``skip_lm_loss`` is overridden by the value derived from ``kd_loss_alpha`` with a ``DeprecationWarning``. The ``--no_skip_lm_loss`` and ``--kd_loss_scale`` flags in ``examples/megatron_bridge/distill.py`` are likewise deprecated and ignored. - Rename the architecture-specific recipe tier from ``modelopt_recipes/huggingface/`` to ``modelopt_recipes/model_type/`` to clarify that it holds recipes shared across every checkpoint of a Hugging Face ``model_type``. Saved ``--recipe huggingface//...`` paths still resolve via a backward-compatibility alias but now emit a ``FutureWarning``, so update them to ``model_type//...`` as the ``huggingface/`` prefix is deprecated. - The single-format quantization CLI flags are deprecated in favour of ``--recipe`` and will be removed in a future release; passing one now emits a ``FutureWarning``. ``examples/hf_ptq``: ``--qformat`` and ``--kv_cache_qformat``. ``examples/megatron_bridge/quantize.py``: ``--quant_cfg``, ``--kv_cache_quant`` and ``--weight_only``. ``examples/torch_onnx/torch_quant_to_onnx.py``: ``--qformat``. A recipe carries the quantization config, the calibration algorithm and the KV-cache setting in one file, so they cannot drift apart the way separate flags can -- and ``--recipe`` already took precedence over all six, silently on ``hf_ptq`` and with a warning on ``megatron_bridge`` -- with one gap the recipe closes rather than inherits: a weight AutoQuantize recipe that omits ``kv_cache`` still falls back to ``--kv_cache_qformat``, so set ``kv_cache`` in the recipe when migrating. Use a recipe from ``modelopt_recipes/general/ptq/``, an architecture-specific one under ``modelopt_recipes/model_type//``, or a checkpoint-specific one under ``modelopt_recipes/models/``. The warning fires only when a flag is passed explicitly: ``--qformat`` defaults to ``fp8`` and ``--kv_cache_qformat`` to ``fp8_cast``, so warning on the defaults would fire on every run, including runs that correctly use ``--recipe``. ``examples/speculative_decoding/scripts/quantize_drafter.py`` keeps ``--qformat`` undeprecated: it has no ``--recipe`` alternative yet. - The TensorRT-LLM checkpoint export format is deprecated and will be removed in 0.49.0: ``export_tensorrt_llm_checkpoint`` and ``torch_to_tensorrt_llm_checkpoint`` now emit a ``DeprecationWarning`` on use. Use ``export_hf_checkpoint``, which exports a unified Hugging Face checkpoint deployable on TensorRT-LLM, vLLM and SGLang. Its implementation moved to ``modelopt.torch.export.trtllm``, so import those two functions from there and the ``ModelConfig`` dataclasses from ``modelopt.torch.export.trtllm.model_config``; both functions remain importable from ``modelopt.torch.export`` for this release only. diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index e1810c19d5e..89b44e6bc5f 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -172,9 +172,23 @@ def get_args(): "--train_iters", type=int, required=True, help="Number of training iterations" ) parser.add_argument( - "--no_skip_lm_loss", action="store_true", help="Disable skipping language model loss" + "--no_skip_lm_loss", + action="store_true", + help="DEPRECATED and ignored. Whether the LM loss is skipped is derived from --kd_loss_alpha " + "(skipped iff alpha == 1.0).", + ) + parser.add_argument( + "--kd_loss_alpha", + type=float, + default=0.9, + help="KD loss weight alpha in (1 - alpha) * lm_loss + alpha * kd_loss. 1.0 skips the LM loss entirely.", + ) + parser.add_argument( + "--kd_loss_scale", + type=float, + default=None, + help="DEPRECATED and ignored. Use --kd_loss_alpha.", ) - parser.add_argument("--kd_loss_scale", type=float, default=1.0, help="KD loss weight") parser.add_argument( "--no_async_save", action="store_true", @@ -188,6 +202,24 @@ def get_args(): help="Restrict the logit KL loss to the teacher's top-k vocabulary entries, " "replacing the full-vocab temporaries with [seq, k] ones.", ) + parser.add_argument( + "--logit_kl_top_p", + type=float, + default=None, + help="Nucleus threshold in (0, 1] applied on top of --logit_kl_topk: only the smallest prefix " + "of the sorted top-k whose cumulative teacher probability reaches this value is distilled.", + ) + parser.add_argument( + "--logit_kl_top_p_min_k", + type=int, + default=1, + help="Minimum number of top-k entries kept per token when --logit_kl_top_p is active.", + ) + parser.add_argument( + "--no_logit_kl_ghost_token", + action="store_true", + help="Disable the residual 'ghost' token (out-of-top-k probability mass) in the top-k KL loss.", + ) parser.add_argument("--lr", type=float, default=1e-4, help="Peak learning rate") parser.add_argument("--min_lr", type=float, default=1e-5, help="Minimum learning rate") parser.add_argument("--lr_warmup_iters", type=int, default=50, help="Number of LR warmup steps") @@ -435,9 +467,11 @@ def _build_model_provider(hf_path, load_weights=True, moe_grouped_gemm=True): ) kd_config = ModelOptDistillConfig( - skip_lm_loss=not args.no_skip_lm_loss, - kd_loss_scale=args.kd_loss_scale, + kd_loss_alpha=args.kd_loss_alpha, logit_kl_topk=args.logit_kl_topk, + logit_kl_top_p=args.logit_kl_top_p, + logit_kl_top_p_min_k=args.logit_kl_top_p_min_k, + logit_kl_ghost_token=not args.no_logit_kl_ghost_token, ) # HF VLM configs expose ``vision_config``; Megatron-Bridge nests the text model under diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py index c93f0961d1f..c6c5af8ab07 100644 --- a/modelopt/torch/distill/plugins/megatron.py +++ b/modelopt/torch/distill/plugins/megatron.py @@ -19,6 +19,7 @@ import logging import re +import warnings from abc import ABCMeta from collections.abc import Callable from dataclasses import dataclass, field @@ -56,18 +57,32 @@ class DistillationConfig: Args: intermediate_layer_pairs: List of tuples of intermediate layer names. logit_layers: Tuple of logit layer names. - skip_lm_loss: Whether to skip computing the standard language model loss (default: ``True``). - kd_loss_scale: Relative scaling factor for the distillation loss if ``skip_lm_loss`` is ``False``. + kd_loss_alpha: Weight of the distillation loss in the convex combination + ``(1 - alpha) * lm_loss + alpha * kd_loss``. Must be in [0, 1]. When ``1.0``, the standard + language model loss is skipped entirely (``skip_lm_loss`` is derived from this value). + skip_lm_loss: DEPRECATED. Derived from ``kd_loss_alpha`` (``True`` iff ``kd_loss_alpha == 1.0``); + any user-provided value is overridden with a warning. + kd_loss_scale: DEPRECATED and ignored. Use ``kd_loss_alpha`` instead. logit_kl_temperature: Temperature for the logit KL-divergence loss. logit_kl_topk: If not None, use TopKLogitsKLLoss instead of LogitsKLLoss with this top-k value. + logit_kl_top_p: Optional nucleus (top-P) threshold applied on top of the teacher's Top-K. + Only the smallest prefix of the (sorted) Top-K whose cumulative teacher probability + reaches this value contributes to the loss. Requires ``logit_kl_topk``. Must be in (0, 1]. + logit_kl_top_p_min_k: Minimum number of Top-K entries kept per token when top-P is active. + logit_kl_ghost_token: Whether ``TopKLogitsKLLoss`` appends a "ghost" token holding the + probability mass outside the kept entries to both distributions (default: ``True``). """ intermediate_layer_pairs: list[tuple[str, ...]] = field(default_factory=list) logit_layers: tuple[str, str] = ("output_layer", "output_layer") - skip_lm_loss: bool = True - kd_loss_scale: float = 1.0 + kd_loss_alpha: float = 0.9 + skip_lm_loss: bool | None = None # deprecated, derived from kd_loss_alpha + kd_loss_scale: float | None = None # deprecated, ignored logit_kl_temperature: float = 1.0 logit_kl_topk: int | None = None + logit_kl_top_p: float | None = None + logit_kl_top_p_min_k: int = 1 + logit_kl_ghost_token: bool = True criterion: Criterion | None = None loss_balancer: mtd.DistillationLossBalancer | None = None @@ -76,8 +91,30 @@ def __post_init__(self): assert all(len(pair) in (2, 3) for pair in self.intermediate_layer_pairs), ( f"{self.intermediate_layer_pairs=}" ) - assert self.kd_loss_scale > 0, f"{self.kd_loss_scale=}" + assert 0 <= self.kd_loss_alpha <= 1, f"{self.kd_loss_alpha=}" + if self.kd_loss_scale is not None: + warnings.warn( + "DistillationConfig.kd_loss_scale is deprecated and ignored. The distillation loss " + "is no longer rescaled to the LM loss magnitude; the total loss is now " + "(1 - kd_loss_alpha) * lm_loss + kd_loss_alpha * kd_loss. Set `kd_loss_alpha` instead.", + DeprecationWarning, + stacklevel=2, + ) + derived_skip_lm_loss = self.kd_loss_alpha == 1.0 + if self.skip_lm_loss is not None: + warnings.warn( + "DistillationConfig.skip_lm_loss is deprecated and is now derived from `kd_loss_alpha` " + f"(skip iff kd_loss_alpha == 1.0). Overriding skip_lm_loss={self.skip_lm_loss} with " + f"{derived_skip_lm_loss} (kd_loss_alpha={self.kd_loss_alpha}).", + DeprecationWarning, + stacklevel=2, + ) + self.skip_lm_loss = derived_skip_lm_loss assert self.logit_kl_temperature > 0, f"{self.logit_kl_temperature=}" + if self.logit_kl_top_p is not None: + assert self.logit_kl_topk is not None, "logit_kl_top_p requires logit_kl_topk" + assert 0 < self.logit_kl_top_p <= 1, f"{self.logit_kl_top_p=}" + assert self.logit_kl_top_p_min_k >= 1, f"{self.logit_kl_top_p_min_k=}" @staticmethod def parse_intermediate_entry(entry: tuple[str, ...]) -> tuple[str, str, Callable]: @@ -130,7 +167,12 @@ def setup_distillation_config( # Use TopKLogitsKLLoss if logit_kl_topk is specified, otherwise use LogitsKLLoss if cfg.logit_kl_topk is not None: criterion[tuple(cfg.logit_layers)] = TopKLogitsKLLoss( - student_cfg, temperature=cfg.logit_kl_temperature, top_k=cfg.logit_kl_topk + student_cfg, + temperature=cfg.logit_kl_temperature, + top_k=cfg.logit_kl_topk, + top_p=cfg.logit_kl_top_p, + top_p_min_k=cfg.logit_kl_top_p_min_k, + add_ghost_token=cfg.logit_kl_ghost_token, ) else: criterion[tuple(cfg.logit_layers)] = LogitsKLLoss( @@ -156,7 +198,8 @@ def setup_distillation_config( if cfg.loss_balancer is None: cfg.loss_balancer = LogitsAndIntermediatesLossBalancer( - kd_loss_scale=cfg.kd_loss_scale, skip_original_loss=cfg.skip_lm_loss + kd_loss_alpha=cfg.kd_loss_alpha, + skip_original_loss=bool(cfg.skip_lm_loss), # always set by __post_init__ ) return cfg @@ -319,56 +362,47 @@ def forward(self, predictions: Tensor, targets: Tensor) -> Tensor: """ predictions, targets = self.pre_forward(predictions, targets) - # Division by temp should happen prior to finding max for both student and teacher. - output_teacher = targets.float() / self._temperature - output_student = predictions.float() / self._temperature + # Temperature-scaled log probabilities (log softmax), globally normalized across TP vocab shards. + p = predictions.float() / self._temperature - self._tp_logsumexp(predictions) + q = targets.float() / self._temperature - self._tp_logsumexp(targets) - # Compute local softmax, and the reweight to compute global softmax. - if self._config.tensor_model_parallel_size > 1: - tp_group = parallel_state.get_tensor_model_parallel_group() + # KL divergence + if self._reverse: + p, q = q, p + loss = torch.sum(F.kl_div(p, q, reduction="none", log_target=True), dim=-1) - # Subtract maximum value along vocab dimension across all GPUs (for stability) - teacher_logits_max, _ = torch.max(output_teacher, dim=-1, keepdim=True) - torch.distributed.all_reduce( - teacher_logits_max, - op=torch.distributed.ReduceOp.MAX, - group=tp_group, - ) - output_teacher -= teacher_logits_max + return self.post_forward(loss, tp_reduce=True) - student_logits_max, _ = torch.max(output_student, dim=-1, keepdim=True) + def _tp_logsumexp(self, logits: Tensor, num_chunks: int = 4) -> Tensor: + """Log-sum-exp of ``logits / self._temperature`` over the vocab dim across all TP shards. + + Accumulates in fp32 over ``num_chunks`` vocab chunks so the transient fp32 working set is a + fraction of the (possibly bf16) full-vocab input. NOTE: for inputs requiring grad, autograd + still retains each chunk's ``exp`` output for backward, so the total saved activation + equals a full-vocab fp32 tensor. Returns shape ``[..., 1]``. + """ + # Max is exact under the monotonic temperature scaling, so take it in the native dtype + # and defer the temperature division to the centered values: lse(x/T) = max/T + log(sum(exp((x-max)/T))). + logits_max = logits.amax(dim=-1, keepdim=True).float() + if self._config.tensor_model_parallel_size > 1: + tp_group = parallel_state.get_tensor_model_parallel_group() torch.distributed.all_reduce( - student_logits_max, - op=torch.distributed.ReduceOp.MAX, - group=tp_group, + logits_max, op=torch.distributed.ReduceOp.MAX, group=tp_group ) - output_student -= student_logits_max.detach() - - # Compute global softmax denominators + logits_max = logits_max.detach() + + denom = None + chunk_size = -(-logits.size(-1) // num_chunks) # ceil division + for chunk in logits.split(chunk_size, dim=-1): + centered = (chunk.float() - logits_max) / self._temperature + partial = torch.exp(centered).sum(dim=-1, keepdim=True) + denom = partial if denom is None else denom + partial + if self._config.tensor_model_parallel_size > 1: # We can't use standard all_reduce function here since the computation # that follows it isn't identical across TP ranks. - denom_teacher = torch.sum(torch.exp(output_teacher), dim=-1, keepdim=True) - denom_teacher = dist_nn.functional.all_reduce(denom_teacher, group=tp_group) + denom = dist_nn.functional.all_reduce(denom, group=tp_group) - denom_student = torch.sum(torch.exp(output_student), dim=-1, keepdim=True) - denom_student = dist_nn.functional.all_reduce(denom_student, group=tp_group) - - # Compute log probabilities (log softmax) - teacher_log_prob = output_teacher - torch.log(denom_teacher) - student_log_prob = output_student - torch.log(denom_student) - - # KL divergence - p, q = student_log_prob, teacher_log_prob - else: - # Compute log probabilities - p, q = F.log_softmax(output_student, dim=-1), F.log_softmax(output_teacher, dim=-1) - - # KL divergence - if self._reverse: - p, q = q, p - loss = torch.sum(F.kl_div(p, q, reduction="none", log_target=True), dim=-1) - - return self.post_forward(loss, tp_reduce=True) + return logits_max / self._temperature + torch.log(denom) class TopKLogitsKLLoss(LogitsKLLoss): @@ -376,6 +410,16 @@ class TopKLogitsKLLoss(LogitsKLLoss): 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 memory and communication. + + 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. Optional refinements: + + * **Top-P (nucleus)**: after sorting the Top-K by teacher probability, only the smallest prefix + whose cumulative teacher mass reaches ``top_p`` (with a floor of ``top_p_min_k`` entries) + contributes to the loss. + * **Ghost token**: a synthetic extra entry holding the probability mass outside the kept + entries, ``log(1 - sum(kept probs))``, is appended to both student and teacher so the loss + also penalizes mass the student places outside the teacher's nucleus. """ def __init__( @@ -384,6 +428,10 @@ def __init__( temperature: float = 1.0, reverse: bool = False, top_k: int = 1024, + *, + top_p: float | None = None, + top_p_min_k: int = 1, + add_ghost_token: bool = True, ): """Constructor. @@ -392,9 +440,19 @@ def __init__( temperature: Divide tensors by this value prior to calculating loss. reverse: Whether to reverse the loss as KLD(teacher, student) instead of KLD(student, teacher) top_k: The number of top vocabulary entries to keep from the teacher's distribution. + top_p: Optional nucleus threshold in (0, 1] applied on top of the Top-K selection. + top_p_min_k: Minimum number of entries kept per token when ``top_p`` is active. + add_ghost_token: Whether to append a residual "ghost" token holding the out-of-Top-K + probability mass to both distributions. """ super().__init__(model_config, temperature, reverse) + assert top_k >= 1, f"{top_k=}" + assert top_p is None or 0 < top_p <= 1, f"{top_p=}" + assert top_p_min_k >= 1, f"{top_p_min_k=}" self.top_k = top_k + self.top_p = top_p + self.top_p_min_k = top_p_min_k + self.add_ghost_token = add_ghost_token def forward(self, predictions: Tensor, targets: Tensor) -> Tensor: """Forward function. @@ -445,14 +503,39 @@ def forward(self, predictions: Tensor, targets: Tensor) -> Tensor: final_teacher_logits = top_teacher_vals final_student_logits = top_student_vals - # Standard (dense) Softmax + KL - p = F.log_softmax(final_student_logits, dim=-1) - q = F.log_softmax(final_teacher_logits, dim=-1) - - # KL divergence + # Log-probs of the Top-K entries under the full-vocab distributions, using global + # (full-vocab) log-normalizers so the entries carry true probabilities. + # NOTE: ``torch.topk`` returns entries sorted descending by teacher value. + teacher_logp = final_teacher_logits - self._tp_logsumexp(targets) + student_logp = final_student_logits - self._tp_logsumexp(predictions) + + # Top-P (nucleus) mask over the sorted Top-K: keep entry i iff cumulative mass *before* it + # is < p. This always keeps the entry that crosses the threshold (and thus top-1). + if self.top_p is not None: + teacher_probs = teacher_logp.exp() + mask = (teacher_probs.cumsum(dim=-1) - teacher_probs) < self.top_p + min_keep = min(self.top_p_min_k, teacher_logp.size(-1)) + mask |= torch.arange(teacher_logp.size(-1), device=mask.device) < min_keep + else: + mask = torch.ones_like(teacher_logp, dtype=torch.bool) + + # Ghost token: residual probability mass outside the kept entries, for both distributions. + if self.add_ghost_token: + eps = 1e-8 + student_kept_mass = (student_logp.exp() * mask).sum(dim=-1, keepdim=True) + teacher_kept_mass = (teacher_logp.exp() * mask).sum(dim=-1, keepdim=True) + student_residual = torch.log((1.0 - student_kept_mass).clamp(min=eps)) + teacher_residual = torch.log((1.0 - teacher_kept_mass).clamp(min=eps)) + student_logp = torch.cat([student_logp, student_residual], dim=-1) + teacher_logp = torch.cat([teacher_logp, teacher_residual], dim=-1) + mask = torch.cat([mask, mask.new_ones((*mask.shape[:-1], 1))], dim=-1) + + # Sparse KL divergence: sum_i q_i * (log q_i - log p_i) over kept entries. + p, q = student_logp, teacher_logp if self._reverse: p, q = q, p - loss = torch.sum(F.kl_div(p, q, reduction="none", log_target=True), dim=-1) + kl = q.exp() * (q - p) + loss = torch.sum(mask * kl, dim=-1) # No need to reduce since all ranks compute same global Top-K return self.post_forward(loss, tp_reduce=False) @@ -461,20 +544,23 @@ def forward(self, predictions: Tensor, targets: Tensor) -> Tensor: class LogitsAndIntermediatesLossBalancer(mtd.DistillationLossBalancer): """LossBalancer implementation for Logit and Intermediate losses. - Dynamically weighs distillation and original losses to balance during training. + Intermediate losses are dynamically rescaled to the magnitude of the logits loss, then the + total distillation loss is combined with the original LM loss as a fixed convex combination + ``(1 - alpha) * lm_loss + alpha * kd_loss`` (matching Megatron-LM's offline cached-logits KD). """ - def __init__(self, kd_loss_scale: float = 1.0, skip_original_loss: bool = False): + def __init__(self, kd_loss_alpha: float = 0.9, skip_original_loss: bool = False): """Constructor. Args: - kd_loss_scale: Multiply distillation losses by this before weighing. - (Not used when `skip_original_loss` is True.) + kd_loss_alpha: Weight of the distillation loss in ``(1 - alpha) * lm + alpha * kd``. + Must be in [0, 1]. (Not used when `skip_original_loss` is True.) skip_original_loss: Used to signal whether the original loss should be used, regardless of whether it was passed into ``mtd.DistillationModel.compute_kd_loss()`` or not. """ super().__init__() - self._kd_loss_scale = kd_loss_scale + assert 0 <= kd_loss_alpha <= 1, f"{kd_loss_alpha=}" + self._kd_loss_alpha = kd_loss_alpha self._skip_original_loss = skip_original_loss def forward(self, loss_dict: dict[str, Tensor]) -> Tensor: @@ -500,13 +586,11 @@ def forward(self, loss_dict: dict[str, Tensor]) -> Tensor: intermediate_loss = logits_loss.new_tensor(intermediate_loss) intermediate_loss_scaled = intermediate_loss + kd_loss = logits_loss + intermediate_loss_scaled if self._skip_original_loss: - total_loss = logits_loss + intermediate_loss_scaled + total_loss = kd_loss else: - kd_loss = logits_loss + intermediate_loss_scaled - if kd_loss > 0 and original_loss > 0: # zero when one CP rank has only context tokens - kd_loss *= original_loss.detach() / kd_loss.detach() - total_loss = original_loss + kd_loss * self._kd_loss_scale + total_loss = (1 - self._kd_loss_alpha) * original_loss + self._kd_loss_alpha * kd_loss out_dict = { "kd_loss": total_loss, 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 7419bb49f86..7f381ae4fec 100644 --- a/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py +++ b/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py @@ -14,9 +14,12 @@ # limitations under the License. from functools import partial +from types import SimpleNamespace +import pytest import torch import torch.nn as nn +import torch.nn.functional as F from _test_utils.torch.megatron.models import get_mcore_gpt_model from _test_utils.torch.megatron.utils import run_mcore_inference_with_dummy_input from _test_utils.torch.misc import set_seed @@ -24,6 +27,9 @@ import modelopt.torch.distill as mtd from modelopt.torch.distill.plugins.megatron import ( DistillationConfig, + LogitsAndIntermediatesLossBalancer, + LogitsKLLoss, + TopKLogitsKLLoss, _mtp_excluded_from_quantization, adjust_distillation_model_for_mcore, setup_distillation_config, @@ -124,7 +130,7 @@ def _test_logits_kl_loss(rank, size): loss["kd_loss"].backward() -def _test_topk_logits_kl_loss(top_k, rank, size): +def _test_topk_logits_kl_loss(kd_kwargs, rank, size): """Test TopKLogitsKLLoss with simple forward/backward pass.""" set_seed(SEED) @@ -169,7 +175,7 @@ def _test_topk_logits_kl_loss(top_k, rank, size): # Setup distillation config with TopKLogitsKLLoss via logit_kl_topk argument distill_cfg = setup_distillation_config( - config_or_path=DistillationConfig(logit_kl_topk=top_k), + config_or_path=DistillationConfig(**kd_kwargs), student_cfg=student_model.config, teacher_cfg=teacher_model.config, ) @@ -211,6 +217,12 @@ def _test_topk_logits_kl_loss(top_k, rank, size): assert isinstance(loss, dict), "Loss should be a dictionary" assert "kd_loss" in loss, "Should contain kd_loss key" + # All TP ranks operate on the same global Top-K, so the loss must be identical across ranks. + gathered = [torch.empty_like(loss["kd_loss"]) for _ in range(size)] + torch.distributed.all_gather(gathered, loss["kd_loss"].detach()) + for other in gathered[1:]: + assert torch.allclose(gathered[0], other), "Top-K KD loss differs across TP ranks" + # Backward pass loss["kd_loss"].backward() @@ -260,7 +272,7 @@ def _test_skip_lm_loss_with_mtp(rank, size): ).cuda() distill_cfg = setup_distillation_config( - config_or_path=DistillationConfig(skip_lm_loss=True), + config_or_path=DistillationConfig(kd_loss_alpha=1.0), # skips LM loss student_cfg=student_model.config, teacher_cfg=teacher_model.config, ) @@ -313,9 +325,136 @@ def test_logits_kl_loss(dist_workers): dist_workers.run(_test_logits_kl_loss) -def test_topk_logits_kl_loss(dist_workers, top_k: int = 5): +@pytest.mark.parametrize( + ("top_p", "top_p_min_k", "ghost_token"), + [(None, 1, True), (None, 1, False), (0.9, 1, True), (0.9, 3, False)], +) +def test_topk_logits_kl_loss(dist_workers, top_p, top_p_min_k, ghost_token, top_k: int = 5): """Test TopKLogitsKLLoss with TP parallelism.""" - dist_workers.run(partial(_test_topk_logits_kl_loss, top_k)) + kd_kwargs = { + "logit_kl_topk": top_k, + "logit_kl_top_p": top_p, + "logit_kl_top_p_min_k": top_p_min_k, + "logit_kl_ghost_token": ghost_token, + } + dist_workers.run(partial(_test_topk_logits_kl_loss, kd_kwargs)) + + +def _make_loss_inputs(seq=4, batch=3, vocab=16): + torch.manual_seed(SEED) + student = torch.randn(seq, batch, vocab, requires_grad=True) + teacher = torch.randn(seq, batch, vocab) * 3 # peaky teacher so top-P actually truncates + return student, teacher + + +def test_topk_logits_kl_loss_numerics_full_vocab_matches_dense(): + """With K = vocab and ghost token, Top-K KL equals the dense full-vocab KL (residual ~0).""" + cfg = SimpleNamespace(tensor_model_parallel_size=1) + student, teacher = _make_loss_inputs() + dense = LogitsKLLoss(cfg)(student, teacher)[0] + topk = TopKLogitsKLLoss(cfg, top_k=student.size(-1), add_ghost_token=True)(student, teacher)[0] + assert torch.allclose(dense, topk, atol=1e-5) + # Without ghost token, the unnormalized Top-K KL over the full vocab is also the dense KL. + topk_no_ghost = TopKLogitsKLLoss(cfg, top_k=student.size(-1), add_ghost_token=False)( + student, teacher + )[0] + assert torch.allclose(dense, topk_no_ghost, atol=1e-5) + + +def test_topk_logits_kl_loss_numerics_ghost_token_reference(): + """Top-K + ghost token matches a hand-written reference on the K+1 bucketed distributions.""" + cfg = SimpleNamespace(tensor_model_parallel_size=1) + student, teacher = _make_loss_inputs() + k = 4 + loss = TopKLogitsKLLoss(cfg, top_k=k, add_ghost_token=True)(student, teacher)[0] + + q_full = F.log_softmax(teacher, dim=-1) + p_full = F.log_softmax(student, dim=-1) + _, idx = torch.topk(teacher, k, dim=-1) + q_k, p_k = q_full.gather(-1, idx), p_full.gather(-1, idx) + q_rest = torch.log1p(-q_k.exp().sum(-1, keepdim=True)) + p_rest = torch.log1p(-p_k.exp().sum(-1, keepdim=True)) + q = torch.cat([q_k, q_rest], -1) + p = torch.cat([p_k, p_rest], -1) + ref = (q.exp() * (q - p)).sum(-1).transpose(0, 1) + assert torch.allclose(loss, ref, atol=1e-5) + # Sanity: total mass within the K+1 buckets is 1 for both distributions. + assert torch.allclose(q.exp().sum(-1), torch.ones_like(q[..., 0]), atol=1e-5) + assert torch.allclose(p.exp().sum(-1), torch.ones_like(p[..., 0]), atol=1e-5) + + +@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.""" + cfg = SimpleNamespace(tensor_model_parallel_size=1) + student, teacher = _make_loss_inputs() + q = F.log_softmax(teacher / temperature, dim=-1) + p = F.log_softmax(student / temperature, dim=-1) + + dense = LogitsKLLoss(cfg, temperature=temperature)(student, teacher)[0] + ref_dense = (q.exp() * (q - p)).sum(-1).transpose(0, 1) + assert torch.allclose(dense, ref_dense, atol=1e-5) + + k = 4 + topk = TopKLogitsKLLoss(cfg, temperature=temperature, top_k=k, add_ghost_token=True)( + student, teacher + )[0] + _, idx = torch.topk(teacher, k, dim=-1) + q_k, p_k = q.gather(-1, idx), p.gather(-1, idx) + q_rest = torch.log1p(-q_k.exp().sum(-1, keepdim=True)) + p_rest = torch.log1p(-p_k.exp().sum(-1, keepdim=True)) + qq = torch.cat([q_k, q_rest], -1) + pp = torch.cat([p_k, p_rest], -1) + ref_topk = (qq.exp() * (qq - pp)).sum(-1).transpose(0, 1) + assert torch.allclose(topk, ref_topk, atol=1e-5) + + +def test_topk_logits_kl_loss_top_p_masks_tail(): + """Top-P zeroes out-of-nucleus entries and honors the min_k floor.""" + cfg = SimpleNamespace(tensor_model_parallel_size=1) + student, teacher = _make_loss_inputs() + k = 8 + q_full = F.log_softmax(teacher, dim=-1) + q_k, idx = torch.topk(q_full, k, dim=-1) + p_k = F.log_softmax(student, dim=-1).gather(-1, idx) + probs = q_k.exp() + keep = (probs.cumsum(-1) - probs) < 0.5 + + # No ghost token: loss is exactly the masked partial KL sum. + loss = TopKLogitsKLLoss(cfg, top_k=k, top_p=0.5, add_ghost_token=False)(student, teacher)[0] + ref = (keep * probs * (q_k - p_k)).sum(-1).transpose(0, 1) + assert torch.allclose(loss, ref, atol=1e-5) + assert not keep.all(), "test inputs should produce some truncation" + + # min_k floor forces at least min_k entries even when nucleus is tiny. + min_k = 3 + loss_min = TopKLogitsKLLoss(cfg, top_k=k, top_p=1e-6, top_p_min_k=min_k, add_ghost_token=False)( + student, teacher + )[0] + ref_min = ((torch.arange(k) < min_k) * probs * (q_k - p_k)).sum(-1).transpose(0, 1) + assert torch.allclose(loss_min, ref_min, atol=1e-5) + + # Ghost token with top-P: residual is mass outside the kept nucleus, distributions sum to 1. + loss_ghost = TopKLogitsKLLoss(cfg, top_k=k, top_p=0.5, add_ghost_token=True)(student, teacher)[ + 0 + ] + q_rest = torch.log1p(-(probs * keep).sum(-1, keepdim=True)) + p_rest = torch.log1p(-(p_k.exp() * keep).sum(-1, keepdim=True)) + ref_ghost = ref + (q_rest.exp() * (q_rest - p_rest)).sum(-1).transpose(0, 1) + assert torch.allclose(loss_ghost, ref_ghost, atol=1e-5) + assert loss_ghost.shape == (student.size(1), student.size(0)) + loss_ghost.sum().backward() + assert student.grad is not None and torch.isfinite(student.grad).all() + + +def test_distillation_config_top_p_validation(): + with pytest.raises(AssertionError): + DistillationConfig(logit_kl_top_p=0.9) # requires logit_kl_topk + with pytest.raises(AssertionError): + DistillationConfig(logit_kl_topk=8, logit_kl_top_p=1.5) + with pytest.raises(AssertionError): + DistillationConfig(logit_kl_topk=8, logit_kl_top_p=0.9, logit_kl_top_p_min_k=0) + DistillationConfig(logit_kl_topk=8, logit_kl_top_p=1.0, logit_kl_top_p_min_k=2) def test_skip_lm_loss_with_mtp(dist_workers): @@ -356,3 +495,44 @@ def _model(*, with_mtp: bool, body_quant: bool, mtp_quant: bool) -> nn.Module: assert not _mtp_excluded_from_quantization( _model(with_mtp=False, body_quant=True, mtp_quant=False) ) + + +def test_loss_balancer_convex_combination(): + """Total loss is (1 - alpha) * lm + alpha * (logits + rescaled intermediate).""" + lm = torch.tensor(2.0) + logits = torch.tensor(0.5) + inter = torch.tensor(4.0) # rescaled to logits magnitude -> contributes 0.5 + key = mtd.loss_balancers.STUDENT_LOSS_KEY + + out = LogitsAndIntermediatesLossBalancer(kd_loss_alpha=0.25)( + {key: lm, "LogitsKLLoss_0": logits, "HiddenStateCosineLoss_0": inter} + ) + assert torch.allclose(out["kd_loss"], torch.tensor(0.75 * 2.0 + 0.25 * (0.5 + 0.5))) + assert torch.allclose(out["logits_loss"], logits) + + # alpha=1 ignores the LM loss entirely; skip_original_loss does the same regardless of alpha. + out = LogitsAndIntermediatesLossBalancer(kd_loss_alpha=1.0)({key: lm, "LogitsKLLoss_0": logits}) + assert torch.allclose(out["kd_loss"], logits) + out = LogitsAndIntermediatesLossBalancer(kd_loss_alpha=0.0, skip_original_loss=True)( + {key: lm, "LogitsKLLoss_0": logits} + ) + assert torch.allclose(out["kd_loss"], logits) + + with pytest.raises(AssertionError): + LogitsAndIntermediatesLossBalancer(kd_loss_alpha=1.5) + with pytest.raises(AssertionError): + DistillationConfig(kd_loss_alpha=-0.1) + + +def test_distillation_config_deprecations(): + """skip_lm_loss is derived from kd_loss_alpha; legacy fields warn.""" + assert DistillationConfig(kd_loss_alpha=1.0).skip_lm_loss is True + assert DistillationConfig(kd_loss_alpha=0.9).skip_lm_loss is False + + with pytest.warns(DeprecationWarning, match="skip_lm_loss is deprecated"): + cfg = DistillationConfig(kd_loss_alpha=0.9, skip_lm_loss=True) + assert cfg.skip_lm_loss is False # user value overridden + + with pytest.warns(DeprecationWarning, match="kd_loss_scale is deprecated"): + cfg = DistillationConfig(kd_loss_scale=2.0) + assert cfg.kd_loss_alpha == 0.9 From 95c9aa7dad1072927544f3962a249144b4539729 Mon Sep 17 00:00:00 2001 From: Asha Anoosheh Date: Thu, 17 Sep 2026 19:10:18 +0200 Subject: [PATCH 2/3] Don't expose ghost token in MBridge script Signed-off-by: Asha Anoosheh --- examples/megatron_bridge/distill.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 89b44e6bc5f..268dc855b3a 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -215,11 +215,6 @@ def get_args(): default=1, help="Minimum number of top-k entries kept per token when --logit_kl_top_p is active.", ) - parser.add_argument( - "--no_logit_kl_ghost_token", - action="store_true", - help="Disable the residual 'ghost' token (out-of-top-k probability mass) in the top-k KL loss.", - ) parser.add_argument("--lr", type=float, default=1e-4, help="Peak learning rate") parser.add_argument("--min_lr", type=float, default=1e-5, help="Minimum learning rate") parser.add_argument("--lr_warmup_iters", type=int, default=50, help="Number of LR warmup steps") @@ -471,7 +466,6 @@ def _build_model_provider(hf_path, load_weights=True, moe_grouped_gemm=True): logit_kl_topk=args.logit_kl_topk, logit_kl_top_p=args.logit_kl_top_p, logit_kl_top_p_min_k=args.logit_kl_top_p_min_k, - logit_kl_ghost_token=not args.no_logit_kl_ghost_token, ) # HF VLM configs expose ``vision_config``; Megatron-Bridge nests the text model under From 3f751f3385e1753cbd2c435497c86d5892f66d3a Mon Sep 17 00:00:00 2001 From: Asha Anoosheh Date: Thu, 17 Sep 2026 19:52:54 +0200 Subject: [PATCH 3/3] Address review comments Signed-off-by: Asha Anoosheh --- CHANGELOG.rst | 2 +- examples/megatron_bridge/distill.py | 12 ++ modelopt/torch/distill/plugins/megatron.py | 128 ++++++++++-------- .../distill/plugins/test_distill_megatron.py | 34 +++-- 4 files changed, 112 insertions(+), 64 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 8b78b03a004..3008c8e6f0f 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -35,7 +35,7 @@ Changelog **Deprecations** -- ``DistillationConfig.kd_loss_scale`` and ``DistillationConfig.skip_lm_loss`` (Megatron distillation plugin) are deprecated. ``kd_loss_scale`` is ignored with a ``DeprecationWarning``; a user-provided ``skip_lm_loss`` is overridden by the value derived from ``kd_loss_alpha`` with a ``DeprecationWarning``. The ``--no_skip_lm_loss`` and ``--kd_loss_scale`` flags in ``examples/megatron_bridge/distill.py`` are likewise deprecated and ignored. +- ``DistillationConfig.kd_loss_scale`` and ``DistillationConfig.skip_lm_loss`` (Megatron distillation plugin) are deprecated. ``kd_loss_scale`` is ignored with a ``FutureWarning``; an explicit ``skip_lm_loss=True`` is translated to ``kd_loss_alpha=1.0`` with a ``FutureWarning``, and otherwise ``skip_lm_loss`` is derived from ``kd_loss_alpha``. The ``--no_skip_lm_loss`` and ``--kd_loss_scale`` flags in ``examples/megatron_bridge/distill.py`` are likewise deprecated and ignored. - Rename the architecture-specific recipe tier from ``modelopt_recipes/huggingface/`` to ``modelopt_recipes/model_type/`` to clarify that it holds recipes shared across every checkpoint of a Hugging Face ``model_type``. Saved ``--recipe huggingface//...`` paths still resolve via a backward-compatibility alias but now emit a ``FutureWarning``, so update them to ``model_type//...`` as the ``huggingface/`` prefix is deprecated. - The single-format quantization CLI flags are deprecated in favour of ``--recipe`` and will be removed in a future release; passing one now emits a ``FutureWarning``. ``examples/hf_ptq``: ``--qformat`` and ``--kv_cache_qformat``. ``examples/megatron_bridge/quantize.py``: ``--quant_cfg``, ``--kv_cache_quant`` and ``--weight_only``. ``examples/torch_onnx/torch_quant_to_onnx.py``: ``--qformat``. A recipe carries the quantization config, the calibration algorithm and the KV-cache setting in one file, so they cannot drift apart the way separate flags can -- and ``--recipe`` already took precedence over all six, silently on ``hf_ptq`` and with a warning on ``megatron_bridge`` -- with one gap the recipe closes rather than inherits: a weight AutoQuantize recipe that omits ``kv_cache`` still falls back to ``--kv_cache_qformat``, so set ``kv_cache`` in the recipe when migrating. Use a recipe from ``modelopt_recipes/general/ptq/``, an architecture-specific one under ``modelopt_recipes/model_type//``, or a checkpoint-specific one under ``modelopt_recipes/models/``. The warning fires only when a flag is passed explicitly: ``--qformat`` defaults to ``fp8`` and ``--kv_cache_qformat`` to ``fp8_cast``, so warning on the defaults would fire on every run, including runs that correctly use ``--recipe``. ``examples/speculative_decoding/scripts/quantize_drafter.py`` keeps ``--qformat`` undeprecated: it has no ``--recipe`` alternative yet. - The TensorRT-LLM checkpoint export format is deprecated and will be removed in 0.49.0: ``export_tensorrt_llm_checkpoint`` and ``torch_to_tensorrt_llm_checkpoint`` now emit a ``DeprecationWarning`` on use. Use ``export_hf_checkpoint``, which exports a unified Hugging Face checkpoint deployable on TensorRT-LLM, vLLM and SGLang. Its implementation moved to ``modelopt.torch.export.trtllm``, so import those two functions from there and the ``ModelConfig`` dataclasses from ``modelopt.torch.export.trtllm.model_config``; both functions remain importable from ``modelopt.torch.export`` for this release only. diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 268dc855b3a..da4cabb70ed 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -23,6 +23,7 @@ import argparse import contextlib import os +import warnings import torch from export_distilled_megatron_to_hf import export_llm_to_hf, save_vlm_to_hf @@ -461,6 +462,17 @@ def _build_model_provider(hf_path, load_weights=True, moe_grouped_gemm=True): f"sizes differ ({padded['student']} vs {padded['teacher']})." ) + if args.kd_loss_scale is not None: + warnings.warn( + "--kd_loss_scale is deprecated and ignored; use --kd_loss_alpha instead.", + FutureWarning, + ) + if args.no_skip_lm_loss: + warnings.warn( + "--no_skip_lm_loss is deprecated and ignored; whether the LM loss is skipped is derived " + "from --kd_loss_alpha (skipped iff 1.0).", + FutureWarning, + ) kd_config = ModelOptDistillConfig( kd_loss_alpha=args.kd_loss_alpha, logit_kl_topk=args.logit_kl_topk, diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py index c6c5af8ab07..1296f399132 100644 --- a/modelopt/torch/distill/plugins/megatron.py +++ b/modelopt/torch/distill/plugins/megatron.py @@ -60,8 +60,8 @@ class DistillationConfig: kd_loss_alpha: Weight of the distillation loss in the convex combination ``(1 - alpha) * lm_loss + alpha * kd_loss``. Must be in [0, 1]. When ``1.0``, the standard language model loss is skipped entirely (``skip_lm_loss`` is derived from this value). - skip_lm_loss: DEPRECATED. Derived from ``kd_loss_alpha`` (``True`` iff ``kd_loss_alpha == 1.0``); - any user-provided value is overridden with a warning. + skip_lm_loss: DEPRECATED. Derived from ``kd_loss_alpha`` (``True`` iff ``kd_loss_alpha == 1.0``). + An explicit ``True`` is translated to ``kd_loss_alpha = 1.0`` with a warning. kd_loss_scale: DEPRECATED and ignored. Use ``kd_loss_alpha`` instead. logit_kl_temperature: Temperature for the logit KL-divergence loss. logit_kl_topk: If not None, use TopKLogitsKLLoss instead of LogitsKLLoss with this top-k value. @@ -97,19 +97,33 @@ def __post_init__(self): "DistillationConfig.kd_loss_scale is deprecated and ignored. The distillation loss " "is no longer rescaled to the LM loss magnitude; the total loss is now " "(1 - kd_loss_alpha) * lm_loss + kd_loss_alpha * kd_loss. Set `kd_loss_alpha` instead.", - DeprecationWarning, + FutureWarning, stacklevel=2, ) - derived_skip_lm_loss = self.kd_loss_alpha == 1.0 if self.skip_lm_loss is not None: - warnings.warn( - "DistillationConfig.skip_lm_loss is deprecated and is now derived from `kd_loss_alpha` " - f"(skip iff kd_loss_alpha == 1.0). Overriding skip_lm_loss={self.skip_lm_loss} with " - f"{derived_skip_lm_loss} (kd_loss_alpha={self.kd_loss_alpha}).", - DeprecationWarning, - stacklevel=2, - ) - self.skip_lm_loss = derived_skip_lm_loss + if self.skip_lm_loss and self.kd_loss_alpha != 1.0: + warnings.warn( + "DistillationConfig.skip_lm_loss is deprecated; translating skip_lm_loss=True to " + "kd_loss_alpha=1.0. Set `kd_loss_alpha` directly instead.", + FutureWarning, + stacklevel=2, + ) + self.kd_loss_alpha = 1.0 + elif not self.skip_lm_loss and self.kd_loss_alpha == 1.0: + warnings.warn( + "DistillationConfig.skip_lm_loss is deprecated, and skip_lm_loss=False conflicts " + "with kd_loss_alpha=1.0, which skips the LM loss. Set `kd_loss_alpha` < 1.0 instead.", + FutureWarning, + stacklevel=2, + ) + else: + warnings.warn( + "DistillationConfig.skip_lm_loss is deprecated and is now derived from " + "`kd_loss_alpha` (skipped iff kd_loss_alpha == 1.0). Stop passing it.", + FutureWarning, + stacklevel=2, + ) + self.skip_lm_loss = self.kd_loss_alpha == 1.0 assert self.logit_kl_temperature > 0, f"{self.logit_kl_temperature=}" if self.logit_kl_top_p is not None: assert self.logit_kl_topk is not None, "logit_kl_top_p requires logit_kl_topk" @@ -362,9 +376,13 @@ def forward(self, predictions: Tensor, targets: Tensor) -> Tensor: """ predictions, targets = self.pre_forward(predictions, targets) - # Temperature-scaled log probabilities (log softmax), globally normalized across TP vocab shards. - p = predictions.float() / self._temperature - self._tp_logsumexp(predictions) - q = targets.float() / self._temperature - self._tp_logsumexp(targets) + # Division by temp should happen prior to finding max for both student and teacher. + output_teacher = targets.float() / self._temperature + output_student = predictions.float() / self._temperature + + # Log probabilities (log softmax), globally normalized across TP vocab shards. + p = output_student - self._tp_logsumexp(output_student) + q = output_teacher - self._tp_logsumexp(output_teacher) # KL divergence if self._reverse: @@ -373,36 +391,28 @@ def forward(self, predictions: Tensor, targets: Tensor) -> Tensor: return self.post_forward(loss, tp_reduce=True) - def _tp_logsumexp(self, logits: Tensor, num_chunks: int = 4) -> Tensor: - """Log-sum-exp of ``logits / self._temperature`` over the vocab dim across all TP shards. + def _tp_logsumexp(self, logits: Tensor) -> Tensor: + """Log-sum-exp over the vocab dim across all TP shards (shape ``[..., 1]``). - Accumulates in fp32 over ``num_chunks`` vocab chunks so the transient fp32 working set is a - fraction of the (possibly bf16) full-vocab input. NOTE: for inputs requiring grad, autograd - still retains each chunk's ``exp`` output for backward, so the total saved activation - equals a full-vocab fp32 tensor. Returns shape ``[..., 1]``. + ``logits`` are expected to be fp32 and already temperature-scaled. """ - # Max is exact under the monotonic temperature scaling, so take it in the native dtype - # and defer the temperature division to the centered values: lse(x/T) = max/T + log(sum(exp((x-max)/T))). - logits_max = logits.amax(dim=-1, keepdim=True).float() - if self._config.tensor_model_parallel_size > 1: - tp_group = parallel_state.get_tensor_model_parallel_group() - torch.distributed.all_reduce( - logits_max, op=torch.distributed.ReduceOp.MAX, group=tp_group - ) + if self._config.tensor_model_parallel_size == 1: + return torch.logsumexp(logits, dim=-1, keepdim=True) + + tp_group = parallel_state.get_tensor_model_parallel_group() + + # Subtract maximum value along vocab dimension across all GPUs (for stability) + logits_max = logits.amax(dim=-1, keepdim=True) + torch.distributed.all_reduce(logits_max, op=torch.distributed.ReduceOp.MAX, group=tp_group) logits_max = logits_max.detach() - denom = None - chunk_size = -(-logits.size(-1) // num_chunks) # ceil division - for chunk in logits.split(chunk_size, dim=-1): - centered = (chunk.float() - logits_max) / self._temperature - partial = torch.exp(centered).sum(dim=-1, keepdim=True) - denom = partial if denom is None else denom + partial - if self._config.tensor_model_parallel_size > 1: - # We can't use standard all_reduce function here since the computation - # that follows it isn't identical across TP ranks. - denom = dist_nn.functional.all_reduce(denom, group=tp_group) + # Compute global softmax denominator. + # We can't use standard all_reduce function here since the computation + # that follows it isn't identical across TP ranks. + denom = torch.exp(logits - logits_max).sum(dim=-1, keepdim=True) + denom = dist_nn.functional.all_reduce(denom, group=tp_group) - return logits_max / self._temperature + torch.log(denom) + return logits_max + torch.log(denom) class TopKLogitsKLLoss(LogitsKLLoss): @@ -471,14 +481,15 @@ def forward(self, predictions: Tensor, targets: Tensor) -> Tensor: f"top_k ({self.top_k}) is larger than total vocab size ({targets.size(-1) * tp_size})" ) - # Take K from each rank, then the global Top-K of those. Reduce before the fp32 cast: - # casting the full vocab first defeats the point. Selection is unchanged (widening is - # exact, temperature scaling monotonic). + # Divide by temperature first + output_teacher = targets.float() / self._temperature + output_student = predictions.float() / self._temperature + + # Extract local Top-K + # 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)) - top_teacher_vals, top_idx = torch.topk(targets, local_top_k, dim=-1) - top_student_vals = torch.gather(predictions, dim=-1, index=top_idx) - top_teacher_vals = top_teacher_vals.float() / self._temperature - top_student_vals = top_student_vals.float() / self._temperature + top_teacher_vals, top_idx = torch.topk(output_teacher, local_top_k, dim=-1) + top_student_vals = torch.gather(output_student, dim=-1, index=top_idx) if tp_size > 1: tp_group = parallel_state.get_tensor_model_parallel_group() @@ -506,8 +517,8 @@ def forward(self, predictions: Tensor, targets: Tensor) -> Tensor: # Log-probs of the Top-K entries under the full-vocab distributions, using global # (full-vocab) log-normalizers so the entries carry true probabilities. # NOTE: ``torch.topk`` returns entries sorted descending by teacher value. - teacher_logp = final_teacher_logits - self._tp_logsumexp(targets) - student_logp = final_student_logits - self._tp_logsumexp(predictions) + teacher_logp = final_teacher_logits - self._tp_logsumexp(output_teacher) + student_logp = final_student_logits - self._tp_logsumexp(output_student) # Top-P (nucleus) mask over the sorted Top-K: keep entry i iff cumulative mass *before* it # is < p. This always keeps the entry that crosses the threshold (and thus top-1). @@ -520,12 +531,18 @@ def forward(self, predictions: Tensor, targets: Tensor) -> Tensor: mask = torch.ones_like(teacher_logp, dtype=torch.bool) # Ghost token: residual probability mass outside the kept entries, for both distributions. + # Computed in log space as log(1 - exp(log_kept)) = log(-expm1(log_kept)), which stays + # accurate and differentiable when the kept mass is close to 1. if self.add_ghost_token: - eps = 1e-8 - student_kept_mass = (student_logp.exp() * mask).sum(dim=-1, keepdim=True) - teacher_kept_mass = (teacher_logp.exp() * mask).sum(dim=-1, keepdim=True) - student_residual = torch.log((1.0 - student_kept_mass).clamp(min=eps)) - teacher_residual = torch.log((1.0 - teacher_kept_mass).clamp(min=eps)) + neg_tiny = -1e-7 # keep log(kept_mass) strictly below 0 so expm1 stays negative + student_log_kept = torch.logsumexp( + student_logp.masked_fill(~mask, float("-inf")), dim=-1, keepdim=True + ).clamp(max=neg_tiny) + teacher_log_kept = torch.logsumexp( + teacher_logp.masked_fill(~mask, float("-inf")), dim=-1, keepdim=True + ).clamp(max=neg_tiny) + student_residual = torch.log(-torch.expm1(student_log_kept)) + teacher_residual = torch.log(-torch.expm1(teacher_log_kept)) student_logp = torch.cat([student_logp, student_residual], dim=-1) teacher_logp = torch.cat([teacher_logp, teacher_residual], dim=-1) mask = torch.cat([mask, mask.new_ones((*mask.shape[:-1], 1))], dim=-1) @@ -580,7 +597,8 @@ def forward(self, loss_dict: dict[str, Tensor]) -> Tensor: intermediate_loss = sum(loss_dict.values()) / max(len(loss_dict), 1) if intermediate_loss > 0: - dynamic_scale = logits_loss.detach() / intermediate_loss.detach() + # abs(): the Top-K partial KL without a ghost token is not a true KL and can be negative. + dynamic_scale = logits_loss.detach().abs() / intermediate_loss.detach() intermediate_loss_scaled = intermediate_loss * dynamic_scale else: intermediate_loss = logits_loss.new_tensor(intermediate_loss) 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 7f381ae4fec..2101e937a49 100644 --- a/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py +++ b/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py @@ -353,7 +353,7 @@ def test_topk_logits_kl_loss_numerics_full_vocab_matches_dense(): student, teacher = _make_loss_inputs() dense = LogitsKLLoss(cfg)(student, teacher)[0] topk = TopKLogitsKLLoss(cfg, top_k=student.size(-1), add_ghost_token=True)(student, teacher)[0] - assert torch.allclose(dense, topk, atol=1e-5) + assert torch.allclose(dense, topk, atol=1e-6) # Without ghost token, the unnormalized Top-K KL over the full vocab is also the dense KL. topk_no_ghost = TopKLogitsKLLoss(cfg, top_k=student.size(-1), add_ghost_token=False)( student, teacher @@ -377,7 +377,7 @@ def test_topk_logits_kl_loss_numerics_ghost_token_reference(): q = torch.cat([q_k, q_rest], -1) p = torch.cat([p_k, p_rest], -1) ref = (q.exp() * (q - p)).sum(-1).transpose(0, 1) - assert torch.allclose(loss, ref, atol=1e-5) + assert torch.allclose(loss, ref, atol=1e-6) # Sanity: total mass within the K+1 buckets is 1 for both distributions. assert torch.allclose(q.exp().sum(-1), torch.ones_like(q[..., 0]), atol=1e-5) assert torch.allclose(p.exp().sum(-1), torch.ones_like(p[..., 0]), atol=1e-5) @@ -406,7 +406,7 @@ def test_logits_kl_losses_temperature_scaling(temperature): qq = torch.cat([q_k, q_rest], -1) pp = torch.cat([p_k, p_rest], -1) ref_topk = (qq.exp() * (qq - pp)).sum(-1).transpose(0, 1) - assert torch.allclose(topk, ref_topk, atol=1e-5) + assert torch.allclose(topk, ref_topk, atol=1e-6) def test_topk_logits_kl_loss_top_p_masks_tail(): @@ -441,7 +441,7 @@ def test_topk_logits_kl_loss_top_p_masks_tail(): q_rest = torch.log1p(-(probs * keep).sum(-1, keepdim=True)) p_rest = torch.log1p(-(p_k.exp() * keep).sum(-1, keepdim=True)) ref_ghost = ref + (q_rest.exp() * (q_rest - p_rest)).sum(-1).transpose(0, 1) - assert torch.allclose(loss_ghost, ref_ghost, atol=1e-5) + assert torch.allclose(loss_ghost, ref_ghost, atol=1e-6) assert loss_ghost.shape == (student.size(1), student.size(0)) loss_ghost.sum().backward() assert student.grad is not None and torch.isfinite(student.grad).all() @@ -510,6 +510,13 @@ def test_loss_balancer_convex_combination(): assert torch.allclose(out["kd_loss"], torch.tensor(0.75 * 2.0 + 0.25 * (0.5 + 0.5))) assert torch.allclose(out["logits_loss"], logits) + # A negative logits loss (possible for Top-K KL without ghost token) must not flip the sign of + # the intermediate-loss contribution. + out = LogitsAndIntermediatesLossBalancer(kd_loss_alpha=1.0)( + {key: lm, "LogitsKLLoss_0": -logits, "HiddenStateCosineLoss_0": inter} + ) + assert torch.allclose(out["kd_loss"], -logits + logits) # -0.5 + abs(-0.5) * (4.0 / 4.0) = 0 + # alpha=1 ignores the LM loss entirely; skip_original_loss does the same regardless of alpha. out = LogitsAndIntermediatesLossBalancer(kd_loss_alpha=1.0)({key: lm, "LogitsKLLoss_0": logits}) assert torch.allclose(out["kd_loss"], logits) @@ -525,14 +532,25 @@ def test_loss_balancer_convex_combination(): def test_distillation_config_deprecations(): - """skip_lm_loss is derived from kd_loss_alpha; legacy fields warn.""" + """skip_lm_loss is derived from kd_loss_alpha; legacy fields warn with FutureWarning.""" assert DistillationConfig(kd_loss_alpha=1.0).skip_lm_loss is True assert DistillationConfig(kd_loss_alpha=0.9).skip_lm_loss is False - with pytest.warns(DeprecationWarning, match="skip_lm_loss is deprecated"): + # Explicit skip_lm_loss=True is translated to kd_loss_alpha=1.0 rather than overridden. + with pytest.warns(FutureWarning, match="translating skip_lm_loss=True"): cfg = DistillationConfig(kd_loss_alpha=0.9, skip_lm_loss=True) - assert cfg.skip_lm_loss is False # user value overridden + assert cfg.kd_loss_alpha == 1.0 and cfg.skip_lm_loss is True + + # skip_lm_loss=False with alpha=1.0 is a conflict; alpha wins. + with pytest.warns(FutureWarning, match="conflicts with kd_loss_alpha=1.0"): + cfg = DistillationConfig(kd_loss_alpha=1.0, skip_lm_loss=False) + assert cfg.skip_lm_loss is True + + # Consistent but deprecated usage still warns. + with pytest.warns(FutureWarning, match="skip_lm_loss is deprecated"): + cfg = DistillationConfig(kd_loss_alpha=0.9, skip_lm_loss=False) + assert cfg.skip_lm_loss is False - with pytest.warns(DeprecationWarning, match="kd_loss_scale is deprecated"): + with pytest.warns(FutureWarning, match="kd_loss_scale is deprecated"): cfg = DistillationConfig(kd_loss_scale=2.0) assert cfg.kd_loss_alpha == 0.9