feat: fix the interface of attn_implementation in embedder.decode_only.*.load_model and reranker.decode_only.*.load_model

Author: lnxtree

FlagEmbedding/finetune/embedder/decoder_only/base/load_model.py

@@ -71,7 +71,7 @@ def get_model(model_args: DecoderOnlyEmbedderModelArguments, output_dir: str, re
         model = AutoModel.from_pretrained(
             model_args.model_name_or_path,
             # torch_dtype=torch.bfloat16,
-            use_flash_attention_2=True if model_args.use_flash_attn else False,
+            attn_implementation = "flash_attention_2" if model_args.use_flash_attn else None,
             token=model_args.token,
             cache_dir=model_args.cache_dir,
             from_tf=bool(".ckpt" in model_args.model_name_or_path),
@@ -152,7 +152,7 @@ def save_merged_model(model_args: DecoderOnlyEmbedderModelArguments, output_dir:
         model = AutoModel.from_pretrained(
             model_args.model_name_or_path,
             # torch_dtype=torch.bfloat16,
-            use_flash_attention_2=True if model_args.use_flash_attn else False,
+            attn_implementation = "flash_attention_2" if model_args.use_flash_attn else None,
             token=model_args.token,
             cache_dir=model_args.cache_dir,
             from_tf=bool(".ckpt" in model_args.model_name_or_path),

FlagEmbedding/finetune/embedder/decoder_only/icl/load_model.py

@@ -71,7 +71,7 @@ def get_model(model_args: DecoderOnlyEmbedderICLModelArguments, output_dir: str,
         model = AutoModel.from_pretrained(
             model_args.model_name_or_path,
             # torch_dtype=torch.bfloat16,
-            use_flash_attention_2=True if model_args.use_flash_attn else False,
+            attn_implementation = "flash_attention_2" if model_args.use_flash_attn else None,
             token=model_args.token,
             cache_dir=model_args.cache_dir,
             from_tf=bool(".ckpt" in model_args.model_name_or_path),
@@ -150,7 +150,7 @@ def save_merged_model(model_args: DecoderOnlyEmbedderICLModelArguments, output_d
         model = AutoModel.from_pretrained(
             model_args.model_name_or_path,
             # torch_dtype=torch.bfloat16,
-            use_flash_attention_2=True if model_args.use_flash_attn else False,
+            attn_implementation = "flash_attention_2" if model_args.use_flash_attn else None,
             token=model_args.token,
             cache_dir=model_args.cache_dir,
             from_tf=bool(".ckpt" in model_args.model_name_or_path),

FlagEmbedding/finetune/reranker/decoder_only/base/load_model.py

@@ -67,7 +67,7 @@ def get_model(model_args: RerankerModelArguments):
         model = AutoModelForCausalLM.from_pretrained(
             model_args.model_name_or_path,
             # torch_dtype=torch.bfloat16,
-            use_flash_attention_2=True if model_args.use_flash_attn else False,
+            attn_implementation = "flash_attention_2" if model_args.use_flash_attn else None,
             token=model_args.token,
             cache_dir=model_args.cache_dir,
             from_tf=bool(".ckpt" in model_args.model_name_or_path),
@@ -135,7 +135,7 @@ def save_merged_model(model_args: RerankerModelArguments, output_dir: str):
         model = AutoModelForCausalLM.from_pretrained(
             model_args.model_name_or_path,
             # torch_dtype=torch.bfloat16,
-            use_flash_attention_2=True if model_args.use_flash_attn else False,
+            attn_implementation = "flash_attention_2" if model_args.use_flash_attn else None,
             token=model_args.token,
             cache_dir=model_args.cache_dir,
             from_tf=bool(".ckpt" in model_args.model_name_or_path),

FlagEmbedding/finetune/reranker/decoder_only/layerwise/load_model.py

@@ -77,7 +77,7 @@ def get_model(model_args: RerankerModelArguments, only_for_one_logit):
             model_args.model_name_or_path,
             trust_remote_code=model_args.trust_remote_code,
             # torch_dtype=torch.float16 if training_args.fp16 else torch.bfloat16,
-            use_flash_attention_2=True if model_args.use_flash_attn else False,
+            attn_implementation = "flash_attention_2" if model_args.use_flash_attn else None,
             token=model_args.token,
             cache_dir=model_args.cache_dir,
             from_tf=bool(".ckpt" in model_args.model_name_or_path),
@@ -131,7 +131,7 @@ def get_model(model_args: RerankerModelArguments, only_for_one_logit):
         model = LayerWiseMiniCPMForCausalLM.from_pretrained(
             model_args.model_name_or_path,
             # torch_dtype=torch.float16 if training_args.fp16 else torch.bfloat16,
-            use_flash_attention_2=True if model_args.use_flash_attn else False,
+            attn_implementation = "flash_attention_2" if model_args.use_flash_attn else None,
             token=model_args.token,
             cache_dir=model_args.cache_dir,
             from_tf=bool(".ckpt" in model_args.model_name_or_path),