Pack Megatron SFT batches with shared prefix trees

[FurtherAI]

Summary

Automatically prefix-pack Megatron SFT batches after tokenization and supervision-mask selection, without changing what callers submit or what constitutes an optimizer update.

• Reuse the existing prefix-tree planner to share identical unsupervised prefixes while preserving every supervised causal context and target. Supervised prediction positions are not shared.
• Pack once on the CPU, off the event loop, before trainer-rank fanout. Each typed packed payload still represents one original optimizer batch, learning rate, and progress event, even when it produces multiple rows.
• Schedule natural-length rows through the existing CP/DP/PP/VPP machinery, introducing only the required schedule dummies. Report logical tokens, executed token-equivalents, and loss-bearing tokens separately.
• Preserve existing local/serverless SFT callers. Do not change RL queues, autotuning, packing lookahead, or the RL packer's default behavior.

batch_size remains examples per optimizer update; packed_sequence_length is a per-row capacity, not an executed minimum. Packing neither silently truncates examples nor combines optimizer batches to fill capacity.

Correctness Fixes Found During Qualification

• Preserve admitted packing geometry. Bin-local rebuilding can move a branch point below the minimum shared-segment length and make a previously admitted bin exceed capacity. SFT retains the already admitted global geometry; RL keeps its existing default.
• Freeze Nemotron-H router bias during LoRA training. MCore updated this frozen base-model buffer independently of optimizer learning rate, including LR-zero comparisons. The bias is not part of LoRA export. Match the existing frozen-bias policy used for GLM.
• Preserve shared-expert stream handoffs across AOT compilation. AOT drops standalone wait_stream() graphs, allowing consumers to read unfinished shared-expert outputs. Keep the two ownership handoffs eager while retaining compiled math and side-stream overlap. The retained CUDA regression uses the real MCore output handoff and a delayed producer; without the fix, a computation expected to return 4 returns stale data (2).
• Apply the existing finite-result guard to SFT before saving/publishing an adapter.

Validation

Qualification is pinned to 78d835a5abc54bb07b8be80603495312f66e094d; this is not a claim of a new all-handler workflow pass.

Check	Result
Literal Nemotron/Lightning packing-invariance stage, unchanged original fixture and four scenarios	Worst logit MAPE 0.005732%, versus 0.005730% in the original PR; limit 0.5%
Six-layer production-width FP32 SFT, all/last assistant masks and two packing capacities	All four loss/gradient comparisons pass; worst gradient MAPE 0.29145%
Full 52-layer IEEE FP32, single packed row	Natural-routing gradient MAPEs 0.006405% / 0.071750%
Full 52-layer IEEE FP32, three packed rows, matching expert selections	Gradient MAPEs 0.056506% / 0.015543%; limit 0.5%
Compiled BF16 reduced Llama, CP1/CP2/DP2/PP2-VPP2	Worst gradient MAPE below 1%
CPU packing semantics	Exact context/target multiplicities, including all/last assistant masks; 100 randomized batches at three capacities
Shared-expert CUDA handoff regression	Pass, 9.40s
Real full-depth 64K SFT job on one H200	Three optimizer updates across two durable generations; warm schedules 19.30s / 19.14s; CPU packing approximately 30ms/update

Existing focused tests (53), formatting/type/lock checks, and the final exact-context test also passed during qualification. The real job verifies learning rates 1e-6/2e-6/3e-6, per-update progress, finite results, and adapter/optimizer persistence.

Numerical Scope

No thresholds were relaxed, production precision was not changed, and natural-routing full-depth comparisons are not reported as all passing.

Two independent effects explain the remaining differences:

1. Disabling PyTorch's TF32 flags alone does not produce an entirely IEEE FP32 reference: TE explicitly requests TF32 for FP32 GEMMs, and Mamba's Triton dots also default to TF32. Separate diagnostic controls were needed for both. A primitive reproducer shows a 0.000001907 input difference becoming a 0.000976563 output difference under TE TF32; IEEE preserves the original difference.
2. MoE top-k selection is discontinuous. The three-row natural FP32 cases reach 1.012% / 1.768% gradient MAPE with 11 changed expert selections. At the first change in the harder case, a 0.000267 score perturbation exceeds a 0.0000877 selection margin. Matching selected experts, while still recomputing differentiable routing probabilities, reduces these errors to 0.056506% / 0.015543%. Raw fixed-route repeats are exact. Full-depth BF16 natural routing reaches 6.52% (legacy preparer: 6.54%); the fixed-selection control is 3.79%, below the unchanged 5% bound.

These controls distinguish packing/normalization/backward correctness from changes in the selected MoE branches. No production route replay was added to SFT.

Diff Scope

13 files, +573/-49 lines: production +247/-49, integration tests +304, documentation +22. The implementation consists of the SFT packing adapter, typed submission and row scheduling, and the small correctness fixes above; it does not introduce another training orchestrator.