Add predict step accumulation (#7767)

* Add eval_accumulation_step and clean distributed eval

* Add TPU test

* Add TPU stuff

* Fix arg name

* Fix Seq2SeqTrainer

* Fix total_size

* Update src/transformers/trainer_pt_utils.py

Co-authored-by: Lysandre Debut <lysandre@huggingface.co>

* Doc and add test to TPU

* Add unit test

* Adapt name

Co-authored-by: Lysandre Debut <lysandre@huggingface.co>

tests/test_trainer.py

@@ -1,3 +1,18 @@
+# coding=utf-8
+# Copyright 2018 the HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import dataclasses
 import os
 import tempfile

tests/test_trainer_distributed.py

@@ -13,15 +13,14 @@
 #   CUDA_VISIBLE_DEVICES=-1 python ./tests/test_trainer_distributed.py
 #
 
-
-import logging
 import sys
 from typing import Dict
 
 from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available
+from transformers.utils import logging
 
 
-logger = logging.getLogger(__name__)
+logger = logging.get_logger(__name__)
 
 
 if is_torch_available():
@@ -101,4 +100,20 @@ if __name__ == "__main__":
             logger.error(p.metrics)
             exit(1)
 
+        trainer.args.eval_accumulation_steps = 2
+
+        metrics = trainer.evaluate()
+        logger.info(metrics)
+        if metrics["eval_success"] is not True:
+            logger.error(metrics)
+            exit(1)
+
+        p = trainer.predict(dataset)
+        logger.info(p.metrics)
+        if p.metrics["eval_success"] is not True:
+            logger.error(p.metrics)
+            exit(1)
+
+        trainer.args.eval_accumulation_steps = None
+
     logger.info("🔥 All distributed tests successful")

tests/test_trainer_tpu.py

@@ -0,0 +1,119 @@
+# This test is meant to be run in on an instance with TPUs like this:
+#
+#   python examples/xla_spawn.py --num_cores=8 tests/test_trainer_tpu.py
+#
+# Replace 8 with the number of TPU cores you have.
+#
+
+import sys
+from typing import Dict
+
+from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+if is_torch_available():
+    import torch
+    from torch import nn
+    from torch.utils.data.dataset import Dataset
+
+    from transformers import Trainer
+
+    class DummyDataset(Dataset):
+        def __init__(self, length: int = 101):
+            self.length = length
+
+        def __len__(self):
+            return self.length
+
+        def __getitem__(self, i) -> int:
+            return i
+
+    class DummyDataCollator:
+        def __call__(self, features):
+            return {"input_ids": torch.tensor(features), "labels": torch.tensor(features)}
+
+    class DummyModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            # Add some (unused) params otherwise DDP will complain.
+            self.fc = nn.Linear(120, 80)
+
+        def forward(self, input_ids, labels=None):
+            if labels is not None:
+                return torch.tensor(0.0, device=input_ids.device), input_ids
+            else:
+                return input_ids
+
+
+def main():
+    parser = HfArgumentParser((TrainingArguments,))
+    sys.argv += ["--output_dir", "./examples"]
+    training_args = parser.parse_args_into_dataclasses()[0]
+
+    logger.warning(
+        "Process rank: %s, device: %s, tpu_num_cores: %s",
+        training_args.local_rank,
+        training_args.device,
+        training_args.tpu_num_cores,
+    )
+
+    # Essentially, what we want to verify in the distributed case is
+    # that we get all samples back, in the right order.
+    # (this is crucial for prediction for instance)
+    for dataset_length in [1001, 256, 15]:
+        dataset = DummyDataset(dataset_length)
+
+        def compute_metrics(p: EvalPrediction) -> Dict:
+            sequential = list(range(len(dataset)))
+            success = p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential
+            return {"success": success}
+
+        trainer = Trainer(
+            model=DummyModel(),
+            args=training_args,
+            data_collator=DummyDataCollator(),
+            eval_dataset=dataset,
+            compute_metrics=compute_metrics,
+        )
+        metrics = trainer.evaluate()
+        logger.info(metrics)
+        if metrics["eval_success"] is not True:
+            logger.error(metrics)
+            exit(1)
+
+        p = trainer.predict(dataset)
+        logger.info(p.metrics)
+        if p.metrics["eval_success"] is not True:
+            logger.error(p.metrics)
+            exit(1)
+
+        trainer.args.eval_accumulation_steps = 2
+
+        metrics = trainer.evaluate()
+        logger.info(metrics)
+        if metrics["eval_success"] is not True:
+            logger.error(metrics)
+            exit(1)
+
+        p = trainer.predict(dataset)
+        logger.info(p.metrics)
+        if p.metrics["eval_success"] is not True:
+            logger.error(p.metrics)
+            exit(1)
+
+        trainer.args.eval_accumulation_steps = None
+
+    logger.info("🔥 All distributed tests successful")
+
+
+def _mp_fn(index):
+    # For xla_spawn (TPUs)
+    main()
+
+
+if __name__ == "__main__":
+    main()

tests/test_trainer_utils.py

@@ -0,0 +1,58 @@
+# coding=utf-8
+# Copyright 2018 the HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+from transformers.file_utils import is_torch_available
+from transformers.testing_utils import require_torch
+
+
+if is_torch_available():
+    from transformers.trainer_pt_utils import DistributedTensorGatherer
+
+
+@require_torch
+class TrainerUtilsTest(unittest.TestCase):
+    def test_distributed_tensor_gatherer(self):
+        # Simulate a result with a dataset of size 21, 4 processes and chunks of lengths 2, 3, 1
+        world_size = 4
+        num_samples = 21
+        input_indices = [
+            [0, 1, 6, 7, 12, 13, 18, 19],
+            [2, 3, 4, 8, 9, 10, 14, 15, 16, 20, 0, 1],
+            [5, 11, 17, 2],
+        ]
+
+        predictions = np.random.normal(size=(num_samples, 13))
+        gatherer = DistributedTensorGatherer(world_size=world_size, num_samples=num_samples)
+        for indices in input_indices:
+            gatherer.add_arrays(predictions[indices])
+        result = gatherer.finalize()
+        self.assertTrue(np.array_equal(result, predictions))
+
+        # With nested tensors
+        gatherer = DistributedTensorGatherer(world_size=world_size, num_samples=num_samples)
+        for indices in input_indices:
+            gatherer.add_arrays([predictions[indices], [predictions[indices], predictions[indices]]])
+        result = gatherer.finalize()
+        self.assertTrue(isinstance(result, list))
+        self.assertTrue(len(result), 2)
+        self.assertTrue(isinstance(result[1], list))
+        self.assertTrue(len(result[1]), 2)
+        self.assertTrue(np.array_equal(result[0], predictions))
+        self.assertTrue(np.array_equal(result[1][0], predictions))
+        self.assertTrue(np.array_equal(result[1][1], predictions))