Update configuration_bert.py

configuration_bert.py

@@ -1,6 +1,7 @@
 # Copyright 2022 MosaicML Examples authors
 # SPDX-License-Identifier: Apache-2.0
 
+from typing import Optional
 from transformers import BertConfig as TransformersBertConfig
 
 
@@ -10,6 +11,7 @@ class BertConfig(TransformersBertConfig):
         self,
         alibi_starting_size: int = 512,
         attention_probs_dropout_prob: float = 0.0,
+        flash_attn_type: Optional[str] = None,
         **kwargs,
     ):
         """Configuration class for MosaicBert.
@@ -20,7 +22,10 @@ class BertConfig(TransformersBertConfig):
                 Defaults to 512.
             attention_probs_dropout_prob (float): By default, turn off attention dropout in Mosaic BERT
                 (otherwise, Flash Attention will be off by default). Defaults to 0.0.
+            flash_attn_type (str): if 'triton' is passed will use ./flash_attn_triton.py.
+                Defaults to None (disabled).
         """
         super().__init__(
             attention_probs_dropout_prob=attention_probs_dropout_prob, **kwargs)
         self.alibi_starting_size = alibi_starting_size
+        self.flash_attn_type = flash_attn_type