feat(speechlm2): distill MTP heads with LK loss

docs/source/speechlm2/configs.rst

@@ -170,6 +170,36 @@ SALMAutomodel-Specific Options
 The SALMAutomodel config exposes a few extra knobs that pass through to NeMo
 Automodel. All are optional — defaults preserve standard behavior.
 
+**Multi-token prediction (MTP):**
+
+.. code-block:: yaml
+
+    model:
+      mtp:
+        enabled: true
+        training_mode: head_only       # "joint" or "head_only"
+        loss_type: lk                  # "cross_entropy" or "lk"
+        lk_lambda: 0.5                 # LK = lambda * KL + (1-lambda) * TV
+        loss_scaling_factor: 0.1
+        num_nextn_predict_layers: 2
+        use_repeated_layer: true
+        replace_existing_head: false
+
+``cross_entropy`` trains each MTP depth against its shifted dataset target and
+supports either joint training or head-only training. ``lk`` distills each MTP
+depth from the backbone's future-token distribution and requires
+``training_mode: head_only``; this keeps the backbone frozen as the teacher.
+``lk_lambda`` must be between 0 and 1. Packed-sequence boundaries are respected,
+and the teacher rows needed by each rank are exchanged when context parallelism
+is enabled.
+
+The progress-bar ``loss`` remains the backbone cross-entropy value for continuity
+with existing runs. The gradient-carrying objective also includes ``mtp_loss``;
+inspect that metric and ``mtp_lk_kl/head_N`` / ``mtp_lk_tv/head_N`` to monitor LK
+training. ``use_repeated_layer`` shares one physical MTP layer across the configured
+logical prediction depths. ``replace_existing_head: false`` preserves compatible
+MTP weights already present in a checkpoint.
+
 **MoE training:**
 
 .. code-block:: yaml

examples/speechlm2/conf/salm_automodel.yaml

@@ -34,6 +34,10 @@ model:
   mtp:
     enabled: false
     training_mode: joint  # joint preserves the recipe freeze policy; head_only trains only llm.mtp
+    # cross_entropy uses dataset labels. lk requires head_only and distills the
+    # frozen backbone distribution with lambda * KL(p||q) + (1-lambda) * TV(p,q).
+    loss_type: cross_entropy
+    lk_lambda: 0.5
     num_nextn_predict_layers: 1
     hybrid_override_pattern: "*"
     # Preserve a checkpoint-native MTP architecture by default. Set true to

nemo/collections/speechlm2/models/salm_automodel.py

@@ -35,6 +35,7 @@
 from nemo.collections.speechlm2.parts.hf_hub import HFHubMixin
 from nemo.collections.speechlm2.parts.mtp import (
     build_mtp_loss_fn,
+    calculate_mtp_lk_loss,
     calculate_mtp_loss_with_per_depth,
     calculate_mtp_teacher_forced_agreement,
     compute_mtp_agreement_lengths,
@@ -451,10 +452,19 @@ def training_step(self, dataloader_iter):
 
     def _training_step_batch(self, batch: dict, batch_idx: int):
         self._current_batch_idx = batch_idx
-        for m in (self.perception.preprocessor, self.perception.encoder, self.llm):
+        for m in (self.perception, self.perception.preprocessor, self.perception.encoder, self.llm):
             if is_frozen(m):
                 m.eval()
 
+        # Keep the root LLM in training mode so Automodel emits MTP states, but
+        # make the frozen teacher path deterministic during head-only LK training.
+        if getattr(self, "_mtp_loss_type", None) == "lk" and self._mtp_enabled:
+            self.llm.model.eval()
+            lm_head = getattr(self.llm, "lm_head", None)
+            if lm_head is not None:
+                lm_head.eval()
+            self.llm.mtp.train()
+
         inputs = self.prepare_inputs(batch)
         self._record_training_stats(batch, inputs)
         forward_outputs = self(
@@ -503,38 +513,56 @@ def _training_step_batch(self, batch: dict, batch_idx: int):
         with torch.no_grad():
             loss_display = loss_sum.detach() / num_frames.clamp(min=1)
 
-        # Multi-Token Prediction auxiliary loss. Compute the aggregate and per-head losses
-        # in one pass so WandB can show each MTP depth without repeating the expensive
-        # lm_head + CE work. ``mtp_loss`` keeps the same meaning as before: the weighted
-        # auxiliary loss added to the training objective after the DP-size correction.
+        # Multi-Token Prediction auxiliary loss. Cross-entropy remains the
+        # backward-compatible default; head-only LK training instead distills
+        # each draft depth from the frozen backbone distribution.
         mtp_metrics = {}
         mtp_h = forward_outputs.get("mtp_per_depth_h", None)
         if mtp_h is not None:
-            # Under packed THD multiple utterances share one token stream, so the
-            # per-depth label roll must not predict the next sequence's first token
-            # from the current sequence's last token. Pass cu_seqlens (empty/None for
-            # BSHD, where each row is already a single sequence) so the loss derives
-            # seq_idx and masks cross-sequence targets.
             mtp_per_depth_targets = inputs.get("mtp_per_depth_targets")
-            mtp_cu_seqlens = (
-                None if mtp_per_depth_targets is not None else inputs.get("llm_kwargs", {}).get("cu_seqlens")
-            )
-            with loss_parallel():
-                mtp_loss_output = calculate_mtp_loss_with_per_depth(
-                    self._mtp_loss_fn,
+            packed_cu_seqlens = inputs.get("llm_kwargs", {}).get("cu_seqlens")
+            mtp_loss_type = getattr(self, "_mtp_loss_type", "cross_entropy")
+            if mtp_loss_type == "lk":
+                cp_group = self._device_mesh["cp"].get_group() if self._context_parallel_size > 1 else None
+                mtp_loss_output = calculate_mtp_lk_loss(
                     mtp_per_depth_targets=mtp_per_depth_targets,
                     mtp_per_depth_h=mtp_h,
+                    teacher_logits=logits,
                     labels=inputs["target_ids"],
                     model=self.llm,
                     scaling_factor=self._mtp_loss_scaling_factor,
+                    lk_lambda=self._mtp_lk_lambda,
                     num_label_tokens=num_frames_global,
-                    grad_reduce_group=dp_group,
-                    cu_seqlens=mtp_cu_seqlens,
-                    return_per_depth=True,
+                    cu_seqlens=packed_cu_seqlens,
+                    context_parallel_group=cp_group,
+                    projection_sync_group=dp_group,
                 )
-                mtp_loss = mtp_loss_output.loss
                 mtp_raw_loss_by_head = mtp_loss_output.per_depth_losses
-            mtp_loss = dp_size * mtp_loss
+                for head_idx, (kl_loss, tv_loss) in enumerate(
+                    zip(mtp_loss_output.per_depth_kl, mtp_loss_output.per_depth_tv), start=1
+                ):
+                    mtp_metrics[f"mtp_lk_kl/head_{head_idx}"] = (dp_size * kl_loss).detach()
+                    mtp_metrics[f"mtp_lk_tv/head_{head_idx}"] = (dp_size * tv_loss).detach()
+            else:
+                # Precomputed targets are authoritative under CP; otherwise
+                # Automodel derives the same shifts from global cu_seqlens.
+                mtp_cu_seqlens = None if mtp_per_depth_targets is not None else packed_cu_seqlens
+                with loss_parallel():
+                    mtp_loss_output = calculate_mtp_loss_with_per_depth(
+                        self._mtp_loss_fn,
+                        mtp_per_depth_targets=mtp_per_depth_targets,
+                        mtp_per_depth_h=mtp_h,
+                        labels=inputs["target_ids"],
+                        model=self.llm,
+                        scaling_factor=self._mtp_loss_scaling_factor,
+                        num_label_tokens=num_frames_global,
+                        grad_reduce_group=dp_group,
+                        cu_seqlens=mtp_cu_seqlens,
+                        return_per_depth=True,
+                    )
+                    mtp_raw_loss_by_head = mtp_loss_output.per_depth_losses
+
+            mtp_loss = dp_size * mtp_loss_output.loss
             mtp_raw_loss_by_head = [dp_size * head_loss for head_loss in mtp_raw_loss_by_head]
             loss = loss + mtp_loss
             mtp_metrics["mtp_loss"] = mtp_loss.detach()
@@ -1204,18 +1232,30 @@ def configure_model(
         mtp_cfg = self.cfg.get("mtp", None)
         mtp_requested = mtp_cfg is not None and mtp_cfg.get("enabled", False)
         mtp_training_mode = str(mtp_cfg.get("training_mode", "joint")) if mtp_requested else "disabled"
+        mtp_loss_type = str(mtp_cfg.get("loss_type", "cross_entropy")) if mtp_requested else "disabled"
         if mtp_requested and mtp_training_mode not in {"joint", "head_only"}:
             raise ValueError(
                 f"Unknown mtp.training_mode {mtp_training_mode!r}; expected 'joint' or 'head_only' when MTP is enabled"
             )
-        logging.info(f"MTP training mode={mtp_training_mode}")
+        if mtp_requested and mtp_loss_type not in {"cross_entropy", "lk"}:
+            raise ValueError(
+                f"Unknown mtp.loss_type {mtp_loss_type!r}; expected 'cross_entropy' or 'lk' when MTP is enabled"
+            )
+        if mtp_loss_type == "lk" and mtp_training_mode != "head_only":
+            raise ValueError("mtp.loss_type='lk' requires mtp.training_mode='head_only' so the teacher stays frozen")
+        mtp_lk_lambda = float(mtp_cfg.get("lk_lambda", 0.5)) if mtp_requested else 0.5
+        if mtp_loss_type == "lk" and not 0.0 <= mtp_lk_lambda <= 1.0:
+            raise ValueError(f"mtp.lk_lambda must be in [0, 1], got {mtp_lk_lambda}")
+        logging.info(f"MTP training mode={mtp_training_mode}, loss type={mtp_loss_type}")
         if mtp_requested:
             # MTP supports both BSHD and packed THD. For THD the MTP loss must
             # receive cu_seqlens so target rolling is masked at packed sequence
             # boundaries (see training_step); the MTP sublayers already get the
             # THD context (qkv_format/cu_seqlens/seq_idx) from the model forward.
             self._mtp_loss_scaling_factor = float(mtp_cfg.get("loss_scaling_factor", 0.1))
-            self._mtp_loss_fn = build_mtp_loss_fn()
+            self._mtp_loss_type = mtp_loss_type
+            self._mtp_lk_lambda = mtp_lk_lambda
+            self._mtp_loss_fn = build_mtp_loss_fn() if mtp_loss_type == "cross_entropy" else None
             requested_depth = int(mtp_cfg.get("num_nextn_predict_layers", 1))
             use_repeated_layer = bool(mtp_cfg.get("use_repeated_layer", False))
             # HF/vLLM exports describe physical layers. A repeated MTP head has one