feat(distillation): add TransferQueue support for On-Policy Distillation

[pthombre]

What does this PR do ?

Adds TransferQueue-backed on-policy distillation so teacher top-k logits can stay in the data plane instead of being returned through the driver.

This PR extends the data-plane support that already exists for GRPO to the distillation workflow. When data_plane.enabled=true, examples/run_distillation.py now builds the student policy as a TQPolicy and dispatches to a new TQ-aware distillation trainer. The legacy distillation path remains the default when the data plane is disabled.

Key pieces added:

• A new distillation_train_sync implementation that orchestrates rollout, teacher top-k inference, and student training through TransferQueue metadata.
• A DistillationRolloutActor that runs generation/rollout processing and seeds the TQ partition with row-aligned distillation training tensors.
• Worker-side get_topk_logits_presharded support so teacher workers write teacher_topk_logits and teacher_topk_indices back to TQ directly.
• Distillation-specific data-plane schema fields for teacher top-k seed inputs and train-time top-k columns.
• TQPolicy extensions for custom partition fields, consumer task names, and custom train fetch fields so the same policy wrapper can serve GRPO and distillation.
• Transport-volume metrics for teacher top-k payloads, valid-token payloads, padding overhead, TQ write volume, and driver bytes avoided.
• A TQ exemplar config and a runnable recipe/test-suite script for Qwen3 32B-to-1.7B distillation.

Issues

None.

Usage

Enable the TQ distillation path by using the new exemplar or recipe:

uv run examples/run_distillation.py \
  --config examples/configs/distillation_math_tq.yaml \
  data_plane.enabled=true

The recipe added in this PR can be launched through the test-suite wrapper:

tests/test_suites/llm/distillation-qwen3-32b-to-1.7b-base-1n8g-fsdp2tp1-tq.v1.sh

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.

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

[pthombre]

/ok to test 5ce45ca

[pthombre]

/ok to test 5ce45ca

[pthombre]

/ok to test d2827ba

[pthombre]

/ok to test ff3a57c

[tianyi-zhang-02]

@pthombre — I built the same entrypoint in #3843 without knowing this PR existed, and I'd rather say so than let two copies sit in the queue.

#3843 adds TQWorkerMixin.get_topk_logits_presharded with the same signature (meta, k, micro_batch_size), writing back the same two columns under the same names — teacher_topk_logits / teacher_topk_indices. The only difference I can see is the return: yours returns a transport ack, mine returns None and relies on the write-back. That's a genuine duplicate, not two PRs that happen to share a file.

Yours is much larger and does things mine doesn't — dispatch_teacher_topk_writeback, topk_payload_nbytes, the transport-metrics aggregation. Mine is 200 lines total and stops at the two-layer primitive.

I'm not asking you to do anything. Flagging it because this PR has been open since May and is currently CONFLICTING against main, and a fair bit has moved under it — notably #3768 landed MOPD on SingleController on 2026-08-26, which added get_teacher_logprobs_presharded right next to where this inserts, plus TEACHER_LP_FIELDS and SC_ROLLOUT_SCHEMA_FIELDS in schema.py.

If you're picking this back up, say so and I'll close #3843 — no point in two. If it's dormant, #3843 is rebased on current main, mergeable, and CPU-tested, and I'd rather it not look like I lifted your design. Either outcome is fine by me. :)