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 to the SingleController data path, so distillation can eventually run there. Nothing calls it yet.

Issues

Groundwork toward distillation on SingleController. No issue to close.

Usage

# Driver side, once the teacher is wired into the train pump:
teacher.get_topk_logits_from_meta(meta, k=cfg.top_k)
# The two tensors land in TQ as `teacher_topk_logits` / `teacher_topk_indices`.
# Nothing comes back through Ray.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Docs unchanged on purpose. There is no user-facing knob yet.

Additional Information

After #3773, distillation is the last algorithm with a rollout loop that SingleController can't run.

It turns out to need very little. The distillation teacher is a Policy, not a separate model class like the PPO critic, so there's no new driver to write — TQPolicy already is the teacher. All that was missing is the top-k entrypoint on the two layers every other forward already has:

• TQWorkerMixin.get_topk_logits_presharded — per-rank fetch → forward → write-back.
• TQPolicy.get_topk_logits_from_meta — the 1-hop dispatch. It reuses LP_SEED_FIELDS, since a teacher forward needs exactly what a logprob forward needs: the student's tokens and their masks. The teacher only adds its own scoring of them.

One thing is different from the sibling entrypoints. get_topk_logits returns two tensors, not one, and both are [B, S, k]. Every other *_presharded method writes back a single [B, S]. _write_back_result_field only validates the batch dimension, so the extra axis passes through fine — but lm_policy.py already carries a comment about from_batches flattening [B,S,k] into [B,S*k], so that axis clearly has a history of getting lost. There's a test pinning it.

I split this out from the wiring so it's testable on its own. The follow-up adds the teacher stage to the train pump and the config validation.

Tests: tests/unit/models/policy/test_tq_teacher_topk.py, same shape as test_tq_value.py — both columns written, the k axis survives, non-leader twins stay quiet, batch-dim mismatch raises. I mutation-tested them: dropping the indices write-back and pointing both writes at the same field each fail.

Not colliding with #3584 / #3608 — those are on the teacher_worker_group path and don't touch these files. :)

[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.