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

unsloth checkpointing fix for latest transformers==4.52.x (#2674)

Browse source
This commit is contained in:
DoubleMathew 2025-06-03 02:06:06 -05:00 • committed by GitHub
commit c33bb76972
1 changed files with 11 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

12
unsloth/models/llama.py
View file

@ -27,6 +27,11 @@ from unsloth_zoo.peft_utils import SKIP_QUANTIZATION_MODULES
transformers_version = Version(transformers_version)
# Transformers moved rotary embeddings out of all attention layers
IS_ATTENTION_REFACTOR = transformers_version > Version("4.47.1")
try:
from transformers.modeling_layers import GradientCheckpointingLayer
except:
GradientCheckpointingLayer = type(None)
from transformers.models.llama.modeling_llama import (
logger,
BaseModelOutputWithPast,
@ -877,7 +882,12 @@ def LlamaModel_fast_forward(
mask = self. GA_mask if use_static_mask else dynamic_GA_mask
pass
if gradient_checkpointing:
try:
is_gradient_checkpointing_layer = isinstance(decoder_layer, GradientCheckpointingLayer)
except:
is_gradient_checkpointing_layer = False
if gradient_checkpointing and not is_gradient_checkpointing_layer:
def create_custom_forward(module):
def custom_forward(*inputs):
return module(*inputs, past_key_value, output_attentions, padding_mask = padding_mask, position_embeddings = position_embeddings)