feat: added functionality to cleave off layers from BERT encoder

modeling_bert.py

@@ -166,6 +166,16 @@ class BertEncoder(nn.Module):
             [create_block(config, layer_idx=i) for i in range(config.num_hidden_layers)]
         )
         self._grad_checkpointing = False
+        self._last_layer_idx = len(self.layers) - 1
+
+    @property
+    def last_layer_idx(self):
+        return self._last_layer_idx
+
+    @last_layer_idx.setter
+    def last_layer_idx(self, idx: int):
+        assert 0 <= idx < len(self.layers)
+        self._last_layer_idx = idx
 
     @property
     def gradient_checkpointing(self):
@@ -186,7 +196,7 @@ class BertEncoder(nn.Module):
             mixer_kwargs = (
                 {"key_padding_mask": key_padding_mask.bool()} if key_padding_mask is not None else None
             )
-            for layer in self.layers:
+            for layer in self.layers[:self.last_layer_idx + 1]:
                 hidden_states = layer(hidden_states, mixer_kwargs=mixer_kwargs)
             if subset_mask is not None:
                 hidden_states = hidden_states[subset_mask]
@@ -197,11 +207,11 @@ 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:
+                for layer in self.layers[:self.last_layer_idx + 1]:
                     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]:
+                for layer in self.layers[:self.last_layer_idx]:
                     hidden_states = layer(hidden_states, mixer_kwargs=mixer_kwargs)
                 if key_padding_mask is not None:
                     subset_idx = torch.nonzero(
@@ -228,7 +238,7 @@ 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)
+                hidden_states = self.layers[self.last_layer_idx](hidden_states_subset, mixer_kwargs=mixer_kwargs)
         return hidden_states