fix pp_seg_method and unfiy training attention with attn_impl (#2572)

Author: cheng221

examples/alignment/dpo/dpo_argument.py

@@ -161,3 +161,6 @@ class DPOModelArgument:
     lora_alpha: int = field(default=-1, metadata={"help": "lora_alpha"})
     rslora_plus: bool = field(default=False, metadata={"help": "Strengthen lora performance"})
     use_quick_lora: bool = field(default=True, metadata={"help": "quick lora"})
+
+    # Attention
+    attn_impl: str = field(default="flashmask", metadata={"help": "Attention implementation"})

examples/alignment/dpo/run_dpo.py

@@ -27,6 +27,7 @@
 )
 
 from paddleformers.datasets.dpo import collate_fn, create_dataset
+from paddleformers.nn.attention import AttentionInterface
 from paddleformers.peft import LoRAConfig, LoRAModel
 from paddleformers.trainer import PdArgumentParser, get_last_checkpoint, set_seed
 from paddleformers.transformers import (
@@ -57,6 +58,11 @@ def main():
 
     paddle.set_device(training_args.device)
     set_seed(training_args.seed)
+
+    avaible_attn_impl = AttentionInterface._global_mapping.keys()
+    if model_args.attn_impl not in avaible_attn_impl:
+        raise ValueError(f"Invalid attn_impl: {model_args.attn_impl}, available attn_impl: {avaible_attn_impl}")
+
     if dpo_config.loss_type == "orpo":
         dpo_config.reference_free = True
         dpo_config.sft_loss_ratio = 1.0
@@ -113,6 +119,8 @@ def main():
         dtype=dtype,
         download_hub=model_args.download_hub,
     )
+    model_config._attn_implementation = model_args.attn_impl
+
     LlmMetaConfig.set_llm_config(model_config, training_args)
 
     if not dpo_config.reference_free and not dpo_config.lora:
@@ -151,11 +159,8 @@ def main():
             ref_model = None
     if training_args.pipeline_parallel_degree > 1:
         model.config.dpo_config = None
-    if model_args.flash_mask and not model.config.use_flash_attention:
-        logger.warning("`flash_mask` must use with zero padding and flash attention.")
-        model.config.use_flash_attention = True
 
-    if model_args.flash_mask and not any(isinstance(model, cls) for cls in flash_mask_support_list):
+    if model_args.attn_impl == "flashmask" 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:
@@ -219,6 +224,7 @@ def main():
         "greedy_intokens": data_args.greedy_intokens,
         "packing": data_args.packing,
         "mix_strategy": data_args.mix_strategy,
+        "encode_one_turn": data_args.encode_one_turn,
     }
     if training_args.do_train and training_args.should_load_dataset:
         train_dataset = create_dataset(

examples/config/ernie4_5/sft_argument_ernie4_5_0p3b.json

@@ -34,10 +34,8 @@
     "tensor_parallel_degree": 2,
     "pipeline_parallel_degree": 2,
     "sharding": "stage2",
-    "zero_padding": true,
-    "flash_mask": true,
     "unified_checkpoint": true,
-    "use_flash_attention": true,
+    "attn_impl": "flashmask",
     "sequence_parallel": true,
     "report_to": "none",
     "convert_from_hf": true,

examples/config/ernie4_5_moe/sft_argument_ernie4_5_21b_a3b.json

@@ -34,10 +34,8 @@
     "tensor_parallel_degree": 4,
     "pipeline_parallel_degree": 2,
     "sharding": "stage2",
-    "zero_padding": true,
-    "flash_mask": true,
     "unified_checkpoint": true,
-    "use_flash_attention": true,
+    "attn_impl": "flashmask",
     "sequence_parallel": true,
     "report_to": "none",
     "convert_from_hf": true,

examples/config/gpt_oss/sft_argument_gptoss_20b.json

@@ -35,7 +35,6 @@
     "tensor_parallel_degree": 4,
     "pipeline_parallel_degree": 1,
     "sharding": "stage2",
-    "zero_padding": false,
     "unified_checkpoint": true,
     "use_flash_attention": false,
     "lora": true,

examples/config/qwen/dpo_argument_qwen2_0p5b.json

@@ -32,15 +32,13 @@
     "load_best_model_at_end": true,
     "tensor_parallel_degree": 1,
     "sharding": "stage1",
-    "use_flash_attention": false,
-    "flash_mask": false,
+    "attn_impl": "flashmask",
     "recompute": true,
     "recompute_granularity": "full",
     "benchmark": false,
     "unified_checkpoint": true,
     "autotuner_benchmark":false,
     "beta": 0.1,
     "loss_type": "sigmoid",
-    "greedy_zero_padding": false,
     "label_smoothing": 0.0
   }

examples/config/qwen/dpo_lora_argument_qwen2_0p5b.json

@@ -32,7 +32,7 @@
     "load_best_model_at_end": true,
     "tensor_parallel_degree": 1,
     "sharding": "stage1",
-    "use_flash_attention": true,
+    "attn_impl": "flashmask",
     "recompute": false,
     "recompute_granularity": "full",
     "beta": 0.1,

examples/config/qwen/lora_argument_qwen2_0p5b.json

@@ -35,10 +35,8 @@
     "pipeline_parallel_degree": 1,
     "sharding": "stage2",
     "lora": true,
-    "zero_padding": true,
-    "flash_mask": true,
     "unified_checkpoint": true,
-    "use_flash_attention": true,
+    "attn_impl": "flashmask",
     "convert_from_hf": false,
     "save_to_hf": false,
     "pissa": false,

examples/config/qwen/sft_argument_qwen2_0p5b.json

@@ -34,10 +34,8 @@
     "tensor_parallel_degree": 1,
     "pipeline_parallel_degree": 1,
     "sharding": "stage2",
-    "zero_padding": true,
-    "flash_mask": true,
     "unified_checkpoint": true,
-    "use_flash_attention": true,
+    "attn_impl": "flashmask",
     "convert_from_hf": false,
     "save_to_hf": false,
     "encode_one_turn": true

examples/run_finetune.py

@@ -22,6 +22,7 @@
 from paddleformers.datasets.data_utils import estimate_training
 from paddleformers.datasets.finetuning import collate_fn
 from paddleformers.datasets.finetuning import create_dataset as create_dataset_sft
+from paddleformers.nn.attention import AttentionInterface
 from paddleformers.peft import LoRAConfig, LoRAModel
 from paddleformers.trainer import (
     IntervalStrategy,
@@ -161,7 +162,12 @@ def main():
         model_config.fuse_attention_qkv = model_args.fuse_attention_qkv
     if model_args.fuse_attention_ffn is not None:
         model_config.fuse_attention_ffn = model_args.fuse_attention_ffn
-    model_config.pp_seg_method = training_args.pp_seg_method
+
+    avaible_attn_impl = AttentionInterface._global_mapping.keys()
+    if model_args.attn_impl not in avaible_attn_impl:
+        raise ValueError(f"Invalid attn_impl: {model_args.attn_impl}, available attn_impl: {avaible_attn_impl}")
+
+    model_config.pp_seg_method = model_args.pp_seg_method
     model_config.seq_length = training_args.max_seq_len
     model_config.max_sequence_length = training_args.max_seq_len
     model_config.num_nextn_predict_layers = model_args.num_nextn_predict_layers
@@ -185,13 +191,7 @@ def main():
     else:
         model = model_class.from_config(model_config, dtype=dtype)
 
-    if model_args.flash_mask and (not data_args.zero_padding or not model.config.use_flash_attention):
-        logger.warning("`flash_mask` must use with zero padding and flash attention.")
-        data_args.zero_padding = True
-        model.config.use_flash_attention = True
-        model.config._attn_implementation = "flashmask"
-
-    if model_args.flash_mask and not any(isinstance(model, cls) for cls in flash_mask_support_list):
+    if model_args.attn_impl == "flashmask" 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: