Fix Gemma3N/Gemma4 gradient_checkpointing_enable crash on positional kwargs

[oobabooga]

Fixes #4886.

GRPO training on a Gemma3N/Gemma4 model crashes with:

TypeError: FastBaseModel.post_patch_model.<locals>._gc_enable_reentrant() takes 0 positional arguments but 1 was given

Problem

For Gemma3N/Gemma4, post_patch_model replaces model.gradient_checkpointing_enable with a wrapper that forces use_reentrant=True (needed to avoid the audio-conformer stride assertion from #4629). The wrapper's signature was def _gc_enable_reentrant(**kwargs), keyword-only.

HF's real method is gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None), and TRL's disable_gradient_checkpointing re-enables it positionally: model.gradient_checkpointing_enable(gradient_checkpointing_kwargs). That positional call hit the keyword-only wrapper and raised the TypeError during GRPO generation.

Fix

Accept gradient_checkpointing_kwargs as a positional-or-keyword arg, matching HF's signature, and force use_reentrant=True regardless of call style. The dict(...) copy also stops the wrapper mutating the caller's config in place.

As a side effect this restores PEFT's capability probe: prepare_model_for_kbit_training checks "gradient_checkpointing_kwargs" in inspect.signature(...).parameters, which the old **kwargs-only signature returned False for, so PEFT was silently dropping the kwargs. The named param now makes it True.

Scope

On current main, the GRPO path that triggers this is usually masked: patch_trl_disable_gradient_checkpointing swaps TRL's disable_gradient_checkpointing for a no-op, so the positional re-enable never runs. The fix still matters for installs where that no-op is absent (the reporter is on an older unsloth) or where its trl.* module rebind misses the call site, for PEFT's introspection above, and for any direct positional caller, since HF's own signature accepts one.

Verification

Reproduced end-to-end on a real unsloth/gemma-3n-E2B-it load with TRL 1.0.0. With TRL's real disable_gradient_checkpointing active, one GRPO step on main crashes with the exact traceback above at _generate_and_score_completions, while the same step on this branch invokes the positional re-enable (confirmed by instrumentation) and completes.

Also exercised the wrapper directly across every caller's call style:

	no-arg	positional (TRL)	keyword (HF Trainer)	PEFT signature probe
main	ok	TypeError (#4886)	ok	False (drops kwargs)
fix	ok	ok, use_reentrant=True	ok, use_reentrant=True	True

[gemini-code-assist]

Code Review

This pull request updates the _gc_enable_reentrant function in unsloth/models/vision.py to accept gradient_checkpointing_kwargs as an explicit parameter. This ensures compatibility when gradient checkpointing is re-enabled positionally by TRL or via keyword arguments by HF/PEFT. There are no review comments, so no additional feedback is provided.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Repo admins can enable using credits for code reviews in their settings.