Fix Gemma3N audio notebooks to use reentrant checkpointing#215
Open
danielhanchen wants to merge 1 commit intomainfrom
Open
Fix Gemma3N audio notebooks to use reentrant checkpointing#215danielhanchen wants to merge 1 commit intomainfrom
danielhanchen wants to merge 1 commit intomainfrom
Conversation
Contributor
|
Note Gemini is unable to generate a review for this pull request due to the file types involved not being currently supported. |
Non-reentrant gradient checkpointing (use_reentrant=False) causes
AOT autograd stride assertion failures during backward pass with
Gemma3N audio conformer due to variable-length audio tensors:
AssertionError: expected size 2==2, stride 1928==1936 at dim=0
Switch to use_reentrant=True which avoids AOT autograd compilation
of the backward pass entirely.
Companion to unslothai/unsloth#4629 which adds a server-side guard.
danielhanchen
force-pushed
the
dh/fix-gemma3n-audio-reentrant-checkpointing
branch
from
c72a24c to
9774858
Compare
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.
Summary
Switch
use_reentrant=Falsetouse_reentrant=Truein Gemma3N Audio notebooks (both Colab and Kaggle versions).Non-reentrant gradient checkpointing triggers AOT autograd compilation of the backward pass, which fails on Gemma3N audio conformer because variable-length audio tensors cause stride mismatches:
The audio conformer's conv/norm layers produce tensors whose strides vary with audio clip duration (e.g. 241 vs 242 mel frames), and the compiled backward graph expects fixed strides.
Changes
nb/Gemma3N_(4B)-Audio.ipynb:use_reentrant=False->use_reentrant=Truenb/Kaggle-Gemma3N_(4B)-Audio.ipynb:use_reentrant=False->use_reentrant=TrueCompanion to unslothai/unsloth#4629 which also adds a server-side guard in
vision.pyto forceuse_reentrant=Truefor Gemma3N regardless of the notebook setting.Test plan
use_reentrant=True: trains successfully (loss 12.9 -> 0.56 over 100 steps on librispeech)