fix packing for eager/sdpa attn

2025-12-16 11:50:35 +08:00 · 2024-07-04 01:52:43 +08:00
parent 87d9b2d005
commit 6fd6aa4530
9 changed files with 51 additions and 20 deletions
--- a/src/llamafactory/extras/constants.py
+++ b/src/llamafactory/extras/constants.py
@@ -79,6 +79,7 @@ TRAINING_STAGES = {
 STAGES_USE_PAIR_DATA = {"rm", "dpo"}

 SUPPORTED_CLASS_FOR_BLOCK_DIAG_ATTN = {
+    "cohere",
    "falcon",
    "gemma",
    "gemma2",