feat: added option for QK normalization

configuration_bert.py

@@ -83,6 +83,7 @@ class JinaBertConfig(PretrainedConfig):
         pad_vocab_size_multiple=1,
         num_tasks=0,
         use_flash_attn=True,
+        use_qk_norm=True,
         **kwargs,
     ):
         assert 'position_embedding_type' not in kwargs
@@ -110,3 +111,4 @@ class JinaBertConfig(PretrainedConfig):
         self.pad_vocab_size_multiple = pad_vocab_size_multiple
         self.num_tasks = num_tasks
         self.use_flash_attn = use_flash_attn
+        self.use_qk_norm = use_qk_norm

modeling_bert.py

@@ -59,9 +59,10 @@ 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))
+    use_flash_attn = config.use_flash_attn
+    use_qk_norm = config.use_qk_norm
+    fused_bias_fc = config.fused_bias_fc
+    window_size = config.window_size
     mixer_cls = partial(
         MHA,
         num_heads=config.num_attention_heads,
@@ -73,6 +74,7 @@ def create_mixer_cls(config, cross_attn=False, return_residual=False):
         return_residual=return_residual,
         use_alibi=True,
         window_size=window_size,
+        qk_norm=use_qk_norm
     )
     return mixer_cls