diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index adcaec96729..76f31c28d90 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -41,6 +41,7 @@ Changelog
 
 *Megatron Framework (M-LM / M-Bridge)*
 
+- Add optional Top-P (nucleus) truncation to the Megatron top-k logits KD loss (now ``TopLogitsKLLoss``) 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.
 - Add ``--mlflow <tracking-uri>`` to the ``examples/megatron_bridge`` scripts that write a checkpoint -- ``prune_minitron.py``, ``quantize.py``, ``distill.py``, ``export_quantized_megatron_to_hf.py`` and ``export_distilled_megatron_to_hf.py`` (MLflow's own ``MLFLOW_TRACKING_URI`` is honoured too). Each run records the invocation, its arguments as searchable params and its log -- a distillation records its training metrics instead of the rank-0 log -- and writes ``.experiment.json`` into the checkpoint it produced, so a pruning, a quantization, the distillation that refines its checkpoint and the export that deploys it can be traced to one another; uploading the checkpoints themselves stays off unless ``--mlflow_log_checkpoints`` is passed.
 - Add ``--ep_size`` to ``examples/megatron_bridge/export_quantized_megatron_to_hf.py`` so large MoE models with grouped-GEMM experts can be exported with their experts sharded across GPUs. Checkpoints built with ``--no_moe_grouped_gemm`` must still be exported at ``--ep_size 1``.
@@ -52,6 +53,8 @@ Changelog
 
 **Backward Breaking Changes**
 
+- ``modelopt.torch.distill.plugins.megatron.TopKLogitsKLLoss`` (``logit_kl_topk`` in ``DistillationConfig``) is renamed to ``TopLogitsKLLoss``, keeping the old name as a deprecated alias. It now normalizes both distributions over the full vocabulary instead of re-normalizing over the Top-K entries, and always 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.
+- ``LogitsAndIntermediatesLossBalancer`` (Megatron distillation plugin) no longer rescales the distillation loss to the magnitude of the LM loss when the LM loss is included. The total is now the fixed convex combination ``(1 - alpha) * lm_loss + alpha * kd_loss`` with ``DistillationConfig.kd_loss_alpha`` (in [0, 1]), matching Megatron-LM's offline cached-logits KD. The default ``kd_loss_alpha=1.0`` skips the LM loss, as before. ``DistillationConfig.skip_lm_loss`` and ``kd_loss_scale`` are removed and raise ``ValueError`` if passed; set ``kd_loss_alpha`` instead. In ``examples/megatron_bridge/distill.py``, ``--kd_loss_alpha`` replaces ``--no_skip_lm_loss`` and ``--kd_loss_scale``.
 - The Megatron-Core DeepSeek-V4 indexer (``CSAIndexer``) is now a quantization module and persists its quantizer state in the checkpoint as ``indexer._extra_state``. A DeepSeek-V4 model quantized with an earlier release resumes from its ``torch_dist`` checkpoint only with a non-strict load (``--dist-ckpt-strictness log_unexpected`` in Megatron-LM) until it is saved again.
 - ``examples/hf_ptq`` no longer detects MTP layers by name. Weights the loader could not place -- an MTP head, an auxiliary tower -- are identified from Transformers' own accounting: the model is loaded with ``from_pretrained(..., output_loading_info=True)`` and the reported ``unexpected_keys`` (present in the checkpoint, not in the model's architecture) are recorded on the model and carried into the export unchanged. Everything the loader *did* place goes through the normal export path. This removes ``load_mtp_weights``, ``mtp_layer_prefixes_from_checkpoint`` and their support matrix of MTP storage conventions, along with ``_add_mtp_exclusions`` and the pre-quantization ``enable: False`` entries ``hf_ptq`` appended to the recipe's ``quant_cfg``. Two consequences: MTP layers now follow the recipe like any other module instead of being force-excluded by the script -- matching ``examples/megatron_bridge``, which has no MTP-specific code at all -- and ``quantization_config.ignore`` can no longer claim a layer is unquantized that the export in fact quantized. Recipes importing ``configs/ptq/units/default_disabled_quantizers`` still disable ``mtp.*``, so their behaviour is unchanged; a recipe omitting that unit will now quantize an MTP the model actually built.
 
diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py
index 8729f9dda8c..f71114cf1f3 100644
--- a/examples/megatron_bridge/distill.py
+++ b/examples/megatron_bridge/distill.py
@@ -208,9 +208,11 @@ 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"
+        "--kd_loss_alpha",
+        type=float,
+        default=1.0,
+        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=1.0, help="KD loss weight")
     parser.add_argument(
         "--no_async_save",
         action="store_true",
@@ -218,11 +220,24 @@ def get_args():
         "CUDA context, which fails when the training process already fills the GPU.",
     )
     parser.add_argument(
-        "--logit_kl_topk",
+        "--logit_kl_top_k",
         type=int,
         default=None,
-        help="Restrict the logit KL loss to the teacher's top-k vocabulary entries, "
-        "replacing the full-vocab temporaries with [seq, k] ones.",
+        help="Restrict the logit KL loss to the teacher's top-k vocabulary entries plus a residual "
+        "bucket for the remaining probability mass (distributions are still normalized over the full vocab).",
+    )
+    parser.add_argument(
+        "--logit_kl_top_p",
+        type=float,
+        default=None,
+        help="Nucleus threshold in (0, 1] applied on top of --logit_kl_top_k: 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")
@@ -474,9 +489,10 @@ 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,
-        logit_kl_topk=args.logit_kl_topk,
+        kd_loss_alpha=args.kd_loss_alpha,
+        logit_kl_topk=args.logit_kl_top_k,
+        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
diff --git a/examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/README.md b/examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/README.md
index 02ab3334be4..6638098494e 100644
--- a/examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/README.md
+++ b/examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/README.md
@@ -176,7 +176,7 @@ srun ... python -u /opt/Model-Optimizer/examples/megatron_bridge/distill.py \
     --gbs 512 \
     --train_iters 500 \
     --lr 1e-5 --min_lr 1e-6 --lr_warmup_iters 50 \
-    --logit_kl_topk 4096 \
+    --logit_kl_top_k 4096 \
     --recompute_granularity full --recompute_method uniform --recompute_num_layers 1 \
     --no_async_save \
     --eval_iters 0 \
@@ -190,7 +190,7 @@ Non-default arguments:
 - `--tp_size 1 --pp_size 1` — **required, not chosen** (see below). `--ep_size 8` must match the PTQ checkpoint. You may increase `--cp_size` to enable context parallelism for longer sequence lengths (`nemo:26.10` container onwards).
 - `--seq_length 32768 --gbs 512` — 16.8M tokens/iteration, 1.7B per 100 iterations.
 - `--lr 1e-5 --min_lr 1e-6` — an order of magnitude below typical distillation LRs: the job is to adapt weights to quantization, not to learn the task.
-- `--logit_kl_topk 4096` — restricts the KD loss to the teacher's top-4096 vocab entries. With a 248,320-token vocabulary the dense `[seq, vocab]` fp32 logits are **30.31 GiB per tensor** at 32K, which OOMs on its own.
+- `--logit_kl_top_k 4096` — restricts the KD loss to the teacher's top-4096 vocab entries. With a 248,320-token vocabulary the dense `[seq, vocab]` fp32 logits are **30.31 GiB per tensor** at 32K, which OOMs on its own.
 - `--recompute_*` / `--no_async_save` / `--eval_iters 0` — all needed to fit. Async save spawns a worker needing its own CUDA context; the validation path computes full-vocab LM and MTP cross-entropy (top-k applies to training only), so eval OOMs at 32K even though training fits.
 
 </details>
@@ -322,7 +322,7 @@ It is not verbosity. It is a **failure to terminate on a small fraction of sub-s
 - The **median** also roughly doubles (+91.6%), so the whole distribution shifted right — this is not *only* a tail effect.
 - Capped rate peaks at **iteration 50** (4.3%) and settles at 3.1% / 3.6% by 300 / 500; it is not gradual drift.
 
-The obvious suspect — that `--logit_kl_topk 4096` leaves the stop tokens outside the loss — **did not hold up**. Probing the BF16 teacher over one runaway trace: `</think>` does fall outside top-4096 at 35% of positions overall, but *in the looping region* the teacher gives `<|im_end|>` a median rank of **5** and `</think>` ~570, both well inside top-k. The teacher is signalling "stop here" at positions the loss did cover, and the student still does not stop. More likely: the blend has few "the answer is written, now stop" positions in this style, and a teacher-forced loss never exercises free-running generation 10K+ tokens deep.
+The obvious suspect — that `--logit_kl_top_k 4096` leaves the stop tokens outside the loss — **did not hold up**. Probing the BF16 teacher over one runaway trace: `</think>` does fall outside top-4096 at 35% of positions overall, but *in the looping region* the teacher gives `<|im_end|>` a median rank of **5** and `</think>` ~570, both well inside top-k. The teacher is signalling "stop here" at positions the loss did cover, and the student still does not stop. More likely: the blend has few "the answer is written, now stop" positions in this style, and a teacher-forced loss never exercises free-running generation 10K+ tokens deep.
 
 **It is fixable at decode time.** Adding `presence_penalty: 1.5` (Qwen's own thinking-mode recommendation for this model) removes nearly all of it, with no retraining. Every cell is SciCode **without → with** the penalty, 8 runs per side. *Capped* = hit the 131,072-token limit with no stop token; almost all such sub-steps return nothing and score zero. Counts are pooled over all 8 runs, so the denominator is 338 × 8 = 2,704 sub-steps:
 
diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py
index 5f6a1de4fb7..3aac2666650 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
@@ -50,6 +51,10 @@
 logger = logging.getLogger(__name__)
 
 
+# Former DistillationConfig fields that now raise if passed (see ``DistillationConfig.__new__``).
+_REMOVED_DISTILLATION_CONFIG_FIELDS = frozenset({"skip_lm_loss", "kd_loss_scale"})
+
+
 @dataclass
 class DistillationConfig:
     """Knowledge-Distillation config.
@@ -57,28 +62,59 @@ 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]. Default: ``1.0``. When ``1.0``,
+            the standard language model loss is skipped entirely (see :attr:`skip_lm_loss`).
         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_topk: If not None, use TopLogitsKLLoss 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.
     """
 
     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 = 1.0
     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
     criterion: Criterion | None = None
     loss_balancer: mtd.DistillationLossBalancer | None = None
 
+    def __new__(cls, *args, **kwargs):
+        """Reject removed fields with a migration hint before the dataclass ``__init__`` runs.
+
+        Done here rather than in ``__init__`` so the check survives subclasses that re-apply
+        ``@dataclass`` (which regenerates ``__init__``).
+        """
+        removed = _REMOVED_DISTILLATION_CONFIG_FIELDS & kwargs.keys()
+        if removed:
+            raise ValueError(
+                f"DistillationConfig {sorted(removed)} have been removed. Use `kd_loss_alpha` "
+                "instead: the total loss is (1 - kd_loss_alpha) * lm_loss + kd_loss_alpha * "
+                "kd_loss, and the LM loss is skipped when kd_loss_alpha == 1.0 (the default, "
+                "equivalent to the old skip_lm_loss=True)."
+            )
+        return super().__new__(cls)
+
+    @property
+    def skip_lm_loss(self) -> bool:
+        """Whether the standard LM loss is skipped, i.e. ``kd_loss_alpha == 1.0``."""
+        return self.kd_loss_alpha == 1.0
+
     def __post_init__(self):
         assert len(self.logit_layers) == 2, f"{self.logit_layers=}"
         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=}"
         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]:
@@ -128,10 +164,14 @@ def setup_distillation_config(
     if cfg.criterion is None:
         criterion = {}
         if parallel_state.is_pipeline_last_stage():
-            # Use TopKLogitsKLLoss if logit_kl_topk is specified, otherwise use LogitsKLLoss
+            # Use TopLogitsKLLoss 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
+                criterion[tuple(cfg.logit_layers)] = TopLogitsKLLoss(
+                    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,
                 )
             else:
                 criterion[tuple(cfg.logit_layers)] = LogitsKLLoss(
@@ -157,7 +197,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=cfg.skip_lm_loss,
         )
 
     return cfg
@@ -324,59 +365,58 @@ def forward(self, predictions: Tensor, targets: Tensor) -> Tensor:
         output_teacher = targets.float() / self._temperature
         output_student = predictions.float() / self._temperature
 
-        # 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()
+        # 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)
 
-            # 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
+        # KL divergence
+        if self._reverse:
+            p, q = q, p
+        loss = torch.sum(F.kl_div(p, q, reduction="none", log_target=True), dim=-1)
 
-            student_logits_max, _ = torch.max(output_student, dim=-1, keepdim=True)
-            torch.distributed.all_reduce(
-                student_logits_max,
-                op=torch.distributed.ReduceOp.MAX,
-                group=tp_group,
-            )
-            output_student -= student_logits_max.detach()
+        return self.post_forward(loss, tp_reduce=True)
 
-            # Compute global softmax denominators
-            # 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)
+    def _tp_logsumexp(self, logits: Tensor) -> Tensor:
+        """Log-sum-exp over the vocab dim across all TP shards (shape ``[..., 1]``).
 
-            denom_student = torch.sum(torch.exp(output_student), dim=-1, keepdim=True)
-            denom_student = dist_nn.functional.all_reduce(denom_student, group=tp_group)
+        ``logits`` are expected to be fp32 and already temperature-scaled.
+        """
+        if self._config.tensor_model_parallel_size == 1:
+            return torch.logsumexp(logits, dim=-1, keepdim=True)
 
-            # Compute log probabilities (log softmax)
-            teacher_log_prob = output_teacher - torch.log(denom_teacher)
-            student_log_prob = output_student - torch.log(denom_student)
+        tp_group = parallel_state.get_tensor_model_parallel_group()
 
-            # 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)
+        # Subtract maximum value along vocab dimension across all GPUs (for stability)
+        # Detached before the (non-autograd, in-place) collective: the max is only a stability shift.
+        logits_max = logits.amax(dim=-1, keepdim=True).detach()
+        torch.distributed.all_reduce(logits_max, op=torch.distributed.ReduceOp.MAX, group=tp_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)
+        # 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 self.post_forward(loss, tp_reduce=True)
+        return logits_max + torch.log(denom)
 
 
-class TopKLogitsKLLoss(LogitsKLLoss):
-    """Calculates KL-Divergence loss restricted to the Teacher's Top-K vocabulary entries.
+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 memory and communication.
+    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.
+
+    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
+    probability mass outside the kept entries, ``log(1 - sum(kept probs))``, is appended to both
+    student and teacher, so the KL is taken between two proper distributions over K + 1 buckets and
+    also penalizes mass the student places outside the teacher's kept entries.
+
+    Optionally, **Top-P (nucleus)** truncation keeps only the smallest prefix of the teacher-sorted
+    Top-K whose cumulative teacher mass reaches ``top_p`` (with a floor of ``top_p_min_k`` entries);
+    the truncated entries' mass moves into the ghost token.
     """
 
     def __init__(
@@ -385,6 +425,9 @@ def __init__(
         temperature: float = 1.0,
         reverse: bool = False,
         top_k: int = 1024,
+        *,
+        top_p: float | None = None,
+        top_p_min_k: int = 1,
     ):
         """Constructor.
 
@@ -393,9 +436,16 @@ 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.
         """
         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
 
     def forward(self, predictions: Tensor, targets: Tensor) -> Tensor:
         """Forward function.
@@ -414,14 +464,21 @@ 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).
+        # 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
+
+        # 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
 
         if tp_size > 1:
             tp_group = parallel_state.get_tensor_model_parallel_group()
@@ -446,36 +503,87 @@ 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 - teacher_lse
+        student_logp = final_student_logits - student_lse
+
+        # 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.
+        # 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.
+        # The floor keeps log(kept_mass) strictly below 0 so expm1 stays negative. Deliberate side
+        # effects once kept mass exceeds 1 - eps: the ghost bucket is pinned at ~eps (so the loss can
+        # dip below 0 by O(eps)), and the clamp stops the ghost term's gradient. Both are negligible,
+        # and the latter is the correct limit when the kept entries hold all of the mass.
+        neg_tiny = -torch.finfo(student_logp.dtype).eps
+        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)
+
+        # 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)
 
 
+class TopKLogitsKLLoss(TopLogitsKLLoss):
+    """Deprecated alias of :class:`TopLogitsKLLoss`."""
+
+    def __init__(self, *args, **kwargs):
+        """Constructor. Emits a ``FutureWarning`` and forwards to :class:`TopLogitsKLLoss`."""
+        warnings.warn(
+            "TopKLogitsKLLoss is deprecated and will be removed in a future release; "
+            "use TopLogitsKLLoss instead.",
+            FutureWarning,
+            stacklevel=2,
+        )
+        super().__init__(*args, **kwargs)
+
+
 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 = 1.0, 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:
@@ -492,22 +600,25 @@ def forward(self, loss_dict: dict[str, Tensor]) -> Tensor:
             if "Logits" in _key:  # class name
                 logits_key = _key  # should only be one
         logits_loss = loss_dict.pop(logits_key)
-        intermediate_loss = sum(loss_dict.values()) / max(len(loss_dict), 1)
-
-        if intermediate_loss > 0:
-            dynamic_scale = logits_loss.detach() / intermediate_loss.detach()
+        # Rescale intermediate losses to the logits-loss magnitude, without a host sync.
+        if loss_dict:
+            intermediate_loss = sum(loss_dict.values()) / len(loss_dict)
+            denom = intermediate_loss.detach()
+            dynamic_scale = torch.where(
+                denom > 0,
+                logits_loss.detach() / denom.clamp(min=torch.finfo(denom.dtype).tiny),
+                torch.zeros_like(denom),
+            )
             intermediate_loss_scaled = intermediate_loss * dynamic_scale
         else:
-            intermediate_loss = logits_loss.new_tensor(intermediate_loss)
+            intermediate_loss = logits_loss.new_zeros(())
             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,
@@ -591,7 +702,8 @@ def _sharded_state_dict(self, *args, **kwargs) -> "ShardedStateDict":
     # An MTP head left out of quantization is exempt from that: there is no quantization
     # error to recover there, and its CE materialises an fp32 [seq, vocab] tensor.
     skip_mtp_loss = _mtp_excluded_from_quantization(model)
-    if distill_cfg.skip_lm_loss and skip_mtp_loss:
+    skip_lm_loss = distill_cfg.skip_lm_loss
+    if skip_lm_loss and skip_mtp_loss:
         # Freeze the untrained MTP head: DDP's overlapped grad reduce asserts on params with no grad.
         warn_rank_0("MTP head is outside quantization and its loss is skipped: freezing it.")
         with model.hide_teacher_model():
@@ -603,7 +715,7 @@ def _compute_student_lm_loss(self, labels, logits) -> Tensor:
         self._lm_loss_call_count += 1
         mtp_num_layers = self.config.mtp_num_layers or 0
         is_mtp_call = self._lm_loss_call_count <= mtp_num_layers
-        if distill_cfg.skip_lm_loss and self.training and (not is_mtp_call or skip_mtp_loss):
+        if skip_lm_loss and self.training and (not is_mtp_call or skip_mtp_loss):
             return torch.zeros_like(labels, dtype=logits.dtype)
         return type(self).compute_language_model_loss(self, labels, logits)
 
diff --git a/tests/examples/megatron_bridge/test_qad.py b/tests/examples/megatron_bridge/test_qad.py
index b54dffa16d1..0360f55b5ef 100644
--- a/tests/examples/megatron_bridge/test_qad.py
+++ b/tests/examples/megatron_bridge/test_qad.py
@@ -118,7 +118,7 @@ def test_qad(tmp_path: Path, num_gpus, create_student):
         seq_length=16,
         mbs=1,
         gbs=4,
-        logit_kl_topk=8,
+        logit_kl_top_k=8,
         train_iters=train_iters,
         lr_warmup_iters=2,
         eval_interval=early_exit_iter,
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 f060dff3fd9..9985f9778b0 100644
--- a/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py
+++ b/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py
@@ -13,11 +13,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import dataclasses
 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
@@ -25,6 +28,10 @@
 import modelopt.torch.distill as mtd
 from modelopt.torch.distill.plugins.megatron import (
     DistillationConfig,
+    LogitsAndIntermediatesLossBalancer,
+    LogitsKLLoss,
+    TopKLogitsKLLoss,
+    TopLogitsKLLoss,
     _mtp_excluded_from_quantization,
     adjust_distillation_model_for_mcore,
     setup_distillation_config,
@@ -116,7 +123,9 @@ def _test_logits_kl_loss(rank, size):
 
     # Compute distillation loss
     loss = distillation_model.compute_kd_loss(
-        student_loss=student_loss, loss_reduction_fn=lambda x: x[0].mean()
+        # Reduce the per-token LM loss to a scalar, as Megatron's loss function does in training.
+        student_loss=student_loss.mean(),
+        loss_reduction_fn=lambda x: x[0].mean(),
     )
     assert isinstance(loss, dict), "Loss should be a dictionary"
     assert "kd_loss" in loss, "Should contain kd_loss key"
@@ -125,8 +134,8 @@ def _test_logits_kl_loss(rank, size):
     loss["kd_loss"].backward()
 
 
-def _test_topk_logits_kl_loss(top_k, rank, size):
-    """Test TopKLogitsKLLoss with simple forward/backward pass."""
+def _test_topk_logits_kl_loss(kd_kwargs, rank, size):
+    """Test TopLogitsKLLoss with simple forward/backward pass."""
     set_seed(SEED)
 
     num_layers = 2
@@ -168,9 +177,9 @@ def _test_topk_logits_kl_loss(top_k, rank, size):
         activation_func="squared_relu",
     ).cuda()
 
-    # Setup distillation config with TopKLogitsKLLoss via logit_kl_topk argument
+    # Setup distillation config with TopLogitsKLLoss 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,
     )
@@ -207,11 +216,19 @@ def _test_topk_logits_kl_loss(top_k, rank, size):
 
     # Compute distillation loss
     loss = distillation_model.compute_kd_loss(
-        student_loss=student_loss, loss_reduction_fn=lambda x: x[0].mean()
+        # Reduce the per-token LM loss to a scalar, as Megatron's loss function does in training.
+        student_loss=student_loss.mean(),
+        loss_reduction_fn=lambda x: x[0].mean(),
     )
     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()
 
@@ -264,7 +281,8 @@ def _test_skip_lm_loss_with_mtp(quantized, skip_lm_loss, rank, size):
         student_model.decoder.weight_quantizer = TensorQuantizer()
 
     distill_cfg = setup_distillation_config(
-        config_or_path=DistillationConfig(skip_lm_loss=skip_lm_loss),
+        # kd_loss_alpha == 1.0 skips the LM loss; any alpha < 1 keeps it.
+        config_or_path=DistillationConfig(kd_loss_alpha=1.0 if skip_lm_loss else 0.5),
         student_cfg=student_model.config,
         teacher_cfg=teacher_model.config,
     )
@@ -324,9 +342,121 @@ 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):
-    """Test TopKLogitsKLLoss with TP parallelism."""
-    dist_workers.run(partial(_test_topk_logits_kl_loss, top_k))
+@pytest.mark.parametrize(
+    ("top_p", "top_p_min_k"),
+    [(None, 1), (0.9, 1), (0.9, 3)],
+)
+def test_topk_logits_kl_loss(dist_workers, top_p, top_p_min_k, top_k: int = 5):
+    """Test TopLogitsKLLoss with TP parallelism."""
+    kd_kwargs = {
+        "logit_kl_topk": top_k,
+        "logit_kl_top_p": top_p,
+        "logit_kl_top_p_min_k": top_p_min_k,
+    }
+    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 = TopLogitsKLLoss(cfg, top_k=student.size(-1))(student, teacher)[0]
+    assert torch.allclose(dense, topk, atol=1e-6)
+
+
+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 = TopLogitsKLLoss(cfg, top_k=k)(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-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)
+
+
+@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 = TopLogitsKLLoss(cfg, temperature=temperature, top_k=k)(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-6)
+
+
+def test_topk_logits_kl_loss_top_p_masks_tail():
+    """Top-P moves out-of-nucleus mass into the ghost token 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)
+
+    def reference(keep):
+        partial = (keep * q_k.exp() * (q_k - p_k)).sum(-1)
+        q_rest = torch.log1p(-(q_k.exp() * keep).sum(-1))
+        p_rest = torch.log1p(-(p_k.exp() * keep).sum(-1))
+        return (partial + q_rest.exp() * (q_rest - p_rest)).transpose(0, 1)
+
+    probs = q_k.exp()
+    keep = (probs.cumsum(-1) - probs) < 0.5
+    assert not keep.all(), "test inputs should produce some truncation"
+    loss = TopLogitsKLLoss(cfg, top_k=k, top_p=0.5)(student, teacher)[0]
+    assert torch.allclose(loss, reference(keep), atol=1e-6)
+    assert loss.shape == (student.size(1), student.size(0))
+
+    # min_k floor forces at least min_k entries even when the nucleus is tiny.
+    min_k = 3
+    loss_min = TopLogitsKLLoss(cfg, top_k=k, top_p=1e-6, top_p_min_k=min_k)(student, teacher)[0]
+    assert torch.allclose(loss_min, reference(torch.arange(k) < min_k), atol=1e-6)
+
+    loss.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)
 
 
 @pytest.mark.parametrize(
@@ -370,3 +500,62 @@ 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)
+
+    # Zero intermediate loss contributes nothing (and no division by zero).
+    out = LogitsAndIntermediatesLossBalancer(kd_loss_alpha=1.0)(
+        {key: lm, "LogitsKLLoss_0": logits, "HiddenStateCosineLoss_0": torch.tensor(0.0)}
+    )
+    assert torch.allclose(out["kd_loss"], logits)
+    assert torch.isfinite(out["intermediate_loss"])
+
+    # 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_removed_fields():
+    """The removed fields raise, and a config can be rebuilt from itself."""
+    cfg = DistillationConfig()
+    assert cfg.kd_loss_alpha == 1.0 and cfg.skip_lm_loss  # pure KD by default
+    assert not DistillationConfig(kd_loss_alpha=0.9).skip_lm_loss
+    assert {"skip_lm_loss", "kd_loss_scale"}.isdisjoint(f.name for f in dataclasses.fields(cfg))
+
+    for removed in ({"skip_lm_loss": True}, {"skip_lm_loss": False}, {"kd_loss_scale": 2.0}):
+        with pytest.raises(ValueError, match="have been removed"):
+            DistillationConfig(**removed)
+
+    # skip_lm_loss is a derived property, not a field, so round-trips do not trip the check.
+    cfg = dataclasses.replace(DistillationConfig(kd_loss_alpha=0.9), logit_kl_topk=4)
+    assert cfg.kd_loss_alpha == 0.9 and cfg.logit_kl_topk == 4
+    assert DistillationConfig(**dataclasses.asdict(cfg)).kd_loss_alpha == 0.9
+
+
+def test_topk_logits_kl_loss_deprecated_alias():
+    """The old class name still works and warns."""
+    cfg = SimpleNamespace(tensor_model_parallel_size=1)
+    with pytest.warns(FutureWarning, match="use TopLogitsKLLoss instead"):
+        loss_fn = TopKLogitsKLLoss(cfg, top_k=4)
+    assert isinstance(loss_fn, TopLogitsKLLoss)