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

Complete disable gradient_checkpointing for vision when use_gradient_checkpointing=False

Browse source
This commit is contained in:
Duc-Viet Hoang 2026-01-12 10:03:54 +07:00 • committed by GitHub
commit 5b422f7a06
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"):