EB4.5 supports SFT dataflow

Author: lshpku

examples/experiments/ernie_pretrain/ernie/pretrain.py

@@ -18,6 +18,7 @@
 import re
 import time
 from dataclasses import dataclass
+from functools import partial
 
 import numpy as np
 import paddle
@@ -60,6 +61,10 @@
     build_train_valid_test_datasets,
     check_data_split,
 )
+from paddleformers.datasets.finetuning import collate_fn
+from paddleformers.datasets.finetuning import create_dataset as create_dataset_sft
+from paddleformers.trainer import TrainingArguments
+from paddleformers.trl import ModelConfig
 
 try:
     from paddleformers.trainer.trainer_utils import log_trainer_start
@@ -459,7 +464,47 @@ def sname_to_tname(pp_model):
 
     logger.info(f"using model={type(model)}, cfg={cfg}")
 
-    train_dataset, eval_dataset, test_dataset, data_collator = create_pretrained_dataset(args)
+    dataset_config = {
+        "tokenizer": tokenizer,
+        "max_seq_len": args.max_seq_length + 1,
+        "random_seed": args.seed,
+        "num_replicas": args.dataset_world_size,
+        "rank": args.dataset_rank,
+        "num_samples_each_epoch": trainer_args.get("num_samples_each_epoch", 6000000),
+        "random_shuffle": True,
+        "greedy_intokens": True,
+        "packing": True,
+        "mix_strategy": "concat",
+        "encode_one_turn": True,
+        "use_template": True,
+        "is_pretraining": False,
+    }
+
+    if trainer_args.get("stage") == "sft":
+        train_dataset = create_dataset_sft(
+            task_group=trainer_args["train_dataset_path"],
+            task_group_prob=trainer_args.get("train_dataset_prob", 1.0),
+            sub_dataset_type=trainer_args.get("train_dataset_type", "erniekit"),
+            **dataset_config,
+        )
+        eval_dataset = create_dataset_sft(
+            task_group=trainer_args["eval_dataset_path"],
+            task_group_prob=trainer_args.get("eval_dataset_prob", 1.0),
+            sub_dataset_type=trainer_args.get("eval_dataset_type", "erniekit"),
+            is_valid=True,
+            **dataset_config,
+        )
+        data_collator = partial(
+            collate_fn,
+            tokenizer=tokenizer,
+            training_args=TrainingArguments(
+                output_dir=args.output_dir, num_nextn_predict_layers=args.multi_token_pred_depth
+            ),
+            model_args=ModelConfig(stage="SFT", use_attn_mask_startend_row_indices=True),
+            max_seq_len=args.max_seq_length + 1,
+        )
+    else:
+        train_dataset, eval_dataset, _, data_collator = create_pretrained_dataset(args)
 
     callbacks = []
     callbacks += [GlobalRNGCallback()]

examples/experiments/ernie_pretrain/models/ernie/modeling.py

@@ -357,7 +357,7 @@ def scaled_dot_product_attention(
             value_states.astype(value_states.dtype),
             startend_row_indices=startend_row_indices,
             dropout=config.attention_probs_dropout_prob,
-            causal=False,
+            causal=True,
         )
         attn_output = attn_output.reshape([bsz, q_len, head_dim * num_heads])
         return attn_output, None
@@ -1104,6 +1104,7 @@ def rope_attn(
         past_key_value=None,
         use_cache=False,
         inbatch_pack_offset=None,
+        attn_mask_startend_row_indices=None,
     ):
         if mix_layer is not None:
             query_states, key_states, value_states = paddle.split(mix_layer, 3, axis=-1)
@@ -1186,6 +1187,7 @@ def rope_attn(
             config=self.config,
             inbatch_pack_offset=inbatch_pack_offset,
             training=self.training,
+            startend_row_indices=attn_mask_startend_row_indices,
         )
         return attn_output, attn_weights, past_key_value
 

examples/experiments/ernie_pretrain/models/ernie/modeling_moe.py

@@ -821,6 +821,7 @@ def forward(
         use_cache: bool = False,
         inbatch_pack_offset: Optional[Tuple[paddle.Tensor]] = None,
         token_type_ids: Optional[Tuple[paddle.Tensor]] = None,
+        attn_mask_startend_row_indices: Optional[paddle.Tensor] = None,
     ) -> Tuple[paddle.Tensor, Optional[paddle.Tensor], Optional[Tuple[paddle.Tensor]]]:
         if token_type_ids is not None:
             token_type_ids = token_type_ids[:, :-1]
@@ -901,6 +902,7 @@ def forward(
                 past_key_value,
                 use_cache,
                 inbatch_pack_offset,
+                attn_mask_startend_row_indices=attn_mask_startend_row_indices,
                 use_reentrant=False,
             )
         else:
@@ -915,6 +917,7 @@ def forward(
                 past_key_value=past_key_value,
                 use_cache=use_cache,
                 inbatch_pack_offset=inbatch_pack_offset,
+                attn_mask_startend_row_indices=attn_mask_startend_row_indices,
             )
         if self.config.sequence_parallel:
             attn_output = attn_output.reshape([-1, attn_output.shape[-1]])
@@ -1152,6 +1155,7 @@ def forward(
         use_cache: Optional[bool] = False,
         inbatch_pack_offset: Optional[paddle.Tensor] = None,
         output_gate_logits=True,
+        attn_mask_startend_row_indices: Optional[paddle.Tensor] = None,
     ) -> Tuple[paddle.Tensor, Optional[Tuple[paddle.Tensor, paddle.Tensor]]]:
         residual = hidden_states
         if token_type_ids is not None:
@@ -1178,6 +1182,7 @@ def forward(
             use_cache=use_cache,
             inbatch_pack_offset=inbatch_pack_offset,
             token_type_ids=token_type_ids,
+            attn_mask_startend_row_indices=attn_mask_startend_row_indices,
         )
 
         if self.use_linear_residual_norm_recompute is True:
@@ -1660,6 +1665,7 @@ def forward(
         output_hidden_states=None,
         return_dict=False,
         inbatch_pack_offset=None,
+        attn_mask_startend_row_indices=None,
         **kwargs,
     ):
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
@@ -1719,6 +1725,12 @@ def forward(
             )
         hidden_states = inputs_embeds
 
+        attn_mask_startend_row_indices_ori = attn_mask_startend_row_indices
+        if attn_mask_startend_row_indices is not None:
+            attn_mask_startend_row_indices = attn_mask_startend_row_indices[
+                :, :, : -self.config.multi_token_pred_depth
+            ]
+
         all_hidden_states = () if output_hidden_states else None
         all_self_attns = () if output_attentions else None
         next_decoder_cache = () if use_cache else None
@@ -1743,6 +1755,7 @@ def forward(
                     past_key_value,
                     use_cache,
                     inbatch_pack_offset,
+                    attn_mask_startend_row_indices=attn_mask_startend_row_indices,
                 )
             else:
                 layer_outputs = decoder_layer(
@@ -1754,6 +1767,7 @@ def forward(
                     past_key_value,
                     use_cache,
                     inbatch_pack_offset,
+                    attn_mask_startend_row_indices=attn_mask_startend_row_indices,
                 )
 
             if isinstance(layer_outputs, (tuple, list)):
@@ -1786,6 +1800,11 @@ def forward(
                     ],
                     axis=1,
                 )
+                attn_mask_startend_row_indices_cur_depth = None
+                if attn_mask_startend_row_indices is not None:
+                    attn_mask_startend_row_indices_cur_depth = attn_mask_startend_row_indices_ori[
+                        :, :, (depth + 1) : inputs_embeds_ori.shape[1] + (depth + 1)
+                    ] - (depth + 1)
 
                 inputs_embeds_cur_depth_norm = self.mtp_emb_norm[depth](inputs_embeds_cur_depth)
                 hidden_states_norm = self.mtp_hidden_norm[depth](hidden_states)
@@ -1809,6 +1828,7 @@ def forward(
                     past_key_value,
                     use_cache,
                     inbatch_pack_offset,
+                    attn_mask_startend_row_indices=attn_mask_startend_row_indices_cur_depth,
                 )
 
                 if isinstance(layer_outputs, (tuple, list)):
@@ -2132,6 +2152,8 @@ def forward(
         data_id=None,
         src_id=None,
         inbatch_pack_offset=None,
+        attn_mask_startend_row_indices=None,
+        **kwargs,
     ):
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
@@ -2151,6 +2173,7 @@ def forward(
             output_hidden_states=output_hidden_states,
             return_dict=True,
             inbatch_pack_offset=inbatch_pack_offset,
+            attn_mask_startend_row_indices=attn_mask_startend_row_indices,
         )
 
         hidden_states = outputs.last_hidden_state