ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #3879 from ducviet00/fix-gc

Browse source 
Disable gradient checkpointing when explicitly off for vision
This commit is contained in:
Daniel Han 2026-01-14 04:32:02 -08:00 • committed by GitHub
commit 5c8ccf0671
2 changed files with 5 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

6
unsloth/models/rl.py
View file

@ -264,17 +264,19 @@ def prepare_for_training_mode(f):
def wrapper(self, *args, **kwargs):
# Enable training mode
_was_training = None
# Get gradient checkpointing setting from training arguments
use_gc = getattr(self.args, 'gradient_checkpointing', True)
if hasattr(self, 'model') and hasattr(self.model, "training"):
_was_training = self.model.training
if hasattr(self, 'model') and hasattr(self.model, "for_training"):
self.model.for_training()
self.model.for_training(use_gradient_checkpointing=use_gc)
output = f(self, *args, **kwargs)
# Restore previous mode when possible
if hasattr(self, 'model') and hasattr(self.model, "for_inference"):
if _was_training is False:
self.model.for_inference()
elif _was_training is True and hasattr(self.model, "for_training"):
self.model.for_training()
self.model.for_training(use_gradient_checkpointing=use_gc)
# Reset gradient checkpointing buffers to free memory while staying ready for next run
try:
reset_unsloth_gradient_checkpointing_buffers()

2
unsloth/models/vision.py
View file

@ -1273,7 +1273,7 @@ class FastBaseModel:
# Since transformers 4.53, must turn on explicitly
for module in model.modules():
if hasattr(module, "gradient_checkpointing"):
module.gradient_checkpointing = True
module.gradient_checkpointing = use_gradient_checkpointing
# Also re-enable training for embeddings for NEFTune
if hasattr(model, "get_input_embeddings"):