Support gradient checkpointing

modeling_hyena.py

@@ -50,8 +50,32 @@ class StripedHyenaModelForCausalLM(StripedHyenaPreTrainedModel):
     def force_dtype(self):
         self.backbone.to_bfloat16_except_poles_residues() 
         
+    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
+        if not self.supports_gradient_checkpointing:
+            raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.")
+    
+        if gradient_checkpointing_kwargs is None:
+            gradient_checkpointing_kwargs = {"use_reentrant": True}
+    
+        # TODO support deepspeed checkpoint
+        gradient_checkpointing_func = functools.partial(
+            torch.utils.checkpoint.checkpoint, **gradient_checkpointing_kwargs
+        )
+    
+        self._set_gradient_checkpointing(
+            enable=True, gradient_checkpointing_func=gradient_checkpointing_func
+        )
+    
+        if getattr(self, "_hf_peft_config_loaded", False):
+            # When using PEFT + gradient checkpointing + Trainer we need to make sure the input has requires_grad=True
+            # we do it also on PEFT: https://github.com/huggingface/peft/blob/85013987aa82aa1af3da1236b6902556ce3e483e/src/peft/peft_model.py#L334
+            # When training with PEFT, only LoRA layers will have requires grad set to True, but the output of frozen layers need to propagate
+            # the gradients to make sure the gradient flows.
+            self.enable_input_require_grads()
+    
     def _set_gradient_checkpointing(self, enable, gradient_checkpointing_func):
         self.backbone.gradient_checkpointing = enable
+        self.backbone._gradient_checkpointing_func = gradient_checkpointing_func
 
     def get_input_embeddings(self):
         return self.backbone.embedding_layer