Skip short trailing windows so batch > 1 collate sees uniform shapes - #87
Open
jeqcho wants to merge 1 commit into
Open
Skip short trailing windows so batch > 1 collate sees uniform shapes#87jeqcho wants to merge 1 commit into
jeqcho wants to merge 1 commit into
Conversation
Windows near an episode tail can trim down to fewer chunks than the rest of the batch. With per_device_train_batch_size > 1 the collate np.stack then sees mixed shapes and crashes. Skip these windows entirely, like the existing <=8-frame case, and give the empty array an integer dtype since it is used as an index downstream.
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
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.
Fixes #80.
Root cause of the trailing-batch crash: near an episode tail,
get_frame_indicestrims the window (unique_sorted = unique_sorted[:-7]), producing a sample with fewer chunks than the rest of the batch. Withper_device_train_batch_size > 1thenp.stackindreamzero_cotrain.pythen sees mixed shapes and raisesValueError: all input arrays must have the same shape. On our corpus this fired roughly once per 580 batches, so short smoke runs pass and long runs die.Fix: skip short windows entirely (same mechanism as the existing
<= 8frames case) instead of trimming to a different chunk count. The empty array needs an integer dtype because it is used as an index downstream. The cost is discarding a small number of episode-tail anchors, which seems preferable to a crash; padding them to full length would change training semantics.Note this is distinct from the variable-chunk design of the synthetic DROID dataset discussed in #69 — this path is the sharded LeRobot loader, where every sample is supposed to be full-length.
Validated on a 12k-step fine-tune of DreamZero-AgiBot at per-device batch 4 (global 32) over 124 bimanual YAM LeRobot repos with zero collate failures: https://github.com/robocurve/dreamzero-yam, released at https://huggingface.co/robocurve/dreamzero-yam-molmoact2.