improve checks for various falcon model sizes

axolotl-ai-cloud · Sep 13, 2023 · 266ce55 · 266ce55
1 parent 7070825
commit 266ce55
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py
@@ -114,7 +114,11 @@ def load_model(
         if cfg.flash_attention:
             replace_btlm_attn_with_flash_attn()
 
-    if hasattr(model_config, "model_type") and model_config.model_type == "falcon":
+    if hasattr(model_config, "model_type") and model_config.model_type in [
+        "falcon",
+        "RefinedWebModel",
+        "RefinedWeb",
+    ]:
         if cfg.flash_attention:
             replace_falcon_attn_with_flash_attn()