feat: add Energon-owned SFT sequence packing

[rohitrango]

Summary

• add minimal Energon-owned multimodal SFT packing
• add greedy and balanced greedy knapsack packing algorithms
• support producer-packed batches in Megatron SFTv2
• add Qwen2.5-VL and Nemotron Omni packing recipes

The implementation adds ~600 nonblank, non-comment Python lines. Configs and tests are excluded from that count.

Validation

• Ruff format, Ruff lint, and the standard allow-listed Pyrefly check passed
• 6 focused knapsack unit tests passed
• 3 focused Energon/Megatron packing unit tests passed with the mcore selector
• Qwen2.5-VL non-packing recipe completed 50 steps from a clean checkpoint; final loss 0.01768
• Qwen2.5-VL packing recipe trained through step 36 and saved step 30 before the validation node was repurposed
• Nemotron Omni non-packing recipe completed 50 steps; final loss 0.08346 and final gradient norm 5.47
• Nemotron Omni packing recipe completed 50 steps; 24,238 valid tokens, final loss 0.03185, and final gradient norm 3.1183
• Nemotron packed W&B run: https://wandb.ai/nvidia/sft-dev/runs/ub57wmfz

Convergence
Qwen2.5VL with and without packing

Nemotron-Omni with and without packing and CP>1

Previous PR and self-review

Previous PR: #4070. A self-review was performed on the previous PR, and the resulting feedback was addressed in this PR.

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

[copy-pr-bot]

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

Contributors can view more details about this message here.

[guyueh1]

prepacked MTP batches lose required position IDs.

When MTP is enabled, process_global_batch() creates mtp_loss_mask, and this prepacked branch transports that mask but unconditionally leaves position_ids=None. The existing NeMo-owned packing path explicitly constructs per-segment position IDs for MTP, and the Nemotron Omni MTP tests document that its hybrid model asserts they are present. setup_sft_v2() does not reject Energon packing combined with mtp_num_layers > 0, so that configuration can fail during forward.

[terrykong]

Checked this and it holds up. Adding the other half of it: nothing rejects the combination at setup time either.

Problem. setup_sft_v2 validates the Energon packing path — it requires sequence_packing.enabled with fuse_loss, checks the algorithm, and rejects dynamic batching. But the string mtp does not appear anywhere in that file. So data.energon.packing_buffer_size set together with policy.megatron_cfg.mtp_num_layers greater than 0 is accepted, and the run gets all the way to the model forward before it dies on the assert that this test describes — "HybridModel asserts position_ids are present whenever its MTP block runs". For a 30B-A3B job that is a whole allocation spent to learn that two config keys disagree.

The prepacked branch already knows MTP is in play: it transports mtp_loss_mask twelve lines earlier, then sets position_ids = None anyway.

AI-1

Add a setup-time check next to the other Energon guards in setup_sft_v2: if packing_buffer_size is set and mtp_num_layers is greater than 0, raise with a message naming both keys. Keep it whichever way the position_ids question lands — if you make the prepacked branch build positions the way the NeMo-owned branch does, the guard comes out in the same commit that adds a test covering the combination; if you do not, the guard is what stops someone hitting it. Failing at setup costs a second.

[guyueh1]

One important refinement: building position_ids here is sufficient only for text-only MTP; it does not make the pinned path correct for multimodal MTP. Nemotron Omni merges projected image/audio features into decoder_input, but the pinned MCore (1e7598cbfae8) reuses that tensor only for the main decoder. Its HybridModel then invokes MTP with input_ids/position_ids, and MTP rolls those IDs and calls the language embedding again. At media positions this conditions MTP on the placeholder-token embedding rather than the projected media embedding seen by the backbone.

Upstream draft Megatron-LM #7253 adds the missing precomputed-embedding MTP path: it rolls decoder_input directly and makes the token/position assertion conditional on no precomputed input. That support is absent from both current Megatron-LM main and this PRs pin.

So I think the setup-time rejection should remain until an equivalent MCore change is pinned. A position_ids-only patch may avoid the assertion but would silently train the multimodal MTP branch with the wrong conditioning embeddings. If support is added instead, it needs an end-to-end prepacked multimodal MTP test, not just a position-ID construction test.

[rohitrango]

This will mean we cannot run the comparison with production runs (which presumably uses the same setup in megatron-lm). Are you ok with that?

[guyueh1]

[P1] Packed padding alters MoE expert-bias updates
The new Nemotron packing recipes inherit moe_router_enable_expert_bias: true. Energon pads every physical pack to max_input_seq_length, but the prepacked Megatron path does not construct or forward a padding mask.
The pinned Nemotron wrapper explicitly drops that mask, while Megatron-Core excludes padding from local_tokens_per_expert only when a mask is supplied. Consequently, unused packing gaps are treated as real tokens when updating expert bias, silently changing MoE routing and training behavior.

i think we should fix this by considering to pass the correct padding mask to the router

[terrykong]

Strong PR. The Energon-owned packing path is cleanly factored, and the evidence behind it is better than most feature PRs get: convergence curves for both models packed and unpacked, loss and gradient norms, and a W&B link.

Worth calling out:

• You self-reviewed #4070 with 18 findings and genuinely fixed 16 of them here. We checked that against the code rather than against the "Fixed in <sha>" replies — including all 9 findings that never got a reply.
• The -104 in tests/unit/data/packing/test_algorithms.py is an improvement, not lost coverage. 17 test functions before, 17 after, identical names. Twelve repeated hardcoded algorithm lists collapsed into ALL_ALGORITHMS = list(PackingAlgorithm), read from the enum, so coverage widened from 4 algorithms to 6 by itself and a 7th would be picked up for free.
• The balanced-packer default fix is real and measured: 34 bins at 0.928 utilisation, tied with MFFD, where the old delta=20 default gave 0.384.
• Both of the PR's own new test files run green here: tests/unit/algorithms/test_sft_v2.py — 20 passed, and tests/unit/data/test_energon_packing.py — 9 passed. The packing one needed no stubbing at all; the packing math is module-level functions with no framework dependency, which is why it has real unit coverage even though all three recipes are disabled.
• The length = sum(...) recompute after truncation in generic_sft.py is a real bug you caught yourself — without it, a truncated sample carries its pre-truncation length into the packer.
• base.py's deleted sample_schema was already dead before this PR, so dropping it is safe. It is an unrelated cleanup riding in a feature PR.

One narrow ask: there is no throughput number. The convergence evidence is good and we are not asking for more of it, but packing exists to buy throughput, and you ran packed and unpacked for both models, so a step-time or tokens-per-second figure probably already exists in the same W&B runs. One sentence on why balanced_greedy_knapsack is the recipe default would help too, since the design doc this PR edits calls MFFD "the default recommendation" — we measured the two as tied, so this is a docs question, not a performance one.

What did and did not get checked, so you can weight the comments below:

• ruff lint and format pass on the changed files at the repo's pinned 0.9.9 (pyproject.toml:281) — All checks passed! and 698 files already formatted.
• That is a floor, not a substitute. It is not pyrefly, and we could not verify the PR's own claim that the allow-listed Pyrefly check passed.
• pre-commit as a whole could not run: uv run is blocked by the linux-only lockfile on macOS.
• No GPU. Nothing in _prepare_prepacked, _slice_prepacked_for_cp or the fused-loss path was executed by anyone, including the adversarial pass.

Two background write-ups while reviewing, if useful: who packs the sequences compares this path against the trainer-side packing already on main, and all packing paths maps it against the offline prepacking proposed in #3380.

Generated by Claude Code

[terrykong]

Checked this and it holds up. Adding the other half of it: nothing rejects the combination at setup time either.

Problem. setup_sft_v2 validates the Energon packing path — it requires sequence_packing.enabled with fuse_loss, checks the algorithm, and rejects dynamic batching. But the string mtp does not appear anywhere in that file. So data.energon.packing_buffer_size set together with policy.megatron_cfg.mtp_num_layers greater than 0 is accepted, and the run gets all the way to the model forward before it dies on the assert that this test describes — "HybridModel asserts position_ids are present whenever its MTP block runs". For a 30B-A3B job that is a whole allocation spent to learn that two config keys disagree.

The prepacked branch already knows MTP is in play: it transports mtp_loss_mask twelve lines earlier, then sets position_ids = None anyway.

AI-1

Add a setup-time check next to the other Energon guards in setup_sft_v2: if packing_buffer_size is set and mtp_num_layers is greater than 0, raise with a message naming both keys. Keep it whichever way the position_ids question lands — if you make the prepacked branch build positions the way the NeMo-owned branch does, the guard comes out in the same commit that adds a test covering the combination; if you do not, the guard is what stops someone hitting it. Failing at setup costs a second.