feat: choose flash attention heuristically if not set explicitly

modeling_bert.py

@@ -66,7 +66,7 @@ logger = logging.getLogger(__name__)
 
 
 def create_mixer_cls(config, cross_attn=False, return_residual=False):
-    use_flash_attn = config.use_flash_attn
+    use_flash_attn = config.use_flash_attn if config.use_flash_attn is not None else torch.cuda.is_available()
     use_qk_norm = config.use_qk_norm
     fused_bias_fc = config.fused_bias_fc
     window_size = config.window_size
@@ -161,7 +161,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.use_flash_attn = config.use_flash_attn if config.use_flash_attn is not None else torch.cuda.is_available()
         self.layers = nn.ModuleList(
             [create_block(config, layer_idx=i) for i in range(config.num_hidden_layers)]
         )