fix: gradient checkpointing compat for transformers 4.40.2 (OOM fix)

root_gainlora/src/t5_gainlora_inflora.py

@@ -1006,9 +1006,11 @@ class T5PreTrainedModel(PreTrainedModel):
             if module.has_relative_attention_bias:
                 module.relative_attention_bias.weight.data.normal_(mean=0.0, std=factor * ((d_model) ** -0.5))
 
-    def _set_gradient_checkpointing(self, module, value=False):
-        if isinstance(module, (T5Attention, T5Stack)):
-            module.gradient_checkpointing = value
+    # NOTE: _set_gradient_checkpointing removed intentionally.
+    # The old format (with 'value' param) causes transformers to silently ignore
+    # gradient_checkpointing_kwargs (including use_reentrant=False).
+    # Without this method, transformers uses the new format which properly
+    # passes the checkpointing function with use_reentrant=False.
 
     def _shift_right(self, input_ids):
         decoder_start_token_id = self.config.decoder_start_token_id
@@ -1415,18 +1417,34 @@ class T5Stack(T5PreTrainedModel):
 
                     return custom_forward
 
-                layer_outputs = checkpoint(
-                    create_custom_forward(layer_module),
-                    hidden_states,
-                    extended_attention_mask,
-                    position_bias,
-                    encoder_hidden_states,
-                    encoder_extended_attention_mask,
-                    encoder_decoder_position_bias,
-                    layer_head_mask,
-                    cross_attn_layer_head_mask,
-                    None,  # past_key_value is always None with gradient checkpointing
-                )
+                gc_fn = getattr(self, '_gradient_checkpointing_func', None)
+                if gc_fn is not None:
+                    layer_outputs = gc_fn(
+                        create_custom_forward(layer_module),
+                        hidden_states,
+                        extended_attention_mask,
+                        position_bias,
+                        encoder_hidden_states,
+                        encoder_extended_attention_mask,
+                        encoder_decoder_position_bias,
+                        layer_head_mask,
+                        cross_attn_layer_head_mask,
+                        None,  # past_key_value is always None with gradient checkpointing
+                    )
+                else:
+                    layer_outputs = checkpoint(
+                        create_custom_forward(layer_module),
+                        hidden_states,
+                        extended_attention_mask,
+                        position_bias,
+                        encoder_hidden_states,
+                        encoder_extended_attention_mask,
+                        encoder_decoder_position_bias,
+                        layer_head_mask,
+                        cross_attn_layer_head_mask,
+                        None,
+                        use_reentrant=False,
+                    )
             else:
                 layer_outputs = layer_module(
                     hidden_states,