预训练在线+离线数据流添加attn mask传入 (#3137)

Author: Jonathans575

paddleformers/cli/train/sft/workflow.py

@@ -19,6 +19,7 @@
 import os
 from functools import partial
 
+import numpy as np
 import paddle
 
 is_sm90 = (
@@ -74,7 +75,7 @@
 )
 
 
-def create_pretrained_dataset(training_args, data_args):
+def create_pretrained_dataset(training_args, data_args, model_args):
     assert data_args.input_dir is not None and len(data_args.input_dir.split()) > 1
 
     check_data_split(
@@ -114,16 +115,40 @@ def create_pretrained_dataset(training_args, data_args):
 
     from paddleformers.data import Stack
 
-    def _collate_data(data, stack_fn=Stack()):
-        tokens_ = stack_fn([x["text"] for x in data])
+    def _collate_data(batch, stack_fn=Stack()):
+        input_keys = ["input_ids", "labels", "position_ids", "attn_mask_startend_row_indices"]
+        return_list = []
+        for batch_sequence in batch:
+            # tokens
+            padded_token_ids = np.array([batch_sequence["text"][:-1]])
+            # labels
+            padded_labels = np.array([batch_sequence["text"][1:]])
+            # position_ids
+            padded_position_ids = np.array([sum(batch_sequence["position_ids"], [])[:-1]])
+            return_list.append(
+                [
+                    padded_token_ids,
+                    padded_labels,
+                    padded_position_ids,
+                ]
+            )
+            # attn mask
+            oral_position_ids = batch_sequence["position_ids"]
+            from paddleformers.datasets.collate import (
+                gen_attn_mask_startend_row_indices,
+            )
 
-        labels = tokens_[:, 1:]
-        tokens = tokens_[:, :-1]
+            return_list[-1].append(
+                gen_attn_mask_startend_row_indices(
+                    oral_position_ids,
+                    data_args.max_seq_len + training_args.num_nextn_predict_layers,
+                    model_args.use_global_causal_attn,
+                )[:, :, :-1, :]
+            )
 
-        return {
-            "input_ids": tokens,
-            "labels": labels,
-        }
+        return_list = [np.concatenate(tensor_list) for tensor_list in zip(*return_list)]
+        input_dict = dict(zip(input_keys, return_list))
+        return input_dict
 
     return train_dataset, valid_dataset, test_dataset, _collate_data
 
@@ -337,7 +362,9 @@ def neft_post_hook(module, input, output):
 
     if data_args.dataset_type == "pretrain":
         training_args.test_iters = training_args.eval_iters * 10
-        train_dataset, eval_dataset, test_dataset, data_collator = create_pretrained_dataset(training_args, data_args)
+        train_dataset, eval_dataset, test_dataset, data_collator = create_pretrained_dataset(
+            training_args, data_args, model_args
+        )
     else:
         train_dataset = create_dataset_sft(
             task_group=data_args.train_dataset_path,

paddleformers/data/causal_dataset.py

@@ -443,6 +443,10 @@ def __getitem__(self, idx):
             sample, mask = self.indexed_dataset.get(
                 self.doc_idx[doc_index_f], offset=offset_f, length=offset_l - offset_f + 1
             )
+
+            # position_ids
+            all_position_ids = []
+            all_position_ids.append(list(range(len(sample))))
         else:
             # Otherwise, get the rest of the initial document.
             doc_ids.append(self.doc_idx[doc_index_f])
@@ -468,6 +472,10 @@ def __getitem__(self, idx):
             sample_list.append(sample)
             if append_mask:
                 mask_list.append(mask)
+            # position_ids
+            all_position_ids = []
+            for item in sample_list:
+                all_position_ids.append(list(range(len(item))))
             sample = np.concatenate(sample_list)
             if append_mask:
                 mask = np.concatenate(mask_list)
@@ -505,6 +513,8 @@ def __getitem__(self, idx):
                 "CPT": self.CPT,
             }
 
+        res.update({"position_ids": all_position_ids})
+
         return res
 
 

paddleformers/datasets/collate.py

@@ -208,17 +208,14 @@ def collate_fn(batch: List[List[Sequence]], tokenizer, training_args, model_args
             padded_nbatch_pack_offset = pad_batch_data([nbatch_pack_offset], pad_idx=0, max_seq_len=max_seq_len)
             return_list[-1].append(padded_nbatch_pack_offset)
 
-        if not model_args.stage.lower() == "pt":
-            if model_args.use_attn_mask_startend_row_indices:
-                return_list[-1].append(
-                    gen_attn_mask_startend_row_indices(
-                        original_token_ids, max_seq_len, model_args.use_global_causal_attn
-                    )
-                )
-            else:
-                return_list[-1].append(
-                    gen_self_attn_mask(original_token_ids, max_seq_len, model_args.use_global_causal_attn)
-                )
+        if model_args.use_attn_mask_startend_row_indices:
+            return_list[-1].append(
+                gen_attn_mask_startend_row_indices(original_token_ids, max_seq_len, model_args.use_global_causal_attn)
+            )
+        else:
+            return_list[-1].append(
+                gen_self_attn_mask(original_token_ids, max_seq_len, model_args.use_global_causal_attn)
+            )
 
     return_list = [np.concatenate(tensor_list) for tensor_list in zip(*return_list)]
     input_dict = dict(zip(input_keys, return_list))