bugfix

PaddlePaddle · Jun 27, 2024 · 74bc7ad · 74bc7ad
1 parent 62b74d8
commit 74bc7ad
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/llm/alignment/dpo/run_dpo.py b/llm/alignment/dpo/run_dpo.py
@@ -137,7 +137,7 @@ def main():
         )
         model.config.use_flash_attention = True
 
-    if not any(isinstance(model, cls) for cls in flash_mask_support_list):
+    if model_args.flash_mask and not any(isinstance(model, cls) for cls in flash_mask_support_list):
         raise NotImplementedError(f"{model.__class__} not support flash mask.")
 
     if model_args.tokenizer_name_or_path is not None:

diff --git a/llm/run_finetune.py b/llm/run_finetune.py
@@ -173,7 +173,7 @@ def main():
         data_args.zero_padding = True
         model.config.use_flash_attention = True
 
-    if not any(isinstance(model, cls) for cls in flash_mask_support_list):
+    if model_args.flash_mask and not any(isinstance(model, cls) for cls in flash_mask_support_list):
         raise NotImplementedError(f"{model.__class__} not support flash mask.")
 
     if training_args.do_train and model_args.neftune: