Unpack DiT block output tuple in the non-checkpointed training branch - #86
Open
jeqcho wants to merge 1 commit into
Open
Unpack DiT block output tuple in the non-checkpointed training branch#86jeqcho wants to merge 1 commit into
jeqcho wants to merge 1 commit into
Conversation
Transformer blocks return (x, kv_cache). Only the gradient-checkpointing branch unpacked the tuple, so running the training forward under torch.no_grad() handed a tuple to the next block and crashed. Unpack it in the non-checkpointed branch too.
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.
Fixes #85.
Blocks return
(x, kv_cache); only the checkpointed branch unpacked it, so running the training forward with gradients disabled handed a tuple to the next block. This unpacks it in the non-checkpointed branch as well, with an assert that the cache slot is unused on this path.Validated during a 12k-step LoRA fine-tune of DreamZero-AgiBot on bimanual YAM data (https://github.com/robocurve/dreamzero-yam, released at https://huggingface.co/robocurve/dreamzero-yam-molmoact2): the grad-disabled forward scored every 500-step milestone, and the deterministic val/loss it produced matched training loss behavior throughout. Checkpointed and grad-enabled paths are untouched.