Merge pull request #4446 from stceum/bug-fix

Bug Fix: `off` is parsed as `False` in yaml file

Former-commit-id: cc452c32c7f882c7f024a9d229352918a9eaa925

src/llamafactory/hparams/model_args.py

@@ -97,7 +97,7 @@ class ModelArguments:
         default=None,
         metadata={"help": "Which scaling strategy should be adopted for the RoPE embeddings."},
     )
-    flash_attn: Literal["off", "sdpa", "fa2", "auto"] = field(
+    flash_attn: Literal["disabled", "sdpa", "fa2", "auto"] = field(
         default="auto",
         metadata={"help": "Enable FlashAttention for faster training and inference."},
     )

src/llamafactory/model/model_utils/attention.py

@@ -32,7 +32,7 @@ def configure_attn_implementation(config: "PretrainedConfig", model_args: "Model
     if model_args.flash_attn == "auto":
         return
 
-    elif model_args.flash_attn == "off":
+    elif model_args.flash_attn == "disabled":
         requested_attn_implementation = "eager"
 
     elif model_args.flash_attn == "sdpa":

tests/model/model_utils/test_attention.py

@@ -29,7 +29,7 @@ INFER_ARGS = {
 
 
 def test_attention():
-    attention_available = ["off"]
+    attention_available = ["disabled"]
     if is_torch_sdpa_available():
         attention_available.append("sdpa")
 
@@ -37,7 +37,7 @@ def test_attention():
         attention_available.append("fa2")
 
     llama_attention_classes = {
-        "off": "LlamaAttention",
+        "disabled": "LlamaAttention",
         "sdpa": "LlamaSdpaAttention",
         "fa2": "LlamaFlashAttention2",
     }