feat(opd): add student top-k on-policy distillation

docs/about/algorithms/mopd.md

@@ -35,6 +35,20 @@ tool / environment tokens contribute zero. Because the advantage subtracts a
 real `prev_logprobs`, MOPD requires the student log-probabilities to actually be
 computed — see [Configuration](#configuration).
 
+### Student top-k reverse KL
+
+Setting `on_policy_distillation.student_topk` replaces the sampled-token MOPD
+loss with a lower-variance reverse-KL estimator. For every token, the student
+selects its highest-probability `k` vocabulary entries. NeMo RL evaluates the
+student and teacher exactly on that support and uses the sampled rollout token
+to estimate the contribution from the remaining vocabulary.
+
+The sampled tail uses the score-function coefficient
+`1 + log π_student − log π_teacher`. The `1` is required because the derivative
+also acts on the student probability multiplying the log-probability gap.
+Support indices and teacher probabilities are treated as constants; gradients
+flow only through the current student probabilities.
+
 ## Configuration
 
 Enable MOPD in two places: select the advantage estimator and add the
@@ -63,6 +77,9 @@ loss_fn:
 
 on_policy_distillation:
   enabled: true
+  # Optional: use the student's top-k vocabulary entries for an exact KL head
+  # and the sampled rollout token for the remaining tail.
+  # student_topk: 64
   # Map each NeMo Gym agent name to a teacher checkpoint.
   teacher_model_by_agent_name:
     default_teacher: Qwen/Qwen3-1.7B
@@ -99,6 +116,20 @@ on_policy_distillation:
 > `grpo.seq_logprob_error_threshold`), because the advantage would silently
 > degrade to `teacher_logprobs − 0`.
 
+> [!NOTE]
+> Student top-k mode currently requires async GRPO, a Megatron policy backend,
+> token-level loss, `loss_fn.disable_ppo_ratio: true`, sequence packing disabled,
+> context parallel size 1, fused linear log-probabilities disabled, and an
+> unfiltered training distribution (`generation.top_k: null`, `top_p: 1.0`).
+> CISPO, dual PPO clipping, and sequence-level importance ratios are unsupported.
+
+Student top-k training reports four additional metrics:
+
+- `opd_topk_head_loss`: exact reverse-KL contribution on the selected support.
+- `opd_topk_tail_loss`: sampled score-function surrogate for the tail gradient.
+- `opd_topk_student_mass`: student probability mass captured by the support.
+- `opd_topk_target_outside_fraction`: fraction of sampled targets outside it.
+
 ### Teacher routing
 
 Each rollout sample carries its NeMo Gym `agent_ref`. At collection time the
@@ -141,6 +172,22 @@ The reference recipe self-distills `Qwen/Qwen3-1.7B` (student == teacher) across
 student and teacher are identical, the OPD loss stays near zero — it is a
 correctness smoke test, not a demonstration of distillation gains.
 
+To exercise the student-top-k path with that recipe, disable its inherited
+sequence packing and set the support size explicitly:
+
+```sh
+uv run examples/nemo_gym/run_grpo_nemo_gym.py \
+  --config examples/configs/recipes/llm/mopd-qwen3-1.7b-3n8g-megatron-pack.yaml \
+  policy.sequence_packing.enabled=false \
+  on_policy_distillation.student_topk=64 \
+  data.train.data_path=/path/to/train.jsonl \
+  data.validation.data_path=/path/to/val.jsonl
+```
+
+This remains a smoke-test configuration. Measure throughput, peak memory, and
+loss/reward convergence against full-vocabulary MOPD before using it as a
+performance or quality baseline.
+
 ## References
 
 - LLM-Core Xiaomi, *MiMo-V2-Flash Technical Report*, which introduces the

nemo_rl/algorithms/async_utils/trajectory_collector.py

@@ -792,6 +792,131 @@ def _get_logprobs_for_group(group_key, indices):
 
         return result, total_time
 
+    def _compute_teacher_support_logprobs(
+        self,
+        input_ids: torch.Tensor,
+        topk_indices: torch.Tensor,
+        agent_refs: list[dict[str, Any]],
+        input_lengths: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, float]:
+        """Route student-selected supports to teachers and stitch their logprobs.
+
+        Args:
+            input_ids: Student token IDs with shape ``[batch, sequence]``.
+            topk_indices: Student-selected vocabulary indices with shape
+                ``[batch, sequence, k]``.
+            agent_refs: Per-sample agent metadata used for teacher routing.
+            input_lengths: Optional unpadded sequence lengths.
+
+        Returns:
+            Teacher support log-probabilities in original batch order and the
+            elapsed teacher-evaluation time.
+        """
+        if topk_indices.ndim != 3 or topk_indices.shape[:2] != input_ids.shape:
+            raise ValueError(
+                "topk_indices must have shape [B, S, K] aligned with input_ids; "
+                f"got {topk_indices.shape} and {input_ids.shape}."
+            )
+        opd_cfg = self.on_policy_distillation_cfg
+        reference_aliases = resolve_reference_aliases(
+            agent_refs,
+            opd_cfg.get("teacher_model_by_agent_name", {}),
+            default_teacher_alias=opd_cfg.get("default_teacher_alias"),
+            strict_agent_name_match=opd_cfg.get("strict_agent_name_match", False),
+        )
+        group_to_indices: dict[str, list[int]] = defaultdict(list)
+        for sample_index, alias in enumerate(reference_aliases):
+            group_key = self.alias_to_group_alias.get(alias, alias)
+            group_to_indices[group_key].append(sample_index)
+
+        batch_size, seq_len, support_size = topk_indices.shape
+        result = torch.zeros(batch_size, seq_len, support_size, dtype=torch.float32)
+        if not group_to_indices:
+            return result, 0.0
+
+        def _evaluate_group(group_key: str, indices: list[int]):
+            teacher = self.teacher_worker_groups[group_key]
+            sub_input_ids = input_ids[indices]
+            sub_topk_indices = topk_indices[indices]
+            sub_lengths = input_lengths[indices] if input_lengths is not None else None
+            actual_batch_size = sub_input_ids.shape[0]
+            dp_size = teacher.sharding_annotations.get_axis_size("data_parallel")
+            remainder = actual_batch_size % dp_size
+            if remainder:
+                pad_count = dp_size - remainder
+                sub_input_ids = torch.cat(
+                    [sub_input_ids, sub_input_ids[-1:].expand(pad_count, -1)], dim=0
+                )
+                sub_topk_indices = torch.cat(
+                    [
+                        sub_topk_indices,
+                        sub_topk_indices[-1:].expand(pad_count, -1, -1),
+                    ],
+                    dim=0,
+                )
+                if sub_lengths is not None:
+                    sub_lengths = torch.cat(
+                        [sub_lengths, sub_lengths[-1:].expand(pad_count)], dim=0
+                    )
+
+            sub_data = BatchedDataDict(
+                {"input_ids": sub_input_ids, "topk_indices": sub_topk_indices}
+            )
+            if sub_lengths is not None:
+                sub_data["input_lengths"] = sub_lengths
+            lock_started = time.time()
+            with self._teacher_locks[group_key]:
+                inference_started = time.time()
+                support_output = teacher.get_logprobs_on_support(sub_data)
+            finished = time.time()
+            print(
+                f"[teacher_support] group={group_key} samples={actual_batch_size} "
+                f"lock_wait={inference_started - lock_started:.2f}s "
+                f"inference={finished - inference_started:.2f}s"
+            )
+            return indices, support_output["support_logprobs"][:actual_batch_size]
+
+        started = time.time()
+        with concurrent.futures.ThreadPoolExecutor(
+            max_workers=len(group_to_indices)
+        ) as executor:
+            futures = [
+                executor.submit(_evaluate_group, group_key, indices)
+                for group_key, indices in group_to_indices.items()
+            ]
+            for future in concurrent.futures.as_completed(futures):
+                indices, support_logprobs = future.result()
+                result[indices] = support_logprobs
+        total_time = time.time() - started
+        return result, total_time
+
+    async def compute_teacher_support_logprobs(
+        self,
+        input_ids: torch.Tensor,
+        topk_indices: torch.Tensor,
+        agent_refs: list[dict[str, Any]],
+        input_lengths: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, float]:
+        """Run teacher support inference without blocking the actor event loop.
+
+        Args:
+            input_ids: Student token IDs with shape ``[batch, sequence]``.
+            topk_indices: Student-selected vocabulary indices with shape
+                ``[batch, sequence, k]``.
+            agent_refs: Per-sample agent metadata used for teacher routing.
+            input_lengths: Optional unpadded sequence lengths.
+
+        Returns:
+            Teacher support log-probabilities and evaluation time.
+        """
+        return await asyncio.to_thread(
+            self._compute_teacher_support_logprobs,
+            input_ids,
+            topk_indices,
+            agent_refs,
+            input_lengths,
+        )
+
     async def _iter_rollout_groups(
         self,
         repeated_batch: BatchedDataDict[DatumSpec],

nemo_rl/algorithms/grpo.py

@@ -382,6 +382,55 @@ def _validate_multimodal_dedup_capability(master_config: MasterConfig) -> None:
             "grpo.deduplicate_multimodal_data=true is currently supported "
             "only when data_plane.enabled=false."
         )
+def _validate_student_topk_config(
+    master_config: MasterConfig, student_topk: Optional[int]
+) -> None:
+    """Reject settings that the student-top-k OPD loss cannot honor."""
+    if student_topk is None:
+        return
+
+    grpo_config = master_config.grpo
+    policy_config = master_config.policy
+    generation_config = policy_config["generation"]
+    megatron_cfg = policy_config.get("megatron_cfg", {})
+
+    if not opd_module.is_opd_enabled(master_config):
+        raise ValueError(
+            "on_policy_distillation.student_topk requires "
+            "on_policy_distillation.enabled=true."
+        )
+    if not grpo_config.async_grpo.enabled:
+        raise NotImplementedError(
+            "Student-top-k OPD currently supports only grpo.async_grpo.enabled=true."
+        )
+    if policy_config["sequence_packing"]["enabled"]:
+        raise NotImplementedError(
+            "Student-top-k OPD does not yet support sequence packing."
+        )
+    if megatron_cfg.get("context_parallel_size", 1) != 1:
+        raise NotImplementedError(
+            "Student-top-k OPD does not yet support context parallelism."
+        )
+    if not megatron_cfg.get("enabled", False):
+        raise NotImplementedError(
+            "Student-top-k OPD currently requires the Megatron policy backend."
+        )
+    if need_top_k_or_top_p_filtering(
+        TrainingSamplingParams(
+            top_k=generation_config["top_k"],
+            top_p=generation_config["top_p"],
+        )
+    ):
+        raise NotImplementedError(
+            "Student-top-k OPD currently requires unfiltered training distributions "
+            "(policy.generation.top_k=null and top_p=1.0)."
+        )
+    if master_config.loss_fn.reference_policy_kl_penalty != 0:
+        raise ValueError(
+            "Student-top-k OPD requires loss_fn.reference_policy_kl_penalty=0."
+        )
+    if grpo_config.adv_estimator.name != "opd":
+        raise ValueError("Student-top-k OPD requires grpo.adv_estimator.name='opd'.")
 
 
 def setup(
@@ -604,6 +653,8 @@ def init_train_dataloader(dataset, suffix: str = ""):
     use_fused_linear_logprobs = bool(
         megatron_cfg.get("enabled") and megatron_cfg.get("use_fused_linear_logprobs")
     )
+    student_topk = opd_module.get_student_topk(master_config)
+    _validate_student_topk_config(master_config, student_topk)
     if use_fused_linear_logprobs:
         # Sequence packing is not yet validated with the fused path: the fused
         # forward rolls labels over the whole (packed) sequence and would mix
@@ -632,7 +683,9 @@ def init_train_dataloader(dataset, suffix: str = ""):
         )
 
     loss_fn = ClippedPGLossFn(
-        loss_config, use_fused_linear_logprobs=use_fused_linear_logprobs
+        loss_config,
+        use_fused_linear_logprobs=use_fused_linear_logprobs,
+        opd_student_topk=student_topk,
     )
 
     # Validate force_on_policy_ratio
@@ -4114,6 +4167,7 @@ def async_grpo_train(
         "Importance sampling correction must be enabled for async GRPO for good convergence due to off-policy samples!"
     )
     max_generation_failures = master_config.grpo.async_grpo.max_generation_failures
+    student_topk = opd_module.get_student_topk(master_config)
 
     if router_replay_enabled(master_config.policy) and (
         master_config.data_plane or {}
@@ -4626,11 +4680,13 @@ def async_grpo_train(
                     # Teacher logprobs are stored in batch dict by collection-time
                     # computation and padded by from_batches. Extract here.
                     trajectory_teacher_logprobs = None
+                    teacher_agent_refs = None
                     if opd_module.is_opd_enabled(master_config):
                         if "teacher_reference_logprobs" in repeated_batch:
                             trajectory_teacher_logprobs = repeated_batch[
                                 "teacher_reference_logprobs"
                             ]
+                        teacher_agent_refs = repeated_batch.get("agent_ref")
 
                     # Aggregate rollout metrics across groups with proper aggregation per metric type
                     per_group_metrics = {}
@@ -4769,6 +4825,40 @@ def async_grpo_train(
                             train_data["generation_logprobs"]
                         )
 
+                    if student_topk is not None:
+                        if trajectory_teacher_logprobs is None:
+                            raise ValueError(
+                                "Student-top-k OPD requires collection-time teacher "
+                                "target logprobs in the replay batch."
+                            )
+                        if not isinstance(teacher_agent_refs, list):
+                            raise ValueError(
+                                "Student-top-k OPD requires one agent_ref per training sample."
+                            )
+                        topk_output = policy.get_topk_logits(
+                            train_data,
+                            k=student_topk,
+                            timer=timer,
+                        )
+                        train_data["prev_topk_indices"] = topk_output["topk_indices"]
+                        (
+                            teacher_support_logprobs,
+                            teacher_support_time,
+                        ) = ray.get(
+                            trajectory_collector.compute_teacher_support_logprobs.remote(
+                                train_data["input_ids"],
+                                train_data["prev_topk_indices"],
+                                teacher_agent_refs,
+                                input_lengths=train_data["input_lengths"],
+                            )
+                        )
+                        train_data["teacher_support_logprobs"] = (
+                            teacher_support_logprobs
+                        )
+                        rollout_metrics["teacher_support_logprob_time"] = (
+                            teacher_support_time
+                        )
+
                     if not skip_reference_logprobs:
                         train_data["reference_policy_logprobs"] = (
                             policy.get_reference_policy_logprobs(
@@ -4806,6 +4896,10 @@ def async_grpo_train(
                     trajectory_teacher_logprobs = _pad_teacher_logprobs(
                         trajectory_teacher_logprobs, train_data["input_ids"].shape[1]
                     )
+                    if student_topk is not None:
+                        train_data["teacher_reference_logprobs"] = (
+                            trajectory_teacher_logprobs
+                        )
 
                 # Compute advantages with adv_estimator using correct mask and logprobs
                 with timer.time("advantage_calculation"):

nemo_rl/algorithms/loss/interfaces.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+# Copyright (c) 2025-2026, NVIDIA CORPORATION.  All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -52,6 +52,7 @@ class MetricNormalizer(enum.Enum):
 class LossInputType(enum.Enum):
     LOGIT = "logit"
     LOGPROB = "logprob"
+    OPD_STUDENT_TOPK = "opd_student_topk"
     DISTILLATION = "distillation"
     DISTILLATION_CROSS_TOKENIZER = "distillation_cross_tokenizer"
     DRAFT = "draft"
@@ -97,6 +98,7 @@ def __call__(
             **kwargs: Loss function input, which varies by input_type:
                 - For LossInputType.LOGPROB: next_token_logprobs (torch.Tensor)
                 - For LossInputType.LOGIT: logits (torch.Tensor)
+                - For LossInputType.OPD_STUDENT_TOPK: next-token and selected-support logprobs
                 - For LossInputType.DISTILLATION: student_topk_logprobs, teacher_topk_logprobs, H_all (torch.Tensor)
                 - For LossInputType.DISTILLATION_CROSS_TOKENIZER: logits (torch.Tensor), teacher_full_logits_by_idx (dict[int, torch.Tensor])
                 - For LossInputType.DRAFT: teacher_logits, student_logits, mask (torch.Tensor)