Energon-owned packing: which half does Energon actually own?

NVIDIA-NeMo/RL #4105 · predecessor #4070 — all links pinned to f31f793.

Four pages on this PR — this one: the line between Energon's code and NeMo-RL's code inside the new path.

"Energon-owned packing" sounds like NeMo-RL handed the whole job away. It did not. Energon decides when to pack and which samples are on the table; NeMo-RL still decides which samples go in which bin and what the packed row looks like. The two sides take turns, four times, in one loop. The upstream Energon package ships no packing algorithms at all — it ships two empty hooks, select_samples_to_pack and pack_selected_samples, and expects you to fill them. That is why this PR adds new packing algorithms to NeMo-RL at the same time as it adds the Energon path: the hooks needed something to call.

The handoffs

ENERGON owns this — the megatron-energon package NEMO-RL owns this — code in this repo 1 Read the shards, decode one sample Energon walks the webdataset and hands over raw conversations 2 Tokenize and mask that one sample generic_sft.py:271 — the task encoder 3 Hold 4 encoded samples in the packing buffer packing_buffer_size=4 — an Energon argument 4 Decide which samples share a bin packing.py:55 — packer.pack(lengths) 5 Call back once per chosen group Energon drives the second hook, one bin at a time 6 Build the packed row packing.py:89 — tokens, labels, cu_seqlens, media 7 Batch the rows, and checkpoint the buffer Energon saves buffered samples and read positions, not the Python RNG Thick green border = the two steps people mean when they say "packing". Both are NeMo-RL code, called by Energon.

The same loop, with numbers

Four samples of 7, 5, 3 and 9 tokens, a bin that holds 18, and the greedy knapsack packer. Read the left column as "who is running right now".

ENERGON packing_buffer_size = 4 ← set under data.energon.* NEMO-RL algorithm = greedy_knapsack ← set under policy.sequence_packing.* NEMO-RL bin_capacity = 18 ← max_input_seq_length, passed to the packer ENERGON buffer fills → lengths [7, 5, 3, 9] NEMO-RL packer.pack(...) → bin A = [9, 7] filling 16 of 18 bin B = [5, 3] filling 8 of 18 ENERGON calls back twice, one call per bin NEMO-RL builds bin A → cu_seqlens [0, 9, 16] NEMO-RL builds bin B → cu_seqlens [0, 5, 8] ENERGON batches the two rows and checkpoints its buffer
DecisionOwnerWhere
How many samples are considered togetherEnergon packing_buffer_size, sft_dataloader.py:469
How many samples come from one shard in a rowEnergon max_samples_per_sequence, sft_dataloader.py:472
Which samples share a binNeMo-RL get_packer builds it, packing.py:55 runs it
How full a bin may getNeMo-RL max_input_seq_length becomes bin_capacity, sft_dataloader.py:376
What the packed row containsNeMo-RL packing.py:89 — new in this PR, 217 lines
What gets saved on checkpointEnergon buffered samples and read positions

Why the split is worth knowing

The back and forth is not a choice this PR made. Upstream Energon's PackingDataset takes two separate callbacks — a pre_packer and a final_packer — and NeMo-RL's two hooks are wired straight into those two slots. Energon asserts that both are overridden as soon as you set packing_buffer_size. There is no one-callback form to switch to.

The reason the two are separate is restart. When Energon rebuilds a single sample after a resume, it replays the final packer on a rebuilt group (packing_dataset.py:396-398) and never replays the pre-packer — which samples were grouped is already written into the restore key. So "choose the group" has to be a different function from "build the row": one is replayed on restore and one is not. Merging them would break restore, which is the thing Energon is there to give you.

The randomness lands on the seam. One of the packers NeMo-RL now exposes to this path is first_fit_shuffle, which calls the process-wide random.shuffle inside step 4. Step 7 — the checkpoint — belongs to Energon, and Energon saves the buffer and the read positions but not the Python RNG. So the thing that makes the choice and the thing that saves the state sit on opposite sides of the line, and a resumed run can group the same buffer differently. That is exactly the review comment on the @stateless decorator: neither side is wrong on its own, the seam is.
Two config blocks, because there are two owners. You set data.energon.packing_buffer_size to turn the loop on, and policy.sequence_packing.algorithm to choose what runs at step 4 — even though nothing about step 4 happens in the policy. The config knobs page walks through which key goes where, and which ones you can leave alone.

What is still open: nothing forces the two hooks to agree. They are configured from two different config blocks and built in two different places — the algorithm from one call to get_packer, the row layout from a function that takes its own padding arguments. One object holding the whole packing policy, with the two hooks as methods on it, would keep the alternation Energon requires while leaving one place where the rules live. That is the improvement available here — not fewer handoffs, fewer sources of truth.

The only way to remove the alternation altogether is to not use the callback protocol: prepack to disk before the run, or pack inside the trainer, which is what main does today. Both trade something real away — a prepacked row cannot react to the sampler, and a trainer-side pack cannot reach across a microbatch boundary. The all packing paths page lays those three side by side.