fix: remove cleaving (#13)

- Revert "feat: cleave off layers from encoder (#11)" (99b812dc5c29ac777aa7b5164d3ec7399d520e06)

modeling_bert.py

@@ -166,25 +166,6 @@ 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 cleaved_layers(self):
-        return len(self.layers) - self.last_layer_idx - 1
-
-    @cleaved_layers.setter
-    def cleaved_layers(self, n: int):
-        assert 0 <= n < len(self.layers)
-        self.last_layer_idx = len(self.layers) - n - 1
 
     @property
     def gradient_checkpointing(self):
@@ -205,7 +186,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[:self.last_layer_idx + 1]:
+            for layer in self.layers:
                 hidden_states = layer(hidden_states, mixer_kwargs=mixer_kwargs)
             if subset_mask is not None:
                 hidden_states = hidden_states[subset_mask]
@@ -216,11 +197,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[:self.last_layer_idx + 1]:
+                for layer in self.layers:
                     hidden_states = layer(hidden_states, mixer_kwargs=mixer_kwargs)
                 hidden_states = pad_input(hidden_states, indices, batch, seqlen)
             else:
-                for layer in self.layers[:self.last_layer_idx]:
+                for layer in self.layers[:-1]:
                     hidden_states = layer(hidden_states, mixer_kwargs=mixer_kwargs)
                 if key_padding_mask is not None:
                     subset_idx = torch.nonzero(
@@ -247,7 +228,7 @@ class BertEncoder(nn.Module):
                     "cu_seqlens_k": cu_seqlens,
                     "max_seqlen_k": max_seqlen_in_batch,
                 }
-                hidden_states = self.layers[self.last_layer_idx](hidden_states_subset, mixer_kwargs=mixer_kwargs)
+                hidden_states = self.layers[-1](hidden_states_subset, mixer_kwargs=mixer_kwargs)
         return hidden_states