feat: updated activation checkpointing (#14)

- wrap every layer in a checkpoint (e0da4c55e7a599407614621df650326c11cafd2f)

modeling_bert.py

@@ -81,7 +81,8 @@ def create_mixer_cls(config, cross_attn=False, return_residual=False):
         return_residual=return_residual,
         use_alibi=True,
         window_size=window_size,
-        qk_norm=use_qk_norm
+        qk_norm=use_qk_norm,
+        checkpointing=False,
     )
     return mixer_cls
 
@@ -174,8 +175,6 @@ class BertEncoder(nn.Module):
     @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.
@@ -187,7 +186,15 @@ class BertEncoder(nn.Module):
                 {"key_padding_mask": key_padding_mask.bool()} if key_padding_mask is not None else None
             )
             for layer in self.layers:
-                hidden_states = layer(hidden_states, mixer_kwargs=mixer_kwargs)
+                if self._grad_checkpointing:
+                    hidden_states = torch.utils.checkpoint.checkpoint(
+                        layer,
+                        hidden_states,
+                        use_reentrant=False,
+                        mixer_kwargs=mixer_kwargs
+                    )
+                else:
+                    hidden_states = layer(hidden_states, mixer_kwargs=mixer_kwargs)
             if subset_mask is not None:
                 hidden_states = hidden_states[subset_mask]
         else:
@@ -198,11 +205,27 @@ class BertEncoder(nn.Module):
             mixer_kwargs = {"cu_seqlens": cu_seqlens, "max_seqlen": max_seqlen_in_batch}
             if subset_mask is None:
                 for layer in self.layers:
-                    hidden_states = layer(hidden_states, mixer_kwargs=mixer_kwargs)
+                    if self._grad_checkpointing:
+                        hidden_states = torch.utils.checkpoint.checkpoint(
+                            layer,
+                            hidden_states,
+                            use_reentrant=False,
+                            mixer_kwargs=mixer_kwargs
+                        )
+                    else:
+                        hidden_states = layer(hidden_states, mixer_kwargs=mixer_kwargs)
                 hidden_states = pad_input(hidden_states, indices, batch, seqlen)
             else:
                 for layer in self.layers[:-1]:
-                    hidden_states = layer(hidden_states, mixer_kwargs=mixer_kwargs)
+                    if self._grad_checkpointing:
+                        hidden_states = torch.utils.checkpoint.checkpoint(
+                            layer,
+                            hidden_states,
+                            use_reentrant=False,
+                            mixer_kwargs=mixer_kwargs
+                        )
+                    else:
+                        hidden_states = layer(hidden_states, mixer_kwargs=mixer_kwargs)
                 if key_padding_mask is not None:
                     subset_idx = torch.nonzero(
                         subset_mask[key_padding_mask], as_tuple=False
@@ -228,7 +251,15 @@ class BertEncoder(nn.Module):
                     "cu_seqlens_k": cu_seqlens,
                     "max_seqlen_k": max_seqlen_in_batch,
                 }
-                hidden_states = self.layers[-1](hidden_states_subset, mixer_kwargs=mixer_kwargs)
+                if self._grad_checkpointing:
+                    torch.utils.checkpoint.checkpoint(
+                        self.layers[-1],
+                        hidden_states_subset,
+                        use_reentrant=False,
+                        mixer_kwargs=mixer_kwargs
+                    )
+                else:
+                    hidden_states = self.layers[-1](hidden_states_subset, mixer_kwargs=mixer_kwargs)
         return hidden_states