feat: try to monkey-patch index_first_axis

modeling_bert.py

@@ -28,12 +28,16 @@ from transformers.models.bert.modeling_bert import (
     BaseModelOutputWithPoolingAndCrossAttentions,
     BertForPreTrainingOutput,
 )
-from .patched_padding_bert import index_first_axis
+from .patched_padding_bert import index_first_axis as index_first_axis_monkey_patch
+import flash_attn.bert_padding
+flash_attn.bert_padding.index_first_axis = index_first_axis_monkey_patch
+"""
 from flash_attn.bert_padding import (
     index_first_axis_residual,
     pad_input,
     unpad_input,
 )
+"""
 from flash_attn.modules.block import Block
 from flash_attn.modules.embedding import BertEmbeddings
 from flash_attn.modules.mha import MHA
@@ -172,14 +176,14 @@ class BertEncoder(nn.Module):
                 hidden_states = hidden_states[subset_mask]
         else:
             batch, seqlen = hidden_states.shape[:2]
-            hidden_states, indices, cu_seqlens, max_seqlen_in_batch = unpad_input(
+            hidden_states, indices, cu_seqlens, max_seqlen_in_batch = flash_attn.bert_padding.unpad_input(
                 hidden_states, key_padding_mask
             )
             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)
-                hidden_states = pad_input(hidden_states, indices, batch, seqlen)
+                hidden_states = flash_attn.bert_padding.pad_input(hidden_states, indices, batch, seqlen)
             else:
                 for layer in self.layers[:-1]:
                     hidden_states = layer(hidden_states, mixer_kwargs=mixer_kwargs)
@@ -197,7 +201,7 @@ class BertEncoder(nn.Module):
                     subset_cu_seqlens = F.pad(
                         torch.cumsum(subset_seqlens, dim=0, dtype=torch.torch.int32), (1, 0)
                     )
-                hidden_states_subset, hidden_states = index_first_axis_residual(
+                hidden_states_subset, hidden_states = flash_attn.bert_padding.index_first_axis_residual(
                     hidden_states, subset_idx
                 )
                 # It's ok to set max_seqlen_q to be much larger