[arnavkomaragiri]

PR Dependencies

Depends on: none within this stack.

Required by: NVIDIA-NeMo/RL#4147.

Summary

Add opt-in shared-prefix training through TreePackedSeqParams, leaving ordinary packed sequences on their existing path.

• Represent DFS segments, ancestor relationships, CP-local queries, and sampled loss edges explicitly.
• Run causal attention over each query segment's ancestors, with autograd-aware CP K/V gathering and per-run activation recomputation.
• Preserve differentiable Mamba convolution/SSM state at branch points.
• Select edge-source hidden states before the vocabulary projection, including TP/SP ownership.
• Reject unsupported attention modifiers, inference paths, custom output processors, and incompatible MTP loss.

The trainer owns exact-call compatibility, fragment sizing, and loss ownership. No dense arbitrary attention mask or benchmark-specific logic is introduced in MCore.

Dependencies

Design discussion: #7386.

This is the sole MCore change required by the RL stack. Gradient-buffer lifecycle and optional MTP bypass are owned by NVIDIA-NeMo/RL#4147, using existing upstream APIs. No changes to TransformerConfig or gradient offload are included.

For Hybrid models with MTP, the caller must explicitly disable auxiliary computation using the existing compute_mtp_loss=False forward argument. MTP-bearing GPT training with tree packing is rejected; this PR does not add a GPT MTP bypass.

Validation And Known Limitations

• September 16 validation without either former MCore prerequisite: async 1x16 Harbor/OpenCode training with TQ, R3, and CP2 completed. All 16 episodes were clean; all 485 successful calls and 653,885 generated tokens matched capture, owned edges and trainer counts. No infrastructure masks; gradient norm 0.921403 and generation KL 0.001088.
• Five focused CP/projection tests passed on each of two GPU ranks. CP1 attention was exact; CP2 output/Q gradients were exact and worst K/V relative L2 error was 0.4792%, within the unchanged 2.5% threshold. TP2 sampled-edge projection backward passed.
• NRL-owned MTP bypass matched upstream forward control bit-exactly in FP32/BF16 logits and base gradients; real DDP accumulation across an intervening logprob phase was also bit-exact.
• BF16 Mamba gradient parity is not fully passing. The control also reproduces a discrepancy in unmodified upstream fused versus unfused execution; FP32 tree comparisons are close. No threshold was relaxed and no speculative precision fix is bundled. A passing training step is not a numerical proof.
• CP K/V gathering scales with fragment length; fragment bounds and interconnect costs need explicit performance review.
• The four known tree-Mamba parity cases were excluded from this lifecycle gate; their tests and tolerances remain unchanged. Full native CI and validation against newly advanced upstream main remain pending.
• Tested base: 54c62dfb6852e9218bf38590be9b90930bf7d589.

Review Status

Draft for author/design review only. Not ready for merge or reviewer assignment. The supported model/backend combinations and final upstream reconciliation remain readiness items.

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