0.19 implementing flash_attn

app.py

@@ -4,6 +4,7 @@ import torch
 import gradio as gr
 import logging
 from huggingface_hub import login
+from flash_attn.flash_attention import FlashAttention
 
 import os
 import traceback
@@ -66,6 +67,10 @@ def load_model_a(model_id):
             device_map="auto",
             trust_remote_code=True,
         ).eval()
+        for name, module in model_a.named_modules():
+            if isinstance(module, torch.nn.MultiheadAttention):
+                module.forward = FlashAttention(module.embed_dim)
+                logging.debug(f'{SPACER} forwarding module of {model_id_a} to flash_attn')
     except Exception as e:
         logging.error(f'{SPACER} Error: {e}, Traceback {traceback.format_exc()}')
 
@@ -83,6 +88,10 @@ def load_model_b(model_id):
             device_map="auto",
             trust_remote_code=True,
         ).eval()
+        for name, module in model_b.named_modules():
+            if isinstance(module, torch.nn.MultiheadAttention):
+                module.forward = FlashAttention(module.embed_dim)
+                logging.debug(f'{SPACER} forwarding module of {model_id_b} to flash_attn')
     except Exception as e:
         logging.error(f'{SPACER} Error: {e}, Traceback {traceback.format_exc()}')
     return gr.update(label=model_id)

requirements.txt

@@ -5,4 +5,5 @@ accelerate==0.33.0
 sentencepiece==0.2.0
 spaces==0.29.2
 gradio==4.39.0
-bitsandbytes==0.43.2
+bitsandbytes==0.43.2
+flash-attn