fix(huggingface): fix huggingface dataloader when using some huggingface third-party tokenizers (#277)

Author: zigzagcai

internlm/data/build_dataloader.py

@@ -125,15 +125,22 @@ def get_hf_train_loader_items(data_cfg):
         model_max_length=data_cfg.seq_len,
         subset_name=data_cfg.get("subset_name", None),
     )
+    pad_token_id = gpc.config.model.get("pad_token_id", 0)
     if gpc.config.model_type == "hf" and not data_cfg.use_packed_dataset:
         train_sampler = StreamingStaticBatchSampler(
             batch_size=data_cfg.micro_num * data_cfg.micro_bsz, rampup_batch_size=data_cfg.rampup_batch_size
         )
         train_collate_fn = partial(
-            nopack_collate_fn, micro_num=data_cfg.micro_num, micro_bsz=data_cfg.micro_bsz, seq_len=data_cfg.seq_len
+            nopack_collate_fn,
+            micro_num=data_cfg.micro_num,
+            micro_bsz=data_cfg.micro_bsz,
+            seq_len=data_cfg.seq_len,
+            pad_token_id=pad_token_id,
         )
     else:
-        train_ds = HuggingFacePackedDataset(dataset=train_ds, seq_len=data_cfg.seq_len, micro_bsz=data_cfg.micro_bsz)
+        train_ds = HuggingFacePackedDataset(
+            dataset=train_ds, seq_len=data_cfg.seq_len, micro_bsz=data_cfg.micro_bsz, pad_token_id=pad_token_id
+        )
         train_sampler = StreamingStaticBatchSampler(
             batch_size=data_cfg.micro_num, rampup_batch_size=data_cfg.rampup_batch_size
         )

internlm/data/streaming/collaters.py

@@ -1,25 +1,34 @@
 import torch
 
 
-def nopack_collate_fn(batch, micro_num, micro_bsz, seq_len):
+def nopack_collate_fn(batch, micro_num, micro_bsz, seq_len, pad_token_id=0):
     input_ids_list = []
     attention_mask_list = []
     labels_list = []
+
     for b in batch:
-        attention_mask = torch.tensor(b["attention_mask"])
-        input_ids = torch.LongTensor(b["input_ids"])
-        input_ids = torch.abs(input_ids * attention_mask)
-        input_ids = torch.nn.functional.pad(input_ids, (0, seq_len - len(input_ids)), mode="constant", value=0)
-        attention_mask = torch.nn.functional.pad(
-            attention_mask, (0, seq_len - len(attention_mask)), mode="constant", value=0
-        )
-        label = torch.LongTensor([w if w > 0 else -100 for w in input_ids.tolist()][1:] + [-100])
-        input_ids_list.append(input_ids)
-        attention_mask_list.append(attention_mask)
-        labels_list.append(label)
+        assert len(b["input_ids"]) > 0
+
+        if "attention_mask" in b:
+            assert len(b["input_ids"]) == len(
+                b["attention_mask"]
+            ), "input_ids and attention_mask should be equal length"
+        else:
+            b["attention_mask"] = [True] * len(b["input_ids"])
+
+        input_ids = b["input_ids"] + [pad_token_id] * (seq_len - len(b["input_ids"]))
+        attention_mask = b["attention_mask"] + [False] * (seq_len - len(b["attention_mask"]))
+        labels = [w if w > 0 else -100 for w in b["input_ids"]][1:] + [-100]
+        labels = labels + [-100] * (seq_len - len(b["input_ids"]))
+
+        input_ids_list.append(torch.LongTensor(input_ids))
+        attention_mask_list.append(torch.BoolTensor(attention_mask))
+        labels_list.append(torch.LongTensor(labels))
+
     input_ids = torch.stack(input_ids_list)
     attention_mask = torch.stack(attention_mask_list)
     labels = torch.stack(labels_list)
+
     return {
         "input_ids": input_ids,
         "attention_mask": attention_mask,

internlm/data/streaming/dataset.py

@@ -47,22 +47,24 @@ def _tokenize(self, samples):
         texts = [sample["text"] for sample in samples]
         tokenized_outputs = self.tokenizer(texts, truncation=True)
         for i in range(len(samples)):
-            yield {key: tokenized_outputs[key][i] for key in tokenized_outputs}
+            assert "input_ids" in tokenized_outputs, "huggingface tokenizer should generate input_ids"
+            if len(tokenized_outputs["input_ids"][i]) > 0:
+                yield {key: tokenized_outputs[key][i] for key in tokenized_outputs}
 
     def __getitem__(self, _):
         return next(self.senior_iterator)
 
 
 class HuggingFacePackedDataset(Dataset):
     """
-    Simple packed dataset for huggingface.
+    Simple packed dataset for huggingface
     """
 
-    def __init__(self, dataset, seq_len, micro_bsz):
+    def __init__(self, dataset, seq_len, micro_bsz, pad_token_id=0):
         self.dataset = dataset
         self.seq_len = seq_len
         self.micro_bsz = micro_bsz
-
+        self.pad_token_id = pad_token_id
         self.senior_iterator = iter(self)
 
     def __iter__(self):
@@ -72,7 +74,7 @@ def __iter__(self):
         for sample in self.dataset:
             if len(input_ids + sample["input_ids"]) > self.micro_bsz * self.seq_len:
                 assert cu_seqlens[-1] <= self.micro_bsz * self.seq_len
-                input_ids = input_ids + [0] * (self.micro_bsz * self.seq_len - len(input_ids))
+                input_ids = input_ids + [self.pad_token_id] * (self.micro_bsz * self.seq_len - len(input_ids))
                 cu_seqlens = (
                     cu_seqlens + [self.micro_bsz * self.seq_len]
                     if cu_seqlens[-1] < self.micro_bsz * self.seq_len
@@ -89,14 +91,15 @@ def __iter__(self):
                 }
                 input_ids = sample["input_ids"]
                 cu_seqlens = [0, len(sample["input_ids"])]
-                labels = sample["input_ids"][1:] + [-100]
+                labels = [w if w > 0 else -100 for w in sample["input_ids"]][1:] + [-100]
             else:
                 input_ids = input_ids + sample["input_ids"]
                 cu_seqlens.append(len(sample["input_ids"]) + cu_seqlens[-1])
-                labels = labels + sample["input_ids"][1:] + [-100]
+                labels = labels + [w if w > 0 else -100 for w in sample["input_ids"]][1:] + [-100]
+
         if input_ids:
             assert cu_seqlens[-1] <= self.micro_bsz * self.seq_len
-            input_ids = input_ids + [0] * (self.micro_bsz * self.seq_len - len(input_ids))
+            input_ids = input_ids + [self.pad_token_id] * (self.micro_bsz * self.seq_len - len(input_ids))
             cu_seqlens = (
                 cu_seqlens + [self.micro_bsz * self.seq_len]
                 if cu_seqlens[-1] < self.micro_bsz * self.seq_len