perf(distillation): defer teacher top-k Ray payloads

[yuchenwang3]

What does this PR do?

Keep DP-local teacher top-k payloads in the Ray object store and pass their references directly to matching student workers instead of materializing the large tensors in the driver.

The fast path is enabled only when teacher and student have compatible static DP layouts. Dynamic batching, mismatched DP sizes, incompatible sequence-packing layouts, and partial batches continue to use the existing materialized path.

Worker methods return a small wrapper containing the worker-local payload reference. The driver still materializes every wrapper, so failures from replicated workers propagate before representative payload references are selected. Ray sharding also preserves existing ObjectRefs rather than wrapping them in another object-store entry.

Issues

Closes #1877

Usage

No user-facing configuration change is required. Compatible distillation runs select the deferred transfer path automatically.

Validation

Exact source under test:

• Base: 71a08314bc3257a15ef04b507309649e3d3f0f68
• Patch: 20cfa59cc405067629d1c99562f10687f396bfb3
• Patch tree: 036dd95966651201da0998d009f939dbec630815
• Source clean before and after testing

Focused tests:

• tests/unit/models/policy/test_deferred_topk.py: 8 passed, 0 skipped
• tests/unit/distributed/test_worker_groups.py: 28 passed, 1 unrelated Nsight environment skip
• tests/unit/algorithms/test_distillation.py: 23 passed, 0 skipped
• The ObjectRef preservation and non-representative worker error-propagation tests both executed and passed

Controlled 8×B200 Ray transfer benchmark:

Payload	Repeats	Base median / p95	Patch median / p95	Speedup	Base / Patch driver RSS delta
0.25 GB	3	1.018 / 1.037 s	0.229 / 0.329 s	4.44×	752.9 MB / 102 KB
1 GB	3	4.194 / 4.215 s	0.929 / 1.212 s	4.51×	2.75 GB / 98 KB
4 GB	3	17.644 / 17.947 s	3.697 / 3.857 s	4.77×	11.00 GB / 94 KB
19.7 GB	2	90.927 / 90.954 s	17.115 / 17.121 s	5.31×	54.18 GB / 303 KB

Every measured run used 8 actors, completed all repeats with exit 0, matched checksums, resolved ObjectRefs automatically, and reported zero object spilling. The final-patch harness models the production path: worker-local ray.put, small wrapper materialization in the driver, then direct payload-reference handoff. A separate probe confirmed that an error from a non-representative worker still fails the driver.

This is a controlled transfer benchmark for the bottleneck mechanism rather than a claim about full end-to-end training-step speedup.

Local checks also passed: Ruff lint, Ruff format check, Python compileall, and git diff --check.

Before your PR is "Ready for review"

• 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? No documentation change is needed because this is an automatic internal transfer optimization with no new user-facing configuration.

Additional Information

The commit is signed off under DCO by Yuchen Wang.

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

[yuchenwang3]

@nvidia-nemo/rl_reviewers_algo @nvidia-nemo/rl_reviewers_infra — final validation is complete on head 20cfa59cc405067629d1c99562f10687f396bfb3: the 8×B200 controlled transfer matrix passed at all four payload sizes with matching checksums, zero spill, and 4.44–5.31× speedup. Focused suites passed 8 + 28 + 23 tests; the only skip was an unrelated Nsight environment check. Could one of you route the review, add CI:L1, and approve /ok to test 20cfa59cc405067629d1c99562f10687f396bfb3? Thanks.