feat(sc): add the teacher top-k forward for distillation on SingleController

[tianyi-zhang-02]

What does this PR do?

Adds the teacher top-k forward primitive needed by distillation on SingleController:

• TQWorkerMixin.get_topk_logits_presharded runs the rank-local forward and writes both top-k tensors back to TransferQueue.
• TQPolicy.get_topk_logits_from_meta dispatches that work from the driver.
• The tests pin both output columns, the [B, S, k] shape, leader-only write-back, and batch-size validation.

This is the first PR in the #3843 → #3846 → #3849 stack. Nothing calls the new entry point until #3846.

Ownership

#2580 already contains the same get_topk_logits_presharded entry point and the same teacher_topk_logits / teacher_topk_indices columns as part of a larger TQ distillation implementation. That PR is still open, so ownership is unresolved. I am holding this stack here and will not add more downstream work until the maintainers choose which direction to keep.

Validation

• Current live head: 744463629719d1e33d49d85ff5f959ac1c4865de.
• Focused policy tests cover the worker and driver entry points.
• The full stack at test(sc): add the distillation recipe and functional test for SingleController #3849 head 907606048aa8c950027dbc199bff1495035e862c passed 765 unit/config cases (6 GPU-only skips) and a real 2-GPU teacher/student checkpoint-restore run.

This base branch is intentionally not being presented as current-main-ready while #2580 ownership remains unresolved; the two downstream branches are kept tested, but no further stack is being built on it.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[tianyi-zhang-02]

#3768 landed while this was open, so worth saying up front how the two relate — they both put a teacher on SingleController and they are not the same thing.

#3768 (MOPD): several teachers, each on its own reserved nodes (non_colocated_teachers), routed per NeMo-Gym agent. Their logprobs become a token-level advantage via adv_estimator.name='opd'. Entrypoint is get_teacher_logprobs_presharded, config lives under on_policy_distillation.

This stack: one frozen teacher, sharing the training GPUs as a second worker group and offloaded between forwards. Its top-k logits feed DistillationLossFn as a loss. Entrypoint is get_topk_logits_presharded, config is the distillation + teacher blocks — the SC port of examples/run_distillation.py.

So: different signal, different resource model, different objective. Both entrypoints now sit next to each other in TQWorkerMixin and the two config paths reject each other explicitly (a config setting both is an error, same as ppo + on_policy_distillation).

Rebased onto #3768. One real fix fell out of it: its new validation reads algo_cfg.adv_estimator, which DistillationConfig does not have, so a distillation config died with an AttributeError inside the validator. That is in #3849.

The one thing I would like a steer on: #3768 established teacher_worker_group as an SC teacher abstraction, and this stack predates it — it reuses TQPolicy instead, since a colocated teacher is a policy that never trains and needed no new driver. If you would rather have one teacher abstraction on this path I am happy to converge on it, but the resource models are different enough (teacher_worker_group reserves nodes; this one shares the trainer's) that I did not want to force it without asking. :)

[tianyi-zhang-02]

Duplicate flag, found while sweeping the open PRs. @pthombre's #2580 ("feat(distillation): add TransferQueue support for On-Policy Distillation") already adds get_topk_logits_presharded with the same signature and the same two column names, teacher_topk_logits / teacher_topk_indices. The only difference is the return value — theirs returns a transport ack, this returns None.

That PR is much larger (it also adds the write-back dispatcher, payload sizing and transport metrics), has been open since 2026-05-27, last touched 2026-06-26, and is CONFLICTING against main.

I've asked there whether it's being picked back up. If it is, I'll close this — no reason to have two. Flagging it here so a reviewer doesn't find the collision cold.