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

@@ -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 (