Fix zero_padding (#9003)

2024-08-26 21:04:22 +08:00
parent 24fa97ee43
commit 56d293dd77
3 changed files with 44 additions and 68 deletions
--- a/llm/run_finetune.py
+++ b/llm/run_finetune.py
@@ -517,14 +517,19 @@ def main():
        model.print_trainable_parameters()

    # Create trainer
-    max_length = (
-        data_args.max_length
-        if training_args.pipeline_parallel_degree > 1 or training_args.autotuner_benchmark
-        else None
-    )  # NOTE(gongenlei): new add autotuner_benchmark
-    padding = (
-        "max_length" if training_args.pipeline_parallel_degree > 1 or training_args.autotuner_benchmark else True
-    )  # NOTE(gongenlei): new add autotuner_benchmark
+
+    if (
+        training_args.pipeline_parallel_degree > 1
+        or training_args.sequence_parallel
+        or training_args.autotuner_benchmark
+    ):
+        # NOTE(gongenlei): new add autotuner_benchmark
+        max_length = data_args.max_length
+        padding = "max_length"
+    else:
+        max_length = None
+        padding = True
+
    if training_args.pipeline_parallel_degree > 1:
        metrics = None
    elif data_args.eval_with_do_generation:
--- a/paddlenlp/datasets/zero_padding_dataset.py
+++ b/paddlenlp/datasets/zero_padding_dataset.py
@@ -53,42 +53,7 @@ class ZeroPadding:
    ]

    @classmethod
-    def _pad_batch_records_to_max_length(cls, batch_records, max_length, pad_token=0):
-        # confirm the at least one item in the pack
-        if len(batch_records) == 0:
-            return batch_records
-        # count all records total length
-        total_length = sum([len(record["input_ids"]) for record in batch_records])
-        reserved_length = max_length - total_length
-
-        # append padding to the max_length
-        if "attn_mask_startend_row_indices" in batch_records[0]:
-            # attn_mask_startend_row_indices is a list of row indices `0`,
-            # which indicates that all tokens are masked.
-            batch_records.append(
-                {
-                    "input_ids": [pad_token] * reserved_length,
-                    "labels": [-100] * reserved_length,
-                    "attn_mask_startend_row_indices": [0] * reserved_length,
-                }
-            )
-        elif "attention_mask" in batch_records[0]:
-            # attention_mask is a fullly masked attention matrix (all False)
-            # which indicates that all tokens are masked.
-            batch_records.append(
-                {
-                    "input_ids": [pad_token] * reserved_length,
-                    "labels": [-100] * reserved_length,
-                    "attention_mask": np.zeros((reserved_length, reserved_length), dtype=bool),
-                }
-            )
-
-        return batch_records
-
-    @classmethod
-    def _pad_batch_records(cls, batch_records, max_length):
-        batch_records = cls._pad_batch_records_to_max_length(batch_records, max_length)
-
+    def _pad_batch_records(cls, batch_records):
        # Only consider supported input keys
        input_keys = [key for key in batch_records[0].keys() if key in cls.supported_input_keys]
        if "attn_mask_startend_row_indices" not in input_keys and "attention_mask" not in input_keys:
@@ -157,7 +122,7 @@ class ZeroPaddingMapDataset(ZeroPadding, Dataset):
                    cur_len_so_far += len(record["input_ids"])
                else:
                    # exceed max length
-                    padded_list = self._pad_batch_records(batch_records, self.max_length)
+                    padded_list = self._pad_batch_records(batch_records)
                    total_data.append(padded_list)
                    # reset
                    batch_records = []
@@ -168,7 +133,7 @@ class ZeroPaddingMapDataset(ZeroPadding, Dataset):

            # remaining data
            if batch_records:
-                padded_list = self._pad_batch_records(batch_records, self.max_length)
+                padded_list = self._pad_batch_records(batch_records)
                total_data.append(padded_list)
        else:
            examples = []
@@ -185,7 +150,7 @@ class ZeroPaddingMapDataset(ZeroPadding, Dataset):
                    generate_packs = generate_greedy_packs(examples, self.max_length)
                    for batch_records in generate_packs:
                        if len(batch_records) > 0:
-                            padded_list = self._pad_batch_records(batch_records, self.max_length)
+                            padded_list = self._pad_batch_records(batch_records)
                            total_data.append(padded_list)
                    examples = [record]
                    i = 1
@@ -193,7 +158,7 @@ class ZeroPaddingMapDataset(ZeroPadding, Dataset):
                generate_packs = generate_greedy_packs(examples, self.max_length)
                for batch_records in generate_packs:
                    if len(batch_records) > 0:
-                        padded_list = self._pad_batch_records(batch_records, self.max_length)
+                        padded_list = self._pad_batch_records(batch_records)
                        total_data.append(padded_list)

        return total_data
@@ -225,7 +190,7 @@ class ZeroPaddingIterableDataset(ZeroPadding, IterableDataset):
                    cur_len_so_far += len(record["input_ids"])
                else:
                    # exceed max length
-                    padded_list = self._pad_batch_records(batch_records, self.max_length)
+                    padded_list = self._pad_batch_records(batch_records)
                    yield padded_list
                    # reset
                    batch_records = []
@@ -235,7 +200,7 @@ class ZeroPaddingIterableDataset(ZeroPadding, IterableDataset):
                    self.zero_padding_global_step += 1
                    cur_len_so_far += len(record["input_ids"])
            if batch_records:
-                padded_list = self._pad_batch_records(batch_records, self.max_length)
+                padded_list = self._pad_batch_records(batch_records)
                yield padded_list
        else:
            examples = []
@@ -253,7 +218,7 @@ class ZeroPaddingIterableDataset(ZeroPadding, IterableDataset):
                    generate_packs = generate_greedy_packs(examples, self.max_length)
                    for batch_records in generate_packs:
                        if len(batch_records) > 0:
-                            padded_list = self._pad_batch_records(batch_records, self.max_length)
+                            padded_list = self._pad_batch_records(batch_records)
                            yield padded_list
                    examples = [record]
                    self.zero_padding_global_step += 1
@@ -262,5 +227,5 @@ class ZeroPaddingIterableDataset(ZeroPadding, IterableDataset):
                generate_packs = generate_greedy_packs(examples, self.max_length)
                for batch_records in generate_packs:
                    if len(batch_records) > 0:
-                        padded_list = self._pad_batch_records(batch_records, self.max_length)
+                        padded_list = self._pad_batch_records(batch_records)
                        yield padded_list
--- a/tests/dataset/test_zero_padding.py
+++ b/tests/dataset/test_zero_padding.py
@@ -19,7 +19,11 @@ import unittest

 import numpy as np

-from paddlenlp.datasets import InTokensIterableDataset, InTokensMapDataset, load_dataset
+from paddlenlp.datasets import (
+    ZeroPaddingIterableDataset,
+    ZeroPaddingMapDataset,
+    load_dataset,
+)
 from paddlenlp.transformers import AutoTokenizer
 from tests.testing_utils import get_tests_dir

@@ -31,7 +35,7 @@ def read_local_dataset(path):
            yield json.loads(line.strip())


-class InTokensTestCommon:
+class ZeroPaddingTestCommon:
    tokenizer = AutoTokenizer.from_pretrained("__internal_testing__/micro-random-llama")
    expected_output = {
        "input_ids": [1, 29871, 30429, 1, 29871, 30429, 2, 1, 29871, 31427, 1, 29871, 31427, 2],
@@ -90,7 +94,7 @@ class InTokensTestCommon:
        return model_inputs


-class TestInTokensMapDataset(InTokensTestCommon, unittest.TestCase):
+class TestZeroPaddingMapDataset(ZeroPaddingTestCommon, unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        fixture_path = get_tests_dir(os.path.join("fixtures", "dummy"))
@@ -111,26 +115,28 @@ class TestInTokensMapDataset(InTokensTestCommon, unittest.TestCase):
        )

    def test_long_max_length(self):
-        inData = InTokensMapDataset(self.dataset, self.tokenizer, max_length=128)
+        inData = ZeroPaddingMapDataset(self.dataset, self.tokenizer, max_length=128)
        self.assertEqual(set(inData[0].keys()), {"input_ids", "labels", "position_ids", "attention_mask"})
        self.assertEqual(len(inData), 1)
        self.assertEqual(type(inData[0]["input_ids"]), list)
        self.assertEqual(np.array(inData[0]["input_ids"]).shape, (70,))

-        inData_input_labels_only = InTokensMapDataset(self.dataset_input_labels_only, self.tokenizer, max_length=128)
+        inData_input_labels_only = ZeroPaddingMapDataset(
+            self.dataset_input_labels_only, self.tokenizer, max_length=128
+        )
        self.assertEqual(set(inData_input_labels_only[0].keys()), {"input_ids", "labels", "attention_mask"})
        self.assertEqual(len(inData_input_labels_only), 1)
        self.assertEqual(type(inData_input_labels_only[0]["input_ids"]), list)
        self.assertEqual(np.array(inData_input_labels_only[0]["input_ids"]).shape, (70,))

    def test_short_max_length(self):
-        inData = InTokensMapDataset(self.dataset, self.tokenizer, max_length=16)
+        inData = ZeroPaddingMapDataset(self.dataset, self.tokenizer, max_length=16)
        self.assertEqual(inData[0]["input_ids"], self.expected_output["input_ids"])
        self.assertEqual(inData[0]["labels"], self.expected_output["labels"])
        self.assertTrue((inData[0]["position_ids"] == self.expected_output["position_ids"]).all())
        self.assertTrue((inData[0]["attention_mask"] == self.expected_output["attention_mask"]).all())

-        inData_input_labels_only = InTokensMapDataset(self.dataset_input_labels_only, self.tokenizer, max_length=16)
+        inData_input_labels_only = ZeroPaddingMapDataset(self.dataset_input_labels_only, self.tokenizer, max_length=16)
        self.assertEqual(inData_input_labels_only[0]["input_ids"], self.expected_output["input_ids"])
        self.assertEqual(inData_input_labels_only[0]["labels"], self.expected_output["labels"])
        self.assertTrue(
@@ -138,19 +144,19 @@ class TestInTokensMapDataset(InTokensTestCommon, unittest.TestCase):
        )

    def test_2d_position_id(self):
-        inData_2d = InTokensMapDataset(self.dataset_position_2d, self.tokenizer, max_length=16)
+        inData_2d = ZeroPaddingMapDataset(self.dataset_position_2d, self.tokenizer, max_length=16)
        self.assertTrue(inData_2d[0]["position_ids"] == self.expected_output["position_ids_2d"])

    def test_missing_data(self):
        orginal_input_ids = [item["input_ids"] for item in self.dataset]
        orginal_input_ids = [sum(orginal_input_ids, [])]
-        inData = InTokensMapDataset(self.dataset, self.tokenizer, max_length=16)
+        inData = ZeroPaddingMapDataset(self.dataset, self.tokenizer, max_length=16)
        tgt_input_ids = [item["input_ids"] for item in inData]
        tgt_input_ids = [sum(tgt_input_ids, [])]
        self.assertEqual(orginal_input_ids, tgt_input_ids)


-class TestInTokensIterableDataset(InTokensTestCommon, unittest.TestCase):
+class TestZeroPaddingIterableDataset(ZeroPaddingTestCommon, unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        fixture_path = get_tests_dir(os.path.join("fixtures", "dummy"))
@@ -168,13 +174,13 @@ class TestInTokensIterableDataset(InTokensTestCommon, unittest.TestCase):
        )

    def test_long_max_length(self):
-        inData = InTokensIterableDataset(self.dataset, self.tokenizer, max_length=128)
+        inData = ZeroPaddingIterableDataset(self.dataset, self.tokenizer, max_length=128)
        example = next(iter(inData))
        self.assertEqual(set(example.keys()), {"input_ids", "labels", "position_ids", "attention_mask"})
        self.assertEqual(type(example["input_ids"]), list)
        self.assertEqual(np.array(example["input_ids"]).shape, (70,))

-        inData_input_labels_only = InTokensIterableDataset(
+        inData_input_labels_only = ZeroPaddingIterableDataset(
            self.dataset_input_labels_only, self.tokenizer, max_length=128
        )
        example = next(iter(inData_input_labels_only))
@@ -183,14 +189,14 @@ class TestInTokensIterableDataset(InTokensTestCommon, unittest.TestCase):
        self.assertEqual(np.array(example["input_ids"]).shape, (70,))

    def test_short_max_length(self):
-        inData = InTokensIterableDataset(self.dataset, self.tokenizer, max_length=16)
+        inData = ZeroPaddingIterableDataset(self.dataset, self.tokenizer, max_length=16)
        example = next(iter(inData))
        self.assertEqual(example["input_ids"], self.expected_output["input_ids"])
        self.assertEqual(example["labels"], self.expected_output["labels"])
        self.assertTrue((example["position_ids"] == self.expected_output["position_ids"]).all())
        self.assertTrue((example["attention_mask"] == self.expected_output["attention_mask"]).all())

-        inData_input_labels_only = InTokensIterableDataset(
+        inData_input_labels_only = ZeroPaddingIterableDataset(
            self.dataset_input_labels_only, self.tokenizer, max_length=16
        )
        example = next(iter(inData_input_labels_only))
@@ -199,13 +205,13 @@ class TestInTokensIterableDataset(InTokensTestCommon, unittest.TestCase):
        self.assertTrue((example["attention_mask"] == self.expected_output["attention_mask"]).all())

    def test_2d_position_id(self):
-        inData_2d = InTokensIterableDataset(self.dataset_position_2d, self.tokenizer, max_length=16)
+        inData_2d = ZeroPaddingIterableDataset(self.dataset_position_2d, self.tokenizer, max_length=16)
        example = next(iter(inData_2d))
        self.assertTrue(example["position_ids"] == self.expected_output["position_ids_2d"])

    def test_missing_data(self):
        orginal_input_ids = [item["input_ids"] for item in self.dataset]
        orginal_input_ids = [sum(orginal_input_ids, [])]
-        inData = InTokensIterableDataset(self.dataset, self.tokenizer, max_length=128)
+        inData = ZeroPaddingIterableDataset(self.dataset, self.tokenizer, max_length=128)
        tgt_input_ids = [item["input_ids"] for item in inData]
        self.assertEqual(orginal_input_ids, tgt_input_ids)