From c33bb769721f502d9dd42c7430d7eb130c16c584 Mon Sep 17 00:00:00 2001 From: DoubleMathew Date: Tue, 3 Jun 2025 02:06:06 -0500 Subject: [PATCH] unsloth checkpointing fix for latest transformers==4.52.x (#2674) --- unsloth/models/llama.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 2587c5a501..53c956ce1a 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -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)