feat: support gradient checkpointing

modeling_bert.py

@@ -154,6 +154,17 @@ class BertEncoder(nn.Module):
         self.layers = nn.ModuleList(
             [create_block(config, layer_idx=i) for i in range(config.num_hidden_layers)]
         )
+        self._grad_checkpointing = False
+
+    @property
+    def gradient_checkpointing(self):
+        return self._grad_checkpointing
+
+    @gradient_checkpointing.setter
+    def gradient_checkpointing(self, value):
+        self._grad_checkpointing = value
+        for block in self.layers:
+            block.mixer.checkpointing = value
 
     def forward(self, hidden_states, key_padding_mask=None, subset_mask=None):
         """If subset_mask is not None, we only want output for the subset of the sequence.
@@ -298,6 +309,11 @@ class BertPreTrainedModel(PreTrainedModel):
     """
     config_class = JinaBertConfig
     base_model_prefix = "bert"
+    supports_gradient_checkpointing = True
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, BertEncoder):
+            module.gradient_checkpointing = value
 
 
 class BertModel(BertPreTrainedModel):