update packing with sdpa and eager attention mode

2024-06-16 02:25:47 +07:00
parent 04315c3d92
commit 238f5c3d99
3 changed files with 148 additions and 73 deletions
--- a/src/llamafactory/extras/constants.py
+++ b/src/llamafactory/extras/constants.py
@@ -66,6 +66,21 @@ STAGES_USE_PAIR_DATA = {"rm", "dpo"}

 SUPPORTED_CLASS_FOR_S2ATTN = {"llama"}

+SUPPORTED_CLASS_FOR_MULTIPACK = [
+    "llama",
+    "mistral",
+    "mixtral",
+    "qwen2",
+    "qwen2_moe",
+    "falcon",
+    "phi",
+    "phi3",
+    "gemma",
+    "gemmoe",
+    "starcoder2",
+    "jamba"
+]
+
 V_HEAD_WEIGHTS_NAME = "value_head.bin"

 V_HEAD_SAFE_WEIGHTS_NAME = "value_head.safetensors"