Upload folder using huggingface_hub

src/hf_model.py

@@ -20,6 +20,7 @@ from .model import HexaTransformer as CoreTransformer
 
 class HexaModel(PreTrainedModel):
     config_class = HexaHFConfig
+    _supports_gradient_checkpointing = True 
     
     def __init__(self, config):
         super().__init__(config)

src/model.py

@@ -97,6 +97,8 @@ class HexaTransformer(nn.Module):
         self.pos_emb = RotaryEmbedding(config.dim_head)
         
         # Transformer Layers
+        self.gradient_checkpointing = False # Default
+
         self.layers = nn.ModuleList([])
         for _ in range(config.depth):
             self.layers.append(TransformerBlock(
@@ -136,7 +138,23 @@ class HexaTransformer(nn.Module):
         
         # Transformer Pass
         for layer in self.layers:
-            x = layer(x, mask=mask, rope_emb=rope_emb)
+            if self.training and self.gradient_checkpointing:
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs)
+                    return custom_forward
+                
+                # Checkpoint requires inputs to have requires_grad=True for at least one input. 
+                # x usually has it.
+                x = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(layer),
+                    x, 
+                    mask, 
+                    rope_emb,
+                    use_reentrant=False 
+                )
+            else:
+                x = layer(x, mask=mask, rope_emb=rope_emb)
             
         x = self.norm_final(x)