Bring offline KD upgrades such as Ghost Token and Top-P to Megatron K…

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 <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/>`_ 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 ``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/<model_type>/...`` paths still resolve via a backward-compatibility alias but now emit a ``FutureWarning``, so update them to ``model_type/<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/<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.

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
@@ -172,9 +173,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 +203,19 @@ 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("--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")
@@ -434,10 +462,22 @@ 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(
-        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,
     )
 
     # HF VLM configs expose ``vision_config``; Megatron-Bridge nests the text model under