From 5c9a22b8166977acd072f224f428bcf072819cc3 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 27 Mar 2026 02:53:21 -0700 Subject: [PATCH] Fix Gemma3N audio training stride assertion with non-reentrant checkpointing (#4629) * Fix Gemma3N audio training stride assertion with non-reentrant checkpointing Gemma3N audio conformer processes variable-length audio tensors that cause stride mismatches in AOT autograd compiled backward when non-reentrant gradient checkpointing is used. The error manifests as: AssertionError: expected size 2==2, stride 1928==1936 at dim=0 This happens because the audio conformer's conv/norm layers produce tensors whose strides vary with audio clip duration, but AOT autograd traces the backward graph assuming fixed strides from the first batch. The notebook sets gradient_checkpointing_kwargs={"use_reentrant": False} and TRL 0.27.0+ also forces this. Both override Unsloth's own use_reentrant=True set during prepare_model_for_training. Fix: intercept gradient_checkpointing_enable on Gemma3N models to always force use_reentrant=True, regardless of what the notebook or TRL passes. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/models/vision.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index a8adba99e7..f558aa3f00 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -1368,6 +1368,24 @@ class FastBaseModel: patch_modules_to_save = True, ) + # Gemma3N audio conformer processes variable-length audio tensors + # that cause stride mismatches in AOT autograd compiled backward + # when non-reentrant checkpointing is used. The notebook or TRL + # may override gradient_checkpointing_kwargs with use_reentrant=False + # after this point, so we intercept gradient_checkpointing_enable + # to always force use_reentrant=True for Gemma3N. + _model_type = getattr(getattr(model, "config", None), "model_type", "") or "" + if "gemma3n" in _model_type.lower(): + _original_gc_enable = model.gradient_checkpointing_enable + + def _gc_enable_reentrant(**kwargs): + gc_kwargs = kwargs.get("gradient_checkpointing_kwargs", {}) or {} + gc_kwargs["use_reentrant"] = True + kwargs["gradient_checkpointing_kwargs"] = gc_kwargs + return _original_gc_enable(**kwargs) + + model.gradient_checkpointing_enable = _gc_enable_reentrant + from transformers.trainer import Trainer if (