[arnavkomaragiri]

PR Dependencies

Depends on: NVIDIA-NeMo/RL#4146, NVIDIA/Megatron-LM#7387.

Required by: NVIDIA-NeMo/RL#4148.

Summary

Connect packed exact-call/tree layouts to policy execution for legacy and TQ training.

• Expand compatible independent calls or bounded tree fragments before DP sharding.
• Preserve logical episode ordering and reassemble sampled-edge logprobs.
• Use explicit tree ancestry/position metadata and select loss edges before vocabulary projection.
• Keep scheduler increments tied to logical samples rather than fragment count.
• Keep accumulated gradient buffers resident during open split steps, following the controller's existing keep_train_buffers contract; reject destructive offload and checkpoint saves while a step is open.
• Own optional MTP-loss bypass in NRL: freeze retained Hybrid MTP weights before DDP/optimizer construction and pass upstream compute_mtp_loss=False through every policy forward path. No checkpoint topology change or MCore config extension.
• Reject MTP bypass for unsupported MTP-bearing models; non-MTP models retain their normal path. Expose Mamba training-state dtype configuration separately.
• Reject unsupported draft/custom-loss/backend combinations instead of silently using ordinary attention.

Dependencies

Stacked after #4146, using representation/transport primitives from #4144 and #4145.

The sole required MCore change is NVIDIA/Megatron-LM#7387, for shared-prefix model execution. Gradient lifetime and Hybrid MTP control use existing upstream APIs and are implemented here, with their tests.

This draft is not merge-ready with the existing MCore pin. The feature imports the new tree API. Bridge has no new implementation in this series; its MCore gitlink and RL dependency pins must be reconciled after the dependencies are accepted. Local fork pins are deliberately not hidden in this feature diff.

Validation And Limitations

Matching policy, model-data, projection, split-step, and routing regressions are included. September 16 validation with only shared-prefix changes in MCore:

• Eight focused NRL config/lifecycle regressions passed.
• Six GPU MTP/lifecycle tests passed per rank: FP32/BF16 Hybrid logits/base gradients match upstream MTP bypass bit-exactly, checkpoint keys remain intact, MTP forward/loss are not executed, and two-chunk DDP accumulation survives an intervening logprob phase bit-exactly.
• Five CP/projection tests passed per rank, with unchanged tolerances.
• The combined async 1x16 TQ/R3/CP2 gate retained all 16 clean episodes, 485 successful calls and 653,885 generated tokens. No capture/sample/logprob masks or infrastructure losses. Gradient norm 0.921403; generation KL 0.001088. One optimizer step and refit completed.

The 1x16 run consumes one prompt group; inter-group gradient lifetime is additionally covered by the focused two-chunk GPU test, not a claim of a multi-step production run.

The known BF16 Mamba gradient discrepancy remains documented in the MCore draft; an upstream fused/unfused control also differs. This patch neither fixes that separate precision issue nor weakens its tests.

Scoped formatting, Python parsing, added-line privacy, and whitespace checks pass. The focused GPU tests above ran in the normal CUDA/TE container; full native CI, complete typechecking and latest-main reconciliation remain pending. Controller heartbeat and terminal-dispatch/shutdown warnings did not affect the consumed group and remain separate cleanup items.

Draft for author review only; no deployment recipes, datasets, or live jobs are changed.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.