update arg name

Former-commit-id: 1509ed550b2060f946ce20e3c5a9e5c49e86e3ab
2024-07-03 23:23:24 +08:00
parent 5acaa476d6
commit 1408aa078d
3 changed files with 20 additions and 34 deletions
--- a/src/llamafactory/extras/constants.py
+++ b/src/llamafactory/extras/constants.py
@@ -78,22 +78,22 @@ TRAINING_STAGES = {

 STAGES_USE_PAIR_DATA = {"rm", "dpo"}

-SUPPORTED_CLASS_FOR_S2ATTN = {"llama"}
-
-SUPPORTED_CLASS_EFFECIENT_PACKING = [
+SUPPORTED_CLASS_FOR_BLOCK_DIAG_ATTN = {
+    "falcon",
+    "gemma",
+    "gemma2",
+    "jamba",
    "llama",
    "mistral",
    "mixtral",
-    "qwen2",
-    "qwen2_moe",
-    "falcon",
    "phi",
    "phi3",
-    "gemma",
-    "gemmoe",
+    "qwen2",
+    "qwen2_moe",
    "starcoder2",
-    "jamba"
-]
+}
+
+SUPPORTED_CLASS_FOR_S2ATTN = {"llama"}

 V_HEAD_WEIGHTS_NAME = "value_head.bin"