Update hf_decoder_model.py

add a6000 support for flash attention
OptimalScale · Nov 6, 2023 · e888923 · e888923
1 parent f846912
commit e888923
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/lmflow/models/hf_decoder_model.py b/src/lmflow/models/hf_decoder_model.py
@@ -71,15 +71,17 @@
 
 GPU_SUPPORT_FLASH_ATTENTION = {
     "A100": ["LlamaForCausalLM", "GPTNeoForCausalLM", "GPT2ForCausalLM", "BloomForCausalLM"],
-    "A40": ["GPTNeoForCausalLM", "GPT2ForCausalLM", "BloomForCausalLM"]
+    "A40": ["GPTNeoForCausalLM", "GPT2ForCausalLM", "BloomForCausalLM"],
+    "A6000": ["LlamaForCausalLM", "GPTNeoForCausalLM", "GPT2ForCausalLM", "BloomForCausalLM"]
 }
 
 try:
     import flash_attn
     if int(flash_attn.__version__.split(".")[0]) == 2:
         GPU_SUPPORT_FLASH_ATTENTION = {
             "A100": ["LlamaForCausalLM", "GPTNeoForCausalLM", "GPT2ForCausalLM", "BloomForCausalLM"],
-            "A40": ["LlamaForCausalLM","GPTNeoForCausalLM", "GPT2ForCausalLM", "BloomForCausalLM"]
+            "A40": ["LlamaForCausalLM","GPTNeoForCausalLM", "GPT2ForCausalLM", "BloomForCausalLM"],
+            "A6000": ["LlamaForCausalLM", "GPTNeoForCausalLM", "GPT2ForCausalLM", "BloomForCausalLM"]
         }
 except:
     pass