feat(sc): run the distillation teacher in the SingleController train pump

examples/run_grpo_single_controller.py

@@ -12,12 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Async GRPO / PPO launcher driven by the SingleController actor.
+"""Async GRPO / PPO / distillation launcher driven by the SingleController actor.
 
 Builds the full SC actor args driver-side via setup_single_controller and hands them
 to SingleControllerActor. Mirrors run_grpo.py for config loading so the same YAML
 files apply. data_plane.enabled=true is mandatory. A config carrying a `ppo:` block
-additionally brings up the PPO critic and trains it alongside the policy.
+additionally brings up the PPO critic and trains it alongside the policy; one
+carrying `distillation:` brings up the frozen teacher instead.
 """
 
 import argparse
@@ -30,10 +31,12 @@
 import ray
 from omegaconf import OmegaConf
 
+from nemo_rl.algorithms.distillation import check_vocab_equality
 from nemo_rl.algorithms.single_controller import SingleControllerActor
 from nemo_rl.algorithms.single_controller_utils import (
     MasterConfig,
     WatchdogConfig,
+    is_distillation_run,
     is_ppo_run,
     setup_single_controller,
 )
@@ -97,6 +100,12 @@ def main() -> None:
 
     if is_ppo_run(config):
         legacy_async_block, legacy_async = "ppo.async_ppo", config.ppo.async_ppo
+    elif is_distillation_run(config):
+        # DistillationConfig carries no legacy async block: distillation never
+        # had a v1 async path, so there is nothing here to reject. Without this
+        # branch the `else` below dereferences `config.grpo`, which is None on a
+        # distillation run.
+        legacy_async_block, legacy_async = "", None
     else:
         legacy_async_block, legacy_async = "grpo.async_grpo", config.grpo.async_grpo
     if legacy_async is not None:
@@ -127,6 +136,18 @@ def main() -> None:
     init_ray()
 
     tokenizer = get_tokenizer(config.policy["tokenizer"])
+    if is_distillation_run(config) and not os.environ.get(
+        "NRL_SKIP_DISTILLATION_TOKENIZER_CHECK"
+    ):
+        # The teacher writes top-k *indices*, which the loss reads as student
+        # vocabulary ids. A teacher on a different vocabulary produces indices
+        # that are silently wrong rather than an error, so check before any
+        # model loads. Same check and same opt-out as distillation.py.
+        check_vocab_equality(
+            tokenizer,
+            config.policy["model_name"],
+            config.teacher["model_name"],  # type: ignore[index]
+        )
     assert config.policy["generation"] is not None, (
         "A generation config is required for SC-driven async GRPO"
     )
@@ -178,6 +199,7 @@ def main() -> None:
             ("Generation", actor_args.gen_handle),
             ("Trainer", actor_args.trainer_handle),
             ("Value", actor_args.value_handle),
+            ("Teacher", getattr(actor_args, "teacher_handle", None)),
         ):
             if resource is None:
                 continue

nemo_rl/algorithms/single_controller.py

@@ -40,6 +40,7 @@
 
 import asyncio
 import contextlib
+import functools
 import hashlib
 import io
 import logging
@@ -86,6 +87,7 @@
     AdvantageConfig,
     MasterConfig,
     algo_config,
+    is_distillation_run,
     is_ppo_run,
     validate_sampler_buffer_capacity,
     validate_single_controller_config,
@@ -101,7 +103,12 @@
 from nemo_rl.data.interfaces import DatumSpec
 from nemo_rl.data_plane import DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, KVBatchMeta
 from nemo_rl.data_plane.async_utils import call_data_plane
-from nemo_rl.data_plane.schema import DP_CALIB_INPUT_FIELDS
+from nemo_rl.data_plane.schema import (
+    DP_CALIB_INPUT_FIELDS,
+    DP_DISTILLATION_TRAIN_FIELDS,
+    DP_TRAIN_FIELDS,
+    TEACHER_TOPK_FIELDS,
+)
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 from nemo_rl.distributed.refit_watchdog import RefitAborted, is_refit_context_lost
 from nemo_rl.environments.nemo_gym import should_use_nemo_gym
@@ -177,6 +184,14 @@ class SingleControllerActor:
     # tick, and it must exist on any instance the watchdog can reach.
     _recovering_from_refit: bool = False
 
+    # Class-level for the same reason -- an instance built with
+    # ``object.__new__``, which is how the unit tests reach a single pump
+    # without standing up a cluster, reads the GRPO/PPO shape by default rather
+    # than raising AttributeError.
+    _is_distillation: bool = False
+    _teacher: Optional["TQPolicy"] = None
+    _train_fields: tuple[str, ...] = DP_TRAIN_FIELDS
+
     def __init__(
         self,
         master_config: MasterConfig,
@@ -199,14 +214,19 @@ def __init__(
         self._algo_cfg = algo_config(master_config)
         self._async_cfg = master_config.async_rl
         self._is_ppo: bool = is_ppo_run(master_config)
+        self._is_distillation: bool = is_distillation_run(master_config)
         # GRPO has no epoch knob: it makes one optimizer step per RL step.
         self._ppo_epochs: int = self._algo_cfg.ppo_epochs if self._is_ppo else 1
 
-        self._policy_logprobs_required = not (
+        # DistillationLossFn takes only the sequence columns and the teacher's
+        # top-k. It forms no importance ratio and no reference KL, so neither
+        # logprob forward has a consumer -- and the two knobs that gate them do
+        # not exist on DistillationConfig.
+        self._policy_logprobs_required = not self._is_distillation and not (
             master_config.loss_fn.force_on_policy_ratio
             and self._algo_cfg.seq_logprob_error_threshold is None
         )
-        self._reference_logprobs_required = bool(
+        self._reference_logprobs_required = not self._is_distillation and bool(
             master_config.loss_fn.reference_policy_kl_penalty > 0
             and not self._algo_cfg.skip_reference_policy_logprobs_calculation
         )
@@ -215,6 +235,12 @@ def __init__(
         self._gen: Generation = actor_args.gen_handle
         self._trainer: TQPolicy = actor_args.trainer_handle
         self._value: Optional[TQValue] = getattr(actor_args, "value_handle", None)
+        self._teacher: Optional[TQPolicy] = getattr(actor_args, "teacher_handle", None)
+        # Distillation's loss reads a different set of columns, and the ones it
+        # skips are never written -- fetching those would error, not read zeros.
+        self._train_fields: tuple[str, ...] = (
+            DP_DISTILLATION_TRAIN_FIELDS if self._is_distillation else DP_TRAIN_FIELDS
+        )
         self._dataloader = actor_args.dataloader
         self._weight_synchronizer = actor_args.weight_synchronizer
         self._advantage_estimator = actor_args.advantage_estimator
@@ -1801,9 +1827,10 @@ async def _train_pump(self) -> None:
                                     self._trainer.get_reference_policy_logprobs_from_meta,
                                     train_meta,
                                 )
-                    elif self._is_ppo:
+                    elif self._is_ppo or self._is_distillation:
                         # prepare_for_lp_inference is skipped here, and it is the only
-                        # other call that parks the policy optimizer before the critic.
+                        # other call that parks the policy optimizer before the critic
+                        # or the teacher.
                         with self._timer.time("value_inference_prep"):
                             await asyncio.to_thread(self._trainer.offload_to_cpu)
 
@@ -1813,12 +1840,22 @@ async def _train_pump(self) -> None:
                             await asyncio.to_thread(self._trainer.finish_inference)
                             train_meta = await self._value_stage(train_meta)
 
-                    # Compute advantages
-                    with self._timer.time("advantage_calculation"):
-                        (
-                            train_meta,
-                            has_valid_training_tokens,
-                        ) = await self._advantage_stage(train_meta)
+                    # Teacher forward
+                    if self._is_distillation:
+                        with self._timer.time("teacher_logprob_inference"):
+                            await asyncio.to_thread(self._trainer.finish_inference)
+                            train_meta = await self._teacher_stage(train_meta)
+
+                    # Compute advantages. Distillation has none: its loss reads the
+                    # teacher's top-k directly and never forms a return, so there is
+                    # no reward to turn into one.
+                    has_valid_training_tokens = True
+                    if not self._is_distillation:
+                        with self._timer.time("advantage_calculation"):
+                            (
+                                train_meta,
+                                has_valid_training_tokens,
+                            ) = await self._advantage_stage(train_meta)
 
                     # A PPO step is this one chunk, so a chunk with nothing left
                     # after filtering is a step that trains neither model.
@@ -1872,8 +1909,11 @@ async def _train_pump(self) -> None:
                                         )
                                         step_open = True
                                     await asyncio.to_thread(
-                                        self._trainer.train_microbatches_from_meta,
-                                        train_meta,
+                                        functools.partial(
+                                            self._trainer.train_microbatches_from_meta,
+                                            train_meta,
+                                            train_fields=self._train_fields,
+                                        )
                                     )
                                     # A PPO step is one chunk: nothing to
                                     # accumulate, so close every epoch here.
@@ -3117,6 +3157,35 @@ async def _value_stage(self, meta: KVBatchMeta) -> KVBatchMeta:
         await asyncio.to_thread(self._value.finish_inference)
         return meta.with_fields([self._advantage_cfg.values_field])
 
+    async def _teacher_stage(self, meta: KVBatchMeta) -> KVBatchMeta:
+        """Run the distillation teacher's top-k forward over the selected chunk.
+
+        Same shape as ``_value_stage``: tensors never touch SC, the workers
+        fetch the sequence columns from DataPlane and commit the teacher's
+        scoring back under ``teacher_topk_logits`` / ``teacher_topk_indices``,
+        which the distillation loss then reads alongside the student's own.
+
+        The teacher is loaded and offloaded around the call so it holds the
+        training GPUs only for the duration of the forward -- ``distillation.py``
+        does the same, and with ``init_optimizer=False`` there is no optimizer
+        state to park.
+
+        Returns:
+            The batch metadata with the two teacher columns recorded on it.
+        """
+        assert self._teacher is not None, (
+            "_teacher_stage requires a teacher; setup builds one only on a "
+            "distillation run."
+        )
+        await asyncio.to_thread(self._teacher.prepare_for_lp_inference)
+        await asyncio.to_thread(
+            self._teacher.get_topk_logits_from_meta,
+            meta,
+            self._algo_cfg.topk_logits_k,
+        )
+        await asyncio.to_thread(self._teacher.offload_after_refit)
+        return meta.with_fields(TEACHER_TOPK_FIELDS)
+
     async def _value_train(self, meta: KVBatchMeta) -> dict[str, Any]:
         """Run one value model optimizer step against this chunk's GAE returns.
 

nemo_rl/algorithms/single_controller_utils/__init__.py

@@ -21,6 +21,7 @@
     RolloutFailureConfig,
     WatchdogConfig,
     algo_config,
+    is_distillation_run,
     is_ppo_run,
 )
 from nemo_rl.algorithms.single_controller_utils.setup import (
@@ -36,6 +37,7 @@
     "SingleControllerActorArgs",
     "WatchdogConfig",
     "algo_config",
+    "is_distillation_run",
     "is_ppo_run",
     "setup_single_controller",
 ]