Merge pull request #670 from yaoguany/flash_attn2

Update hf_decoder_model.py
OptimalScale · Nov 6, 2023 · f49123d · f49123d
2 parents f846912 + e888923
commit f49123d
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