feat: added back option not to use flash attention

configuration_bert.py

@@ -81,6 +81,7 @@ class JinaBertConfig(PretrainedConfig):
         fused_dropout_add_ln=False,
         fused_bias_fc=False,
         pad_vocab_size_multiple=1,
+        use_flash_attn=True,
         **kwargs,
     ):
         assert 'position_embedding_type' not in kwargs
@@ -106,3 +107,4 @@ class JinaBertConfig(PretrainedConfig):
         self.fused_dropout_add_ln = fused_dropout_add_ln
         self.fused_bias_fc = fused_bias_fc
         self.pad_vocab_size_multiple = pad_vocab_size_multiple
+        self.use_flash_attn = use_flash_attn

modeling_bert.py

@@ -59,6 +59,7 @@ logger = logging.getLogger(__name__)
 
 
 def create_mixer_cls(config, cross_attn=False, return_residual=False):
+    use_flash_attn = getattr(config, "use_flash_attn", False)
     fused_bias_fc = getattr(config, "fused_bias_fc", False)
     window_size = getattr(config, "window_size", (-1, -1))
     mixer_cls = partial(
@@ -68,7 +69,7 @@ def create_mixer_cls(config, cross_attn=False, return_residual=False):
         dropout=config.attention_probs_dropout_prob,
         causal=False,
         fused_bias_fc=fused_bias_fc,
-        use_flash_attn=True,
+        use_flash_attn=use_flash_attn,
         return_residual=return_residual,
         use_alibi=True,
         window_size=window_size,
@@ -151,6 +152,7 @@ def _init_weights(module, initializer_range=0.02):
 class BertEncoder(nn.Module):
     def __init__(self, config: JinaBertConfig):
         super().__init__()
+        self.use_flash_attn = getattr(config, "use_flash_attn", False)
         self.layers = nn.ModuleList(
             [create_block(config, layer_idx=i) for i in range(config.num_hidden_layers)]
         )
@@ -171,7 +173,7 @@ class BertEncoder(nn.Module):
         This means that we only compute the last layer output for these tokens.
         subset_mask: (batch, seqlen), dtype=torch.bool
         """
-        if key_padding_mask is None:
+        if key_padding_mask is None or not self.use_flash_attn:
             mixer_kwargs = (
                 {"key_padding_mask": key_padding_mask} if key_padding_mask is not None else None
             )