[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. Direct-packed SFT inputs, 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

October 8 main sync

• Head: 3487617a7b6a8512d9daa9afdda23088d631b099; main: fc4f3bcf57cebef3c453bc149be04231163420e1.
• Resolved the Policy training conflict while preserving direct-packed SFT's strided DP rows and microbatch override. Deferred teacher payloads retain their row-order checks; direct-packed inputs use the materialized fallback.
• Five CPU regression cases cover DP row ordering, valid/invalid packed microbatch overrides, and rejecting incompatible deferred packed layouts.
• Local tests: 82 passed, 1 Nsight skip, 2 collective-test failures on importing nccl, which is unavailable in this macOS environment. The run includes policy/deferred, packed SFT, distillation-driver, collective, and real CPU Ray worker-group tests; --noconftest omits GPU-monitor fixtures. Python 3.13.14, PyTorch 2.13.0, Ray 2.58.0; the dedicated 8-CPU Ray cluster was shut down afterward.
• Changed-file lint, format, import ordering, syntax, recipe minimization and diff checks pass. Full Pyrefly was not rerun. GPU/functional validation, transfer benchmarks and TQ integration remain pending on this head.

October 4 rebase

• Base: 31b4c0295cff085e74f2724e2cf057e5b3e3ae04
• Tested head: 733659d1f600d2d0bd2c8782362e81241cbc8b6f
• Keep the retired DTensor worker deleted; retain the deferred hook on DTensor v2 and Megatron, alongside the new upstream NVFP4 imports.
• The policy/deferred and distillation driver suites pass 33 tests through the normal package imports. The worker-group suite passes 30 tests with real local CPU Ray actors; one Nsight test is skipped because nsys is unavailable.
• These commands use macOS, Python 3.13.14, PyTorch 2.13.0, the configured test environment, and --noconftest to omit the repository autouse GPU-monitor fixtures. Policy/driver cases mock model loading and workers. The worker suite starts a dedicated local Ray cluster with 8 CPU resources and shuts it down afterward.
• Changed-file formatting/lint hooks, source parsing for all nine changed Python files, and git diff --check pass. The recipe-minimization hook was skipped locally. A targeted Pyrefly attempt exited 137 without diagnostics; it is not a type-check pass.
• GPU worker/functional testing, the transfer benchmark, and TQ integration remain pending on this head.

Earlier validation

The unit-test and B200 measurements below apply to patch 20cfa59cc405067629d1c99562f10687f396bfb3 and its stated base, rather than the October 4 rebase.

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 on the current head? 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.

Assisted-by: OpenAI Codex

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

[yuchenwang3]

/ok to test 5db71aa

[yuchenwang3]

@zpqiu, this PR addresses the teacher-result transfer bottleneck you reported in #1877 by deferring compatible DP-local top-k payloads through Ray object references. The PR documents controlled transfer measurements and fallback coverage; those measurements are not an end-to-end training speedup claim. Could you take a look at the approach or help identify the appropriate distillation reviewer? Maintainer CI authorization is also still needed. Thanks.

[yuchenwang3]

Synced with current main and resolved the policy-worker conflict, retaining both deferred top-k transfer and the new pre-refit synchronization hook. No change to the transfer algorithm.

Changed-file Ruff/format checks pass. Local isolated CPU checks confirm the deferred methods are unchanged and the upstream refit hook is preserved. The earlier Ray/GPU measurements remain results for the previously documented revision; I have not rerun that benchmark on this merge. Upstream runner validation still needs maintainer authorization.

[yuchenwang3]

@ananthsub Rebased at eff2617084, keeping the deferred top-k transfer logic and upstream pre-refit sync hook. Could you review the approach for #1877 or help route it? Local lint passes; the earlier Ray/GPU measurements haven't been rerun on this head.

[zpqiu]

Thank you for this contribution! Could we evaluate reusing the existing data-plane/TQ path for teacher top-k payloads? get_teacher_logprobs_presharded() already writes full teacher payloads directly to the data plane. I think @ZhiyuLi-Nvidia @zyzhou5 might have some suggestions on the TQ path.

[yuchenwang3]

Thanks @zpqiu, I traced the TQ path and agree it looks suitable for top-k too. The existing pack_jagged_fields() / materialize() codec handled [N, S, K] payloads in four local CPU checks: BF16/FP32 logits, int64 indices, unequal sequence lengths, and reordered rows. That probe explicitly marks the two fields token-aligned; it doesn't exercise a TQ backend or training.

The remaining wiring would be:

• Register teacher_topk_logits / teacher_topk_indices in the partition and token-aligned field registry.
• Add a top-k presharded entrypoint using the existing fetch, globally planned packing, and worker write-back helpers. get_teacher_logprobs_presharded() currently dispatches token logprobs or full hidden-state/logit payloads.
• Reuse TQPolicy.train_from_meta(train_fields=...) with distillation's field set, and connect partition setup/cleanup in the trainer. Distillation currently constructs plain Policy instances and passes BatchedDataDict batches.

I'd favor reusing those components over maintaining another transport. Would you prefer that I migrate distillation's step lifecycle to TQ in this PR and remove the deferred-Ray path? @ZhiyuLi-Nvidia @zyzhou5, guidance on that scope would help before I rework the driver.

[yuchenwang3]

@zpqiu @ZhiyuLi-Nvidia Could you confirm whether the TQ change should include the distillation driver lifecycle in this PR, or start with top-k field and worker support as a prerequisite? The codec probe is complete; the remaining work is the TQ worker entrypoint and trainer setup/cleanup. I'd like to settle that scope before replacing the deferred-Ray implementation.

[yuchenwang3]

Rebased onto 31b4c0295cff at 733659d1f600, resolving the merge conflicts. The retired DTensor worker stays deleted; the deferred hook remains on DTensor v2 and Megatron, with the upstream NVFP4 imports preserved.

The current-head CPU suites now pass 63 tests: 33 policy/deferred and distillation-driver cases, plus 30 worker-group cases with real local Ray actors. One Nsight case is skipped because nsys is unavailable. These use normal package imports on macOS/Python 3.13.14, with --noconftest omitting the autouse GPU-monitor fixtures and a dedicated 8-CPU Ray cluster. Changed-file lint/format and source parsing pass.

GPU worker/functional tests, the transfer benchmark, full type checking, and TQ integration remain open. The PR description separates these new checks from the earlier B200 results. The TQ lifecycle scope discussed above remains unresolved.

[yuchenwang3]

@zpqiu Synced with main and fixed the conflict. Should this PR migrate distillation setup/train/cleanup to TQ, or should I add the top-k worker support first?