Included gradient checkpointing

modeling_codesage.py

@@ -156,6 +156,7 @@ class CodeSageBlock(nn.Module):
 class CodeSagePreTrainedModel(PreTrainedModel):
     config_class = CodeSageConfig
     base_model_prefix = "transformer"
+    supports_gradient_checkpointing = True
 
     def _init_weights(self, module):
         """Initialize the weights."""
@@ -183,6 +184,8 @@ class CodeSageModel(CodeSagePreTrainedModel):
         self.h = nn.ModuleList([CodeSageBlock(config) for _ in range(config.num_hidden_layers)])
         self.ln_f = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon)
 
+        self.gradient_checkpointing = False
+
         self.init_weights()
 
     def get_input_embeddings(self):
@@ -247,12 +250,22 @@ class CodeSageModel(CodeSagePreTrainedModel):
             if output_hidden_states:
                 all_hidden_states = all_hidden_states + (hidden_states,)
 
-            outputs = block(
-                hidden_states,
-                attention_mask=extended_attention_mask,
-                head_mask=head_mask[i],
-                output_attentions=output_attentions,
-            )
+            # Gradient checkpointing
+            if self.gradient_checkpointing and self.training:
+                outputs = self._gradient_checkpointing_func(
+                    block.__call__,
+                    hidden_states,
+                    extended_attention_mask,
+                    head_mask[i],
+                    output_attentions,
+                )
+            else:
+                outputs = block(
+                    hidden_states,
+                    attention_mask=extended_attention_mask,
+                    head_mask=head_mask[i],
+                    output_attentions=output_attentions,
+                )
 
             hidden_states = outputs[0]
             if output_attentions: