fix general pipeline model

Author: cheng221

examples/config/ernie4_5/sft_argument_ernie4_5_0p3b.json

@@ -0,0 +1,47 @@
+{
+    "model_name_or_path": "baidu/ERNIE-4.5-0.3B-PT",
+    "train_dataset_path": "./data/train.json",
+    "train_dataset_prob": "1.0",
+    "train_dataset_type": "erniekit",
+    "eval_dataset_path": "./data/dev.json",
+    "eval_dataset_prob": "1.0",
+    "eval_dataset_type": "erniekit",
+    "packing": true,
+    "mix_strategy": "random",
+    "output_dir": "./checkpoints/ernie4_5_paddle_sft_ckpts",
+    "max_seq_len": 8192,
+    "per_device_train_batch_size": 1,
+    "gradient_accumulation_steps": 4,
+    "per_device_eval_batch_size": 8,
+    "eval_accumulation_steps":16,
+    "num_train_epochs": 1,
+    "learning_rate": 3e-05,
+    "warmup_steps": 10,
+    "logging_steps": 1,
+    "max_steps": 100,
+    "evaluation_strategy": "epoch",
+    "save_strategy": "epoch",
+    "src_length": 1024,
+    "max_length": 2048,
+    "bf16": true,
+    "fp16_opt_level": "O2",
+    "do_train": true,
+    "do_eval": true,
+    "disable_tqdm": true,
+    "load_best_model_at_end": true,
+    "eval_with_do_generation": false,
+    "metric_for_best_model": "accuracy",
+    "recompute": true,
+    "save_total_limit": 1,
+    "tensor_parallel_degree": 2,
+    "pipeline_parallel_degree": 2,
+    "sharding": "stage2",
+    "zero_padding": true,
+    "flash_mask": true,
+    "unified_checkpoint": true,
+    "use_flash_attention": true,
+    "sequence_parallel": true,
+    "report_to": "none",
+    "convert_from_hf": true,
+    "pp_seg_method": "layer:DecoderLayer|EmptyLayer"
+  }

examples/config/ernie4_5_moe/sft_argument_ernie4_5_21b_a3b.json

@@ -0,0 +1,47 @@
+{
+    "model_name_or_path": "baidu/ERNIE-4.5-21B-A3B-PT",
+    "train_dataset_path": "./data/train.json",
+    "train_dataset_prob": "1.0",
+    "train_dataset_type": "erniekit",
+    "eval_dataset_path": "./data/dev.json",
+    "eval_dataset_prob": "1.0",
+    "eval_dataset_type": "erniekit",
+    "packing": true,
+    "mix_strategy": "random",
+    "output_dir": "./checkpoints/ernie4_5_paddle_sft_ckpts",
+    "max_seq_len": 8192,
+    "per_device_train_batch_size": 1,
+    "gradient_accumulation_steps": 4,
+    "per_device_eval_batch_size": 8,
+    "eval_accumulation_steps":16,
+    "num_train_epochs": 1,
+    "learning_rate": 3e-05,
+    "warmup_steps": 10,
+    "logging_steps": 1,
+    "max_steps": 100,
+    "evaluation_strategy": "epoch",
+    "save_strategy": "epoch",
+    "src_length": 1024,
+    "max_length": 2048,
+    "bf16": true,
+    "fp16_opt_level": "O2",
+    "do_train": true,
+    "do_eval": true,
+    "disable_tqdm": true,
+    "load_best_model_at_end": true,
+    "eval_with_do_generation": false,
+    "metric_for_best_model": "accuracy",
+    "recompute": true,
+    "save_total_limit": 1,
+    "tensor_parallel_degree": 4,
+    "pipeline_parallel_degree": 2,
+    "sharding": "stage2",
+    "zero_padding": true,
+    "flash_mask": true,
+    "unified_checkpoint": true,
+    "use_flash_attention": true,
+    "sequence_parallel": true,
+    "report_to": "none",
+    "convert_from_hf": true,
+    "pp_seg_method": "layer:DecoderLayer|EmptyLayer"
+  }

examples/run_finetune.py

@@ -31,6 +31,10 @@
     DeepseekV2ForCausalLMPipe,
     DeepseekV3ForCausalLM,
     DeepseekV3ForCausalLMPipe,
+    Ernie4_5_MoeForCausalLM,
+    Ernie4_5_MoeForCausalLMPipe,
+    Ernie4_5ForCausalLM,
+    Ernie4_5ForCausalLMPipe,
     Llama3Tokenizer,
     LlamaForCausalLM,
     LlamaForCausalLMPipe,
@@ -53,6 +57,10 @@
     DeepseekV2ForCausalLMPipe,
     DeepseekV3ForCausalLM,
     DeepseekV3ForCausalLMPipe,
+    Ernie4_5ForCausalLM,
+    Ernie4_5ForCausalLMPipe,
+    Ernie4_5_MoeForCausalLM,
+    Ernie4_5_MoeForCausalLMPipe,
     LlamaForCausalLM,
     LlamaForCausalLMPipe,
     Qwen2ForCausalLM,
@@ -138,11 +146,11 @@ 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
     model_config.seq_length = data_args.max_length
+    model_config.max_sequence_length = training_args.max_seq_length
     model_config.num_nextn_predict_layers = model_args.num_nextn_predict_layers
     logger.info(f"Final model config: {model_config}")
-
     logger.info("Creating model")
 
     model_class = AutoModelForCausalLM
@@ -157,7 +165,7 @@ def main():
             model_args.model_name_or_path,
             config=model_config,
             download_hub=model_args.download_hub,
-            convert_from_hf=False,  # run paddle weights
+            convert_from_hf=training_args.convert_from_hf,  # run paddle weights
         )
     else:
         model = model_class.from_config(model_config, dtype=dtype)
@@ -166,6 +174,7 @@ def main():
         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):
         raise NotImplementedError(f"{model.__class__} not support flash mask.")

paddleformers/datasets/finetuning.py

@@ -698,4 +698,4 @@ def gen_attn_mask_startend_row_indices(batch_token_ids: List[List[int]], max_seq
     if offset < max_seq_len:
         attn_mask_startend_row_indices.extend(list(range(offset, max_seq_len)))
     # NOTE(hehuang): The dtype of attn_mask_startend_row_indices must be np.int32
-    return np.array(attn_mask_startend_row_indices, dtype=np.int32)[None, None]
+    return np.array(attn_mask_startend_row_indices, dtype=np.int32)[None, None, ..., None]  # add dimension modify

paddleformers/nn/attention/flashmask_attention.py

@@ -26,7 +26,7 @@ def flashmask_attention_forward(
     query: paddle.Tensor,
     key: paddle.Tensor,
     value: paddle.Tensor,
-    attn_mask_start_row_indices: paddle.Tensor,
+    attn_mask_startend_row_indices: paddle.Tensor,
     dropout: float = 0.0,
     sink: Optional[paddle.Tensor] = None,
     scaling: Optional[float] = None,
@@ -39,7 +39,7 @@ def flashmask_attention_forward(
             query,
             key,
             value,
-            startend_row_indices=attn_mask_start_row_indices,
+            startend_row_indices=attn_mask_startend_row_indices,
             causal=True,
         )
     else:
@@ -48,7 +48,7 @@ def flashmask_attention_forward(
             key,
             value,
             sink,
-            startend_row_indices=attn_mask_start_row_indices,
+            startend_row_indices=attn_mask_startend_row_indices,
             dropout_p=dropout,
             softmax_scale=scaling,
             causal=is_causal,

paddleformers/nn/attention/sdpa_attention.py

@@ -27,19 +27,19 @@ def sdpa_attention_forward(
     key: paddle.Tensor,
     value: paddle.Tensor,
     attention_mask: Optional[paddle.Tensor] = None,
-    attn_mask_start_row_indices=None,
+    attn_mask_startend_row_indices=None,
     dropout: float = 0.0,
     sink: Optional[paddle.Tensor] = None,
     scaling: Optional[float] = None,
     is_causal: Optional[bool] = None,
     **kwargs,
 ):
     # query: b l h d
-    if is_causal is None and attn_mask_start_row_indices is None:
+    if is_causal is None and attn_mask_startend_row_indices is None:
         is_causal = query.shape[1] > 1 and attention_mask is None and getattr(module, "is_causal", True)
-    elif attn_mask_start_row_indices is not None:
+    elif attn_mask_startend_row_indices is not None:
         is_causal = False
-        attention_mask = _gen_from_sparse_attn_mask_indices(attn_mask_start_row_indices, query.dtype)
+        attention_mask = _gen_from_sparse_attn_mask_indices(attn_mask_startend_row_indices, query.dtype)
 
     if sink is None:
         attn_output = nn.functional.scaled_dot_product_attention(

paddleformers/nn/pp_model.py

@@ -254,7 +254,7 @@ def forward(self, args):
         emb = self.embed_tokens(input_ids).astype(self.embed_tokens.weight.dtype)
         if position_ids is None and not self.config.fuse_rope:
             position_ids = (
-                paddle.range(
+                paddle.arange(
                     0,
                     input_ids.shape[1],
                     dtype="int64",
@@ -410,13 +410,13 @@ def forward(self, args):
             max_seq_len = hidden_states.shape[0] * self.config.tensor_parallel_degree
         if attention_mask is None:
             tgt_mask = None
-            attn_mask_start_row_indices = None
+            attn_mask_startend_row_indices = None
         elif attention_mask.dtype == paddle.int32:
             tgt_mask = None
-            attn_mask_start_row_indices = attention_mask[:, :, :max_seq_len]
+            attn_mask_startend_row_indices = attention_mask[:, :, :max_seq_len]
         else:
             tgt_mask = attention_mask[:, :, :max_seq_len, :max_seq_len]
-            attn_mask_start_row_indices = None
+            attn_mask_startend_row_indices = None
             assert len(tgt_mask.shape) == 4, f"Attention mask should be 4D tensor, but got {tgt_mask.shape}."
 
         position_ids_decoder = None
@@ -436,7 +436,7 @@ def forward(self, args):
                 self,
                 hidden_states,
                 attention_mask=tgt_mask,
-                attn_mask_start_row_indices=attn_mask_start_row_indices,
+                attn_mask_startend_row_indices=attn_mask_startend_row_indices,
                 position_ids=position_ids_decoder,
                 position_embeddings=tuple_position_embeddings,
                 use_reentrant=self.config.recompute_use_reentrant,
@@ -446,7 +446,7 @@ def forward(self, args):
                 self,
                 hidden_states=hidden_states,
                 attention_mask=tgt_mask,
-                attn_mask_start_row_indices=attn_mask_start_row_indices,
+                attn_mask_startend_row_indices=attn_mask_startend_row_indices,
                 position_ids=position_ids_decoder,
                 position_embeddings=tuple_position_embeddings,
             )
@@ -492,36 +492,44 @@ def forward(self, logits, labels):
 
 
 class GeneralModelForCausalLMPipe(PipelinePretrainedModel, PipelineLayer):
+    _decoder_layer_cls = None
+    _get_tensor_parallel_mappings = None
+    _init_weights = None
+    _keep_in_fp32_modules = None
     _tied_weights_keys = ["lm_head.weight"]
+    config_class = PretrainedConfig
+    transpose_weight_keys = None
 
-    def __init__(self, config: PretrainedConfig, decoder_layer, **kwargs):
+    def __init__(self, config: PretrainedConfig, **kwargs):
         # dynamic inherit DecoderLayer
-        DecoderLayerPipe = make_decoder_layer_pipe(decoder_layer)
+        if self._decoder_layer_cls is None:
+            raise ValueError("_decoder_layer_cls must be set before init.")
+        DecoderLayerPipe = make_decoder_layer_pipe(self._decoder_layer_cls)
+
         new_initializer_range = math.sqrt(0.3333 / config.hidden_size)
         logger.info(f"change initializer-range from {config.initializer_range} to {new_initializer_range}")
         config.initializer_range = new_initializer_range
 
-        if config.get("moe_group", "") == "mp":
+        moe_group = config.get("moe_group", "dummy")
+        if moe_group == "mp":
             assert config.sequence_parallel
 
-        if config.moe_group in {"mp", "model", "tp", "mpdp"}:
+        if moe_group in {"mp", "model", "tp", "mpdp"}:
             assert config.sequence_parallel
-            logger.info(f"disable FFN tensor model parallel, moe-group={config.moe_group}")
+            logger.info(f"disable FFN tensor model parallel, moe-group={moe_group}")
             config.disable_ffn_model_parallel = True
 
-        config.moe_group_origin = config.moe_group
-        config.moe_group = _parse_moe_group(config.moe_group)
+        config.moe_group_origin = moe_group
+        config.moe_group = _parse_moe_group(moe_group)
         config.moe_world_size = dist.get_world_size(config.moe_group)
         if config.moe_world_size < 0:
             config.moe_world_size = 1
         config.moe_rank = dist.get_rank(config.moe_group)
 
         self.config = config
-
         hcg = get_hcg()
         tensor_parallel_degree = max(hcg.get_model_parallel_world_size(), 1)
         tensor_parallel_rank = max(hcg.get_model_parallel_rank(), 0)
-
         config.tensor_parallel_degree = tensor_parallel_degree
         config.tensor_parallel_rank = tensor_parallel_rank
 
@@ -607,7 +615,7 @@ def __init__(self, config: PretrainedConfig, decoder_layer, **kwargs):
         )
 
     def get_loss_fn(self, config):
-        if config.dpo_config is not None:
+        if config.get("dpo_config", None) is not None:
             loss_fn = CriterionLayerPipe(config, use_infohub=True)
         else:
             loss_fn = CriterionLayerPipe(config)
@@ -633,7 +641,7 @@ def register_cls_attr(cls, config_class=None, pretrained_model_class=None):
     def _prepare_pipeline_inputs_func(cls, inputs):
         first_stage_keys = [
             "input_ids",
-            "attn_mask_start_row_indices",
+            "attn_mask_startend_row_indices",
             "position_ids",
             "nbatch_pack_offset",
         ]

paddleformers/trainer/training_args.py

@@ -1088,6 +1088,7 @@ class TrainingArguments:
         default=False,
         metadata={"help": "Save model to HuggingFace safetensors."},
     )
+    pp_seg_method: Optional[str] = field(default=None, metadata={"help": "PP Segmentation Method"})
 
     def __post_init__(self):
         world_size = paddle.distributed.get_world_size()

paddleformers/transformers/__init__.py

@@ -169,10 +169,10 @@
     ],
     "deepseek_v3.modeling_pp": ["DeepseekV3ForCausalLMPipe"],
     "ernie4_5.configuration": ["Ernie4_5Config"],
-    "ernie4_5.modeling": ["Ernie4_5Model", "Ernie4_5ForCausalLM"],
+    "ernie4_5.modeling": ["Ernie4_5Model", "Ernie4_5ForCausalLM", "Ernie4_5ForCausalLMPipe"],
     "ernie4_5.tokenizer": ["Ernie4_5Tokenizer"],
     "ernie4_5_moe.configuration": ["Ernie4_5_MoeConfig"],
-    "ernie4_5_moe.modeling": ["Ernie4_5_MoeModel", "Ernie4_5_MoeForCausalLM"],
+    "ernie4_5_moe.modeling": ["Ernie4_5_MoeModel", "Ernie4_5_MoeForCausalLM", "Ernie4_5_MoeForCausalLMPipe"],
     "export": ["export_model"],
     "llama.configuration": [
         "LLAMA_PRETRAINED_INIT_CONFIGURATION",

paddleformers/transformers/ernie4_5/__init__.py

@@ -20,7 +20,7 @@
 import_structure = {
     "tokenizer": ["Ernie4_5Tokenizer"],
     "configuration": ["Ernie4_5Config"],
-    "modeling": ["Ernie4_5DecoderLayer", "Ernie4_5Model", "Ernie4_5ForCausalLM"],
+    "modeling": ["Ernie4_5DecoderLayer", "Ernie4_5Model", "Ernie4_5ForCausalLM", "Ernie4_5ForCausalLMPipe"],
 }
 
 if TYPE_CHECKING: