feat(experience): pack microbatches at collate time via AutoModel's THD helper - #80
Draft
HuiyingLi wants to merge 3 commits into
Draft
feat(experience): pack microbatches at collate time via AutoModel's THD helper#80HuiyingLi wants to merge 3 commits into
HuiyingLi wants to merge 3 commits into
Conversation
…HD helper The replay buffer already stores unpadded per-sample records, so the padded batch make_experience_batch rebuilds exists only to be un-padded again by pack_padded_batch inside the forward and re-padded by unpack_to_padded after it. make_packed_experience_batch skips that round trip: it packs the loose records once into the flat [1, total_tokens] THD layout the model consumes. Token packing delegates to AutoModel's pack_features_for_thd + packed_sequence_thd_collater instead of re-deriving the schema here, so the packed batch is identical to the one AutoModel's own THD/CP pipeline builds. Action-side [T-1] fields get one trailing masked slot so every per-token tensor rides the same [T] axis, which is what makes a token-mean loss layout-invariant (asserted in the tests). Not wired into collate_fn yet: R3 routed_experts ([layers, topk, T]) and VLM media are rejected, and the EP token-count equalization is a distributed step that stays in the forward. The padded path is untouched. Tests: parity with pack_padded_batch on ids / positions / cu_seqlens / max_seqlen and on every scored target position. Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
Replaces the standalone make_packed_experience_batch with a packed= flag on make_experience_batch, so the metadata/episode-field handling is shared instead of duplicated and there is no second entry point. Packing now covers routed_experts too: it is stored sequence-last ([layers, topk, T]), so it concatenates on that axis like the token fields rather than being rejected. position_ids and packed_seq_lens become Experience fields, which is what lets the forward consume a pack without re-deriving the layout. Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
… packer The forward no longer packs: a packed Experience arrives with its flat [1, total] tokens, per-sequence position_ids and packed_seq_lens, so _forward_backbone only describes the boundaries to the attention kernel and adds the EP-equalized suffix, whose length is a collective and therefore cannot be known at collate time. Deleted: - pack_padded_batch (74 lines) -> packed_attn_kwargs (~35), which no longer removes padding, builds positions, or rolls targets; the collate does. - unpack_to_padded entirely. A packed batch needs no inverse: per-token outputs keep the flat axis their inputs already share, so _restore_full_sequence only handles the cp>1 gather, and the routing targets no longer need the packer's to reorder. - _forward_backbone returns a 5-tuple instead of threading out. The per-sample shift also disappears: the collate gives each sequence a trailing masked slot, so log_probs/entropy/values stay elementwise-aligned with action_mask and the [:, :-1] slice applies to padded batches only. Verified on one GPU with a tiny Qwen3: padded vs packed action_log_probs agree to 0.0 (logs/actor_parity.py), with a control showing the comparison fails by 1.13 when position_ids are dropped. 218 CPU tests pass. Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Packs each microbatch once, at collate time, instead of padding it and then
un-padding it again inside the model forward.
The replay buffer already stores unpadded per-sample records. Today
make_experience_batchre-pads them to[B, T],pack_padded_batchstrips thatpadding back out inside
_forward_backbone, andunpack_to_paddedre-pads theoutputs before the loss.
make_experience_batch(packed=True)produces the flat[1, total_tokens]THD batch directly, and the forward consumes it.The token schema comes from AutoModel's
pack_features_for_thd+packed_sequence_thd_collaterrather than being re-derived here, so the pack isidentical to the one AutoModel's own THD/CP pipeline builds. That matters more
than it looks: the load-bearing field is
position_idsrestarting per sequence,which is what the flash-attention backend uses to infer sequence boundaries.
Net effect
12 files changed, 124 insertions(+), 164 deletions(-)pack_padded_batch(74 lines) →packed_attn_kwargs(~35): it no longerremoves padding, builds positions, or rolls targets — the collate does.
unpack_to_paddeddeleted outright. A packed batch needs no inverse: per-tokenoutputs keep the flat axis their inputs already share, so
_restore_full_sequencehandles only the cp>1 gather.packed order, and
_forward_backbonereturns a 5-tuple instead of threading it out.collate time, so
log_probs/entropy/valuesstay elementwise-alignedwith
action_mask, and the[:, :-1]slice now applies to padded batches only.What did not disappear is the EP token-count equalization: its length is a
collective, so it stays a forward-time step that extends an existing pack.
Validation
logs/actor_parity.py): the Actor'saction_log_probsfor padded vs packed input agree to 0.0 on every scoredposition. A control (
logs/parity_control.py) shows the same comparison failsby 1.13 when
position_idsare dropped, i.e. the test can fail.compileallclean.here), and VLM (
packing_samplesis already force-disabled for VLM).Dependency
Blocked on NVIDIA-NeMo/Automodel#3514, which adds
pack_features_for_thd. CI stays red until that merges andrequirements.txtbumps the
nemo-automodelpin.