Support gradient checkpointing

config.json

@@ -87,4 +87,4 @@
   "use_flashfft": false,
   "use_interpolated_rotary_pos_emb": true,
   "vocab_size": 512
-}
+}

model.py

@@ -22,6 +22,8 @@ try:
 except ImportError:
     "could not import swap_mha_rope from positional_embeddings.py"
 
+from flashfftconv import FlashDepthWiseConv1d
+
 # dummy import to force huggingface to bundle the tokenizer
 from .tokenizer import ByteTokenizer
 
@@ -64,6 +66,7 @@ class AttentionBlock(nn.Module):
         self.inner_mha_cls.rotary_emb.register_buffer("inv_freq", self.inner_mha_cls.rotary_emb.inv_freq)
 
         self.mlp = ParallelGatedMLP(config).to(dtype=mlp_dtype)
+        self.filter_output = None
 
     def forward(self, u, inference_params=None, padding_mask=None, *args, **kwargs):
         if (
@@ -71,13 +74,12 @@ class AttentionBlock(nn.Module):
         ):  # workaround for masking bug in FA. This works because Wqkv does not have bias
             # and attention scores will be also automatically zeroed.
             u = u * padding_mask[..., None]
-        u = (
-            self.inner_mha_cls(
+        w = self.inner_mha_cls(
                 self.pre_norm(u),
                 inference_params=inference_params,
-            )
-            + u
         )
+        self.filter_output = w
+        u = w + u
         if type(padding_mask) == torch.Tensor:  # guard against bias
             u = u * padding_mask[..., None]
         u = self.mlp(self.post_norm(u)) + u
@@ -120,7 +122,7 @@ class ParallelHyenaFilter(nn.Module):
         self.data_dtype = None
 
         if self.use_flash_depthwise:
-            self.fir_fn = FlashDepthwiseConv1d(
+            self.fir_fn = FlashDepthWiseConv1d(
                 channels=3 * self.hidden_size,
                 kernel_size=self.short_filter_length,
                 padding=self.short_filter_length - 1,
@@ -287,6 +289,7 @@ class ParallelGatedConvBlock(nn.Module):
 
         self.proj_norm_fn = self.proj_norm
         self.res_mlp_norm_fn = self.res_mlp_norm
+        self.filter_output = None
 
         if self.config.get("compile", False):
             self.proj_norm_fn = torch.compile(self.proj_norm, fullgraph=True, dynamic=False, mode="reduce-overhead")
@@ -308,6 +311,8 @@ class ParallelGatedConvBlock(nn.Module):
 
         z, inference_params = self.filter(z, inference_params=inference_params, padding_mask=padding_mask)
 
+        self.filter_output = z
+
         z_in = self.out_filter_dense(z) + u
 
         if type(padding_mask) == torch.Tensor:  # guard against bias
@@ -343,13 +348,15 @@ class StripedHyena(nn.Module):
                 from flashfftconv import FlashFFTConv
             except:
                 raise ImportError
-            self.flash_fft = FlashFFTConv(2 * config.seqlen, dtype=torch.bfloat16)
+            self.flash_fft = FlashFFTConv(2 * config.max_seqlen, dtype=torch.bfloat16)
         else:
             self.flash_fft = None
 
         self.blocks = nn.ModuleList(
             get_block(config, layer_idx, flash_fft=self.flash_fft) for layer_idx in range(config.num_layers)
         )
+        self.gradient_checkpointing = False
+        self._gradient_checkpointing_func = None
 
     def forward(self, x, inference_params_dict=None, padding_mask=None):
         L = x.shape[1]
@@ -379,7 +386,11 @@ class StripedHyena(nn.Module):
             x = x * padding_mask[..., None]
 
         for _, block in enumerate(self.blocks):
-            x, _ = block(x, inference_params=None, padding_mask=padding_mask)
+            if self.gradient_checkpointing and self.training:
+                x, _ = self._gradient_checkpointing_func(block.__call__, x, None, padding_mask)
+            else:
+                x, _ = block(x, inference_params=None, padding_mask=padding_mask)
+        
         return x, None
 
     def initialize_inference_params(self):

modeling_hyena.py

@@ -2,6 +2,7 @@
 """StripedHyena custom code port for the Hugging Face Hub"""
 
 import torch
+import functools
 from torch.nn import functional as F
 from .configuration_hyena import StripedHyenaConfig
 from transformers import PreTrainedModel
@@ -50,8 +51,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