trainer.py 44.9 KB
Newer Older
Julien Chaumond's avatar
Julien Chaumond committed
1
import logging
2
import math
Julien Chaumond's avatar
Julien Chaumond committed
3
4
5
import os
import re
import shutil
6
import warnings
Julien Chaumond's avatar
Julien Chaumond committed
7
8
from contextlib import contextmanager
from pathlib import Path
9
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
Julien Chaumond's avatar
Julien Chaumond committed
10
11
12

import numpy as np
import torch
13
from packaging import version
Julien Chaumond's avatar
Julien Chaumond committed
14
15
16
17
from torch import nn
from torch.utils.data.dataloader import DataLoader
from torch.utils.data.dataset import Dataset
from torch.utils.data.distributed import DistributedSampler
18
from torch.utils.data.sampler import RandomSampler, Sampler, SequentialSampler
19
from tqdm.auto import tqdm, trange
Julien Chaumond's avatar
Julien Chaumond committed
20

21
from .data.data_collator import DataCollator, default_data_collator
22
from .file_utils import is_torch_tpu_available
Julien Chaumond's avatar
Julien Chaumond committed
23
24
from .modeling_utils import PreTrainedModel
from .optimization import AdamW, get_linear_schedule_with_warmup
25
26
27
28
29
30
31
32
from .trainer_utils import (
    PREFIX_CHECKPOINT_DIR,
    EvalPrediction,
    PredictionOutput,
    TrainOutput,
    is_wandb_available,
    set_seed,
)
Patrick von Platen's avatar
Patrick von Platen committed
33
from .training_args import TrainingArguments
Julien Chaumond's avatar
Julien Chaumond committed
34
35


36
37
38
39
40
41
42
43
44
45
46
47
48
_use_native_amp = False
_use_apex = False

# Check if Pytorch version >= 1.6 to switch between Native AMP and Apex
if version.parse(torch.__version__) < version.parse("1.6"):
    from transformers.file_utils import is_apex_available

    if is_apex_available():
        from apex import amp
    _use_apex = True
else:
    _use_native_amp = True
    from torch.cuda.amp import autocast
Julien Chaumond's avatar
Julien Chaumond committed
49
50


51
if is_torch_tpu_available():
Lysandre Debut's avatar
Lysandre Debut committed
52
53
54
55
    import torch_xla.core.xla_model as xm
    import torch_xla.debug.metrics as met
    import torch_xla.distributed.parallel_loader as pl

Julien Chaumond's avatar
Julien Chaumond committed
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
try:
    from torch.utils.tensorboard import SummaryWriter

    _has_tensorboard = True
except ImportError:
    try:
        from tensorboardX import SummaryWriter

        _has_tensorboard = True
    except ImportError:
        _has_tensorboard = False


def is_tensorboard_available():
    return _has_tensorboard


73
if is_wandb_available():
74
75
76
    import wandb


Julien Chaumond's avatar
Julien Chaumond committed
77
78
79
80
81
82
logger = logging.getLogger(__name__)


@contextmanager
def torch_distributed_zero_first(local_rank: int):
    """
83
    Decorator to make all processes in distributed training wait for each local_master to do something.
84
85
86

    Args:
        local_rank (:obj:`int`): The rank of the local process.
Julien Chaumond's avatar
Julien Chaumond committed
87
88
89
90
91
92
93
94
    """
    if local_rank not in [-1, 0]:
        torch.distributed.barrier()
    yield
    if local_rank == 0:
        torch.distributed.barrier()


95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
class SequentialDistributedSampler(Sampler):
    """
    Distributed Sampler that subsamples indicies sequentially,
    making it easier to collate all results at the end.

    Even though we only use this sampler for eval and predict (no training),
    which means that the model params won't have to be synced (i.e. will not hang
    for synchronization even if varied number of forward passes), we still add extra
    samples to the sampler to make it evenly divisible (like in `DistributedSampler`)
    to make it easy to `gather` or `reduce` resulting tensors at the end of the loop.
    """

    def __init__(self, dataset, num_replicas=None, rank=None):
        if num_replicas is None:
            if not torch.distributed.is_available():
                raise RuntimeError("Requires distributed package to be available")
            num_replicas = torch.distributed.get_world_size()
        if rank is None:
            if not torch.distributed.is_available():
                raise RuntimeError("Requires distributed package to be available")
            rank = torch.distributed.get_rank()
        self.dataset = dataset
        self.num_replicas = num_replicas
        self.rank = rank
        self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.num_replicas))
        self.total_size = self.num_samples * self.num_replicas

    def __iter__(self):
        indices = list(range(len(self.dataset)))

        # add extra samples to make it evenly divisible
        indices += indices[: (self.total_size - len(indices))]
Teven's avatar
Teven committed
127
128
129
        assert (
            len(indices) == self.total_size
        ), f"Indices length {len(indices)} and total size {self.total_size} mismatched"
130
131
132

        # subsample
        indices = indices[self.rank * self.num_samples : (self.rank + 1) * self.num_samples]
Teven's avatar
Teven committed
133
134
135
        assert (
            len(indices) == self.num_samples
        ), f"Indices length {len(indices)} and and sample number {self.num_samples} mismatched"
136
137
138
139
140
141
142

        return iter(indices)

    def __len__(self):
        return self.num_samples


Lysandre Debut's avatar
Lysandre Debut committed
143
144
145
146
147
148
def get_tpu_sampler(dataset: Dataset):
    if xm.xrt_world_size() <= 1:
        return RandomSampler(dataset)
    return DistributedSampler(dataset, num_replicas=xm.xrt_world_size(), rank=xm.get_ordinal())


Julien Chaumond's avatar
Julien Chaumond committed
149
150
151
class Trainer:
    """
    Trainer is a simple but feature-complete training and eval loop for PyTorch,
152
153
154
155
156
157
158
159
160
161
    optimized for 馃 Transformers.

    Args:
        model (:class:`~transformers.PreTrainedModel`):
            The model to train, evaluate or use for predictions.
        args (:class:`~transformers.TrainingArguments`):
            The arguments to tweak training.
        data_collator (:obj:`DataCollator`, `optional`, defaults to :func:`~transformers.default_data_collator`):
            The function to use to from a batch from a list of elements of :obj:`train_dataset` or
            :obj:`eval_dataset`.
Sylvain Gugger's avatar
Sylvain Gugger committed
162
        train_dataset (:obj:`torch.utils.data.dataset.Dataset`, `optional`):
163
            The dataset to use for training.
Sylvain Gugger's avatar
Sylvain Gugger committed
164
        eval_dataset (:obj:`torch.utils.data.dataset.Dataset`, `optional`):
165
166
167
168
169
170
171
172
173
174
175
176
            The dataset to use for evaluation.
        compute_metrics (:obj:`Callable[[EvalPrediction], Dict]`, `optional`):
            The function that will be used to compute metrics at evaluation. Must take a
            :class:`~transformers.EvalPrediction` and return a dictionary string to metric values.
        prediction_loss_only (:obj:`bool`, `optional`, defaults to `False`):
            When performing evaluation and predictions, only returns the loss.
        tb_writer (:obj:`SummaryWriter`, `optional`):
            Object to write to TensorBoard.
        optimizers (:obj:`Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR`, `optional`):
            A tuple containing the optimizer and the scheduler to use. Will default to an instance of
            :class:`~transformers.AdamW` on your model and a scheduler given by
            :func:`~transformers.get_linear_schedule_with_warmup` controlled by :obj:`args`.
Julien Chaumond's avatar
Julien Chaumond committed
177
178
179
180
181
182
183
184
185
186
187
    """

    def __init__(
        self,
        model: PreTrainedModel,
        args: TrainingArguments,
        data_collator: Optional[DataCollator] = None,
        train_dataset: Optional[Dataset] = None,
        eval_dataset: Optional[Dataset] = None,
        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
        prediction_loss_only=False,
188
        tb_writer: Optional["SummaryWriter"] = None,
189
        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
Julien Chaumond's avatar
Julien Chaumond committed
190
    ):
191
        self.model = model.to(args.device)
Julien Chaumond's avatar
Julien Chaumond committed
192
        self.args = args
193
        self.data_collator = data_collator if data_collator is not None else default_data_collator
Julien Chaumond's avatar
Julien Chaumond committed
194
195
196
197
        self.train_dataset = train_dataset
        self.eval_dataset = eval_dataset
        self.compute_metrics = compute_metrics
        self.prediction_loss_only = prediction_loss_only
198
199
200
        self.optimizer, self.lr_scheduler = optimizers
        self.tb_writer = tb_writer
        if tb_writer is None and is_tensorboard_available() and self.is_world_process_zero():
Julien Chaumond's avatar
Julien Chaumond committed
201
202
203
204
205
            self.tb_writer = SummaryWriter(log_dir=self.args.logging_dir)
        if not is_tensorboard_available():
            logger.warning(
                "You are instantiating a Trainer but Tensorboard is not installed. You should consider installing it."
            )
206
        if is_wandb_available():
207
            self.setup_wandb()
208
        elif os.environ.get("WANDB_DISABLED") != "true":
209
            logger.info(
210
211
                "You are instantiating a Trainer but W&B is not installed. To use wandb logging, "
                "run `pip install wandb; wandb login` see https://docs.wandb.com/huggingface."
212
            )
Julien Chaumond's avatar
Julien Chaumond committed
213
214
        set_seed(self.args.seed)
        # Create output directory if needed
215
        if self.is_world_process_zero():
Julien Chaumond's avatar
Julien Chaumond committed
216
            os.makedirs(self.args.output_dir, exist_ok=True)
217
        if is_torch_tpu_available():
Lysandre Debut's avatar
Lysandre Debut committed
218
219
220
            # Set an xla_device flag on the model's config.
            # We'll find a more elegant and not need to do this in the future.
            self.model.config.xla_device = True
221
222
223
224
225
226
227
228
229
        if not callable(self.data_collator) and callable(getattr(self.data_collator, "collate_batch", None)):
            self.data_collator = self.data_collator.collate_batch
            warnings.warn(
                (
                    "The `data_collator` should now be a simple callable (function, class with `__call__`), classes "
                    + "with a `collate_batch` are deprecated and won't be supported in a future version."
                ),
                FutureWarning,
            )
230
231
        self.global_step = None
        self.epoch = None
232
233
        if self.args.fp16 and _use_native_amp:
            self.scaler = torch.cuda.amp.GradScaler()
Julien Chaumond's avatar
Julien Chaumond committed
234

235
    def _get_train_sampler(self) -> Optional[torch.utils.data.sampler.Sampler]:
236
        if isinstance(self.train_dataset, torch.utils.data.IterableDataset):
237
            return None
238
        elif is_torch_tpu_available():
239
            return get_tpu_sampler(self.train_dataset)
Lysandre Debut's avatar
Lysandre Debut committed
240
        else:
241
            return (
Lysandre Debut's avatar
Lysandre Debut committed
242
243
244
245
                RandomSampler(self.train_dataset)
                if self.args.local_rank == -1
                else DistributedSampler(self.train_dataset)
            )
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260

    def get_train_dataloader(self) -> DataLoader:
        """
        Returns the training :class:`~torch.utils.data.DataLoader`.

        Will use no sampler if :obj:`self.train_dataset` is a :obj:`torch.utils.data.IterableDataset`, a random sampler
        (adapted to distributed training if necessary) otherwise.

        Subclass and override this method if you want to inject some custom behavior.
        """
        if self.train_dataset is None:
            raise ValueError("Trainer: training requires a train_dataset.")
        train_sampler = self._get_train_sampler()

        return DataLoader(
Julien Chaumond's avatar
Julien Chaumond committed
261
262
263
            self.train_dataset,
            batch_size=self.args.train_batch_size,
            sampler=train_sampler,
264
            collate_fn=self.data_collator,
Setu Shah's avatar
Setu Shah committed
265
            drop_last=self.args.dataloader_drop_last,
Julien Chaumond's avatar
Julien Chaumond committed
266
267
        )

268
269
270
271
272
273
274
275
276
    def _get_eval_sampler(self, eval_dataset: Dataset) -> Optional[torch.utils.data.sampler.Sampler]:
        if isinstance(eval_dataset, torch.utils.data.IterableDataset):
            return None
        elif is_torch_tpu_available():
            return SequentialDistributedSampler(eval_dataset, num_replicas=xm.xrt_world_size(), rank=xm.get_ordinal())
        elif self.args.local_rank != -1:
            return SequentialDistributedSampler(eval_dataset)
        else:
            return SequentialSampler(eval_dataset)
Lysandre Debut's avatar
Lysandre Debut committed
277

Julien Chaumond's avatar
Julien Chaumond committed
278
    def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoader:
279
280
281
        """
        Returns the evaluation :class:`~torch.utils.data.DataLoader`.

282
283
284
285
286
        Will use no sampler if :obj:`self.eval_dataset` is a :obj:`torch.utils.data.IterableDataset`, a sequential
        sampler (adapted to distributed training if necessary) otherwise.

        Subclass and override this method if you want to inject some custom behavior.

287
        Args:
288
289
            eval_dataset (:obj:`torch.utils.data.dataset.Dataset`, `optional`):
                If provided, will override :obj:`self.eval_dataset`.
290
        """
Julien Chaumond's avatar
Julien Chaumond committed
291
292
        if eval_dataset is None and self.eval_dataset is None:
            raise ValueError("Trainer: evaluation requires an eval_dataset.")
Lysandre Debut's avatar
Lysandre Debut committed
293

294
        eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset
295
        eval_sampler = self._get_eval_sampler(eval_dataset)
296

297
        return DataLoader(
298
            eval_dataset,
299
            sampler=eval_sampler,
Julien Chaumond's avatar
Julien Chaumond committed
300
            batch_size=self.args.eval_batch_size,
301
            collate_fn=self.data_collator,
Setu Shah's avatar
Setu Shah committed
302
            drop_last=self.args.dataloader_drop_last,
Julien Chaumond's avatar
Julien Chaumond committed
303
304
305
        )

    def get_test_dataloader(self, test_dataset: Dataset) -> DataLoader:
306
307
308
        """
        Returns the test :class:`~torch.utils.data.DataLoader`.

309
310
311
312
313
        Will use no sampler if :obj:`test_dataset` is a :obj:`torch.utils.data.IterableDataset`, a sequential
        sampler (adapted to distributed training if necessary) otherwise.

        Subclass and override this method if you want to inject some custom behavior.

314
        Args:
315
316
            eval_dataset (:obj:`torch.utils.data.dataset.Dataset`, `optional`):
                The test dataset to use.
317
        """
318
        test_sampler = self._get_eval_sampler(test_dataset)
Lysandre Debut's avatar
Lysandre Debut committed
319

320
321
        # We use the same batch_size as for eval.
        return DataLoader(
Julien Chaumond's avatar
Julien Chaumond committed
322
            test_dataset,
323
            sampler=test_sampler,
Julien Chaumond's avatar
Julien Chaumond committed
324
            batch_size=self.args.eval_batch_size,
325
            collate_fn=self.data_collator,
326
            drop_last=self.args.dataloader_drop_last,
Julien Chaumond's avatar
Julien Chaumond committed
327
        )
Lysandre Debut's avatar
Lysandre Debut committed
328

329
    def create_optimizer_and_scheduler(self, num_training_steps: int):
330
331
332
        """
        Setup the optimizer and the learning rate scheduler.

333
        We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the
334
        Trainer's init through :obj:`optimizers`, or subclass and override this method in a subclass.
335
        """
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
        if self.optimizer is None:
            no_decay = ["bias", "LayerNorm.weight"]
            optimizer_grouped_parameters = [
                {
                    "params": [p for n, p in self.model.named_parameters() if not any(nd in n for nd in no_decay)],
                    "weight_decay": self.args.weight_decay,
                },
                {
                    "params": [p for n, p in self.model.named_parameters() if any(nd in n for nd in no_decay)],
                    "weight_decay": 0.0,
                },
            ]
            self.optimizer = AdamW(
                optimizer_grouped_parameters,
                lr=self.args.learning_rate,
                betas=(self.args.adam_beta1, self.args.adam_beta2),
                eps=self.args.adam_epsilon,
            )
        if self.lr_scheduler is None:
            self.lr_scheduler = get_linear_schedule_with_warmup(
                self.optimizer, num_warmup_steps=self.args.warmup_steps, num_training_steps=num_training_steps
            )
Julien Chaumond's avatar
Julien Chaumond committed
358

359
    def setup_wandb(self):
360
361
362
        """
        Setup the optional Weights & Biases (`wandb`) integration.

363
364
        One can subclass and override this method to customize the setup if needed. Find more information
        `here <https://docs.wandb.com/huggingface>`__. You can also override the following environment variables:
365
366
367
368
369
370
371
372
373

        Environment:
            WANDB_WATCH:
                (Optional, ["gradients", "all", "false"]) "gradients" by default, set to "false" to disable gradient logging
                or "all" to log gradients and parameters
            WANDB_PROJECT:
                (Optional): str - "huggingface" by default, set this to a custom string to store results in a different project
            WANDB_DISABLED:
                (Optional): boolean - defaults to false, set to "true" to disable wandb entirely
374
        """
375
376
377
378
379
380
381
        if hasattr(self, "_setup_wandb"):
            warnings.warn(
                "The `_setup_wandb` method is deprecated and won't be called in a future version, define `setup_wandb` in your subclass.",
                FutureWarning,
            )
            return self._setup_wandb()

382
        if self.is_world_process_zero():
383
384
            logger.info(
                'Automatic Weights & Biases logging enabled, to disable set os.environ["WANDB_DISABLED"] = "true"'
385
            )
386
387
388
389
            combined_dict = {**self.model.config.to_dict(), **self.args.to_sanitized_dict()}
            wandb.init(
                project=os.getenv("WANDB_PROJECT", "huggingface"), config=combined_dict, name=self.args.run_name
            )
390
391
            # keep track of model topology and gradients, unsupported on TPU
            if not is_torch_tpu_available() and os.getenv("WANDB_WATCH") != "false":
392
393
394
                wandb.watch(
                    self.model, log=os.getenv("WANDB_WATCH", "gradients"), log_freq=max(100, self.args.logging_steps)
                )
395

396
    def num_examples(self, dataloader: DataLoader) -> int:
397
        """
398
        Helper to get number of samples in a :class:`~torch.utils.data.DataLoader` by accessing its dataset.
399
        """
400
        return len(dataloader.dataset)
401

Julien Chaumond's avatar
Julien Chaumond committed
402
403
404
405
406
    def train(self, model_path: Optional[str] = None):
        """
        Main training entry point.

        Args:
407
408
409
            model_path (:obj:`str`, `optional`):
                Local path to the model if the model to train has been instantiated from a local path. If present,
                training will resume from the optimizer/scheduler states loaded here.
Julien Chaumond's avatar
Julien Chaumond committed
410
411
412
413
414
415
416
417
418
419
420
        """
        train_dataloader = self.get_train_dataloader()
        if self.args.max_steps > 0:
            t_total = self.args.max_steps
            num_train_epochs = (
                self.args.max_steps // (len(train_dataloader) // self.args.gradient_accumulation_steps) + 1
            )
        else:
            t_total = int(len(train_dataloader) // self.args.gradient_accumulation_steps * self.args.num_train_epochs)
            num_train_epochs = self.args.num_train_epochs

421
        self.create_optimizer_and_scheduler(num_training_steps=t_total)
Julien Chaumond's avatar
Julien Chaumond committed
422
423
424
425
426
427
428
429

        # Check if saved optimizer or scheduler states exist
        if (
            model_path is not None
            and os.path.isfile(os.path.join(model_path, "optimizer.pt"))
            and os.path.isfile(os.path.join(model_path, "scheduler.pt"))
        ):
            # Load in optimizer and scheduler states
430
            self.optimizer.load_state_dict(
431
432
                torch.load(os.path.join(model_path, "optimizer.pt"), map_location=self.args.device)
            )
433
            self.lr_scheduler.load_state_dict(torch.load(os.path.join(model_path, "scheduler.pt")))
Julien Chaumond's avatar
Julien Chaumond committed
434
435

        model = self.model
436
        if self.args.fp16 and _use_apex:
Julien Chaumond's avatar
Julien Chaumond committed
437
438
            if not is_apex_available():
                raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
439
            model, self.optimizer = amp.initialize(model, self.optimizer, opt_level=self.args.fp16_opt_level)
Julien Chaumond's avatar
Julien Chaumond committed
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455

        # multi-gpu training (should be after apex fp16 initialization)
        if self.args.n_gpu > 1:
            model = torch.nn.DataParallel(model)

        # Distributed training (should be after apex fp16 initialization)
        if self.args.local_rank != -1:
            model = torch.nn.parallel.DistributedDataParallel(
                model,
                device_ids=[self.args.local_rank],
                output_device=self.args.local_rank,
                find_unused_parameters=True,
            )

        if self.tb_writer is not None:
            self.tb_writer.add_text("args", self.args.to_json_string())
456
            self.tb_writer.add_hparams(self.args.to_sanitized_dict(), metric_dict={})
Julien Chaumond's avatar
Julien Chaumond committed
457
458

        # Train!
459
        if is_torch_tpu_available():
Lysandre Debut's avatar
Lysandre Debut committed
460
461
462
463
464
            total_train_batch_size = self.args.train_batch_size * xm.xrt_world_size()
        else:
            total_train_batch_size = (
                self.args.train_batch_size
                * self.args.gradient_accumulation_steps
465
                * (torch.distributed.get_world_size() if self.args.local_rank != -1 else 1)
Lysandre Debut's avatar
Lysandre Debut committed
466
            )
Julien Chaumond's avatar
Julien Chaumond committed
467
        logger.info("***** Running training *****")
468
        logger.info("  Num examples = %d", self.num_examples(train_dataloader))
Julien Chaumond's avatar
Julien Chaumond committed
469
        logger.info("  Num Epochs = %d", num_train_epochs)
470
        logger.info("  Instantaneous batch size per device = %d", self.args.per_device_train_batch_size)
Lysandre Debut's avatar
Lysandre Debut committed
471
        logger.info("  Total train batch size (w. parallel, distributed & accumulation) = %d", total_train_batch_size)
Julien Chaumond's avatar
Julien Chaumond committed
472
473
474
        logger.info("  Gradient Accumulation steps = %d", self.args.gradient_accumulation_steps)
        logger.info("  Total optimization steps = %d", t_total)

475
476
        self.global_step = 0
        self.epoch = 0
Julien Chaumond's avatar
Julien Chaumond committed
477
478
479
480
481
482
        epochs_trained = 0
        steps_trained_in_current_epoch = 0
        # Check if continuing training from a checkpoint
        if model_path is not None:
            # set global_step to global_step of last saved checkpoint from model path
            try:
483
484
485
                self.global_step = int(model_path.split("-")[-1].split("/")[0])
                epochs_trained = self.global_step // (len(train_dataloader) // self.args.gradient_accumulation_steps)
                steps_trained_in_current_epoch = self.global_step % (
Julien Chaumond's avatar
Julien Chaumond committed
486
487
488
489
490
                    len(train_dataloader) // self.args.gradient_accumulation_steps
                )

                logger.info("  Continuing training from checkpoint, will skip to saved global_step")
                logger.info("  Continuing training from epoch %d", epochs_trained)
491
                logger.info("  Continuing training from global step %d", self.global_step)
Julien Chaumond's avatar
Julien Chaumond committed
492
493
                logger.info("  Will skip the first %d steps in the first epoch", steps_trained_in_current_epoch)
            except ValueError:
494
                self.global_step = 0
Julien Chaumond's avatar
Julien Chaumond committed
495
496
497
498
499
500
                logger.info("  Starting fine-tuning.")

        tr_loss = 0.0
        logging_loss = 0.0
        model.zero_grad()
        train_iterator = trange(
501
            epochs_trained, int(num_train_epochs), desc="Epoch", disable=not self.is_local_process_zero()
Julien Chaumond's avatar
Julien Chaumond committed
502
503
        )
        for epoch in train_iterator:
504
505
506
            if isinstance(train_dataloader, DataLoader) and isinstance(train_dataloader.sampler, DistributedSampler):
                train_dataloader.sampler.set_epoch(epoch)

507
            if is_torch_tpu_available():
508
509
510
                parallel_loader = pl.ParallelLoader(train_dataloader, [self.args.device]).per_device_loader(
                    self.args.device
                )
511
                epoch_iterator = tqdm(parallel_loader, desc="Iteration", disable=not self.is_local_process_zero())
512
            else:
513
                epoch_iterator = tqdm(train_dataloader, desc="Iteration", disable=not self.is_local_process_zero())
514

515
516
517
518
            # Reset the past mems state at the beginning of each epoch if necessary.
            if self.args.past_index >= 0:
                self._past = None

Julien Chaumond's avatar
Julien Chaumond committed
519
520
521
522
523
524
525
            for step, inputs in enumerate(epoch_iterator):

                # Skip past any already trained steps if resuming training
                if steps_trained_in_current_epoch > 0:
                    steps_trained_in_current_epoch -= 1
                    continue

526
                tr_loss += self.training_step(model, inputs)
Julien Chaumond's avatar
Julien Chaumond committed
527
528
529
530
531
532

                if (step + 1) % self.args.gradient_accumulation_steps == 0 or (
                    # last step in epoch but step is always smaller than gradient_accumulation_steps
                    len(epoch_iterator) <= self.args.gradient_accumulation_steps
                    and (step + 1) == len(epoch_iterator)
                ):
533
                    if self.args.fp16 and _use_native_amp:
534
                        self.scaler.unscale_(self.optimizer)
535
536
                        torch.nn.utils.clip_grad_norm_(model.parameters(), self.args.max_grad_norm)
                    elif self.args.fp16 and _use_apex:
537
                        torch.nn.utils.clip_grad_norm_(amp.master_params(self.optimizer), self.args.max_grad_norm)
Julien Chaumond's avatar
Julien Chaumond committed
538
539
540
                    else:
                        torch.nn.utils.clip_grad_norm_(model.parameters(), self.args.max_grad_norm)

541
                    if is_torch_tpu_available():
542
                        xm.optimizer_step(self.optimizer)
543
                    if self.args.fp16 and _use_native_amp:
544
                        self.scaler.step(self.optimizer)
545
                        self.scaler.update()
Lysandre Debut's avatar
Lysandre Debut committed
546
                    else:
547
                        self.optimizer.step()
Lysandre Debut's avatar
Lysandre Debut committed
548

549
                    self.lr_scheduler.step()
Julien Chaumond's avatar
Julien Chaumond committed
550
                    model.zero_grad()
551
552
                    self.global_step += 1
                    self.epoch = epoch + (step + 1) / len(epoch_iterator)
Julien Chaumond's avatar
Julien Chaumond committed
553

554
555
556
557
558
                    if (self.args.logging_steps > 0 and self.global_step % self.args.logging_steps == 0) or (
                        self.global_step == 1 and self.args.logging_first_step
                    ):
                        logs: Dict[str, float] = {}
                        logs["loss"] = (tr_loss - logging_loss) / self.args.logging_steps
559
560
                        # backward compatibility for pytorch schedulers
                        logs["learning_rate"] = (
561
                            self.lr_scheduler.get_last_lr()[0]
562
                            if version.parse(torch.__version__) >= version.parse("1.4")
563
                            else self.lr_scheduler.get_lr()[0]
564
                        )
565
566
                        logging_loss = tr_loss

567
                        self.log(logs)
568

569
570
                    if self.args.evaluate_during_training and self.global_step % self.args.eval_steps == 0:
                        self.evaluate()
571

572
573
574
575
                    if self.args.save_steps > 0 and self.global_step % self.args.save_steps == 0:
                        # In all cases (even distributed/parallel), self.model is always a reference
                        # to the model we want to save.
                        if hasattr(model, "module"):
Teven's avatar
Teven committed
576
577
578
                            assert (
                                model.module is self.model
                            ), f"Module {model.module} should be a reference to self.model"
579
                        else:
Teven's avatar
Teven committed
580
                            assert model is self.model, f"Model {model} should be a reference to self.model"
581
582
583
584
585
                        # Save model checkpoint
                        output_dir = os.path.join(self.args.output_dir, f"{PREFIX_CHECKPOINT_DIR}-{self.global_step}")

                        self.save_model(output_dir)

586
                        if self.is_world_process_zero():
Julien Chaumond's avatar
Julien Chaumond committed
587
                            self._rotate_checkpoints()
588

589
                        if is_torch_tpu_available():
590
                            xm.rendezvous("saving_optimizer_states")
591
592
593
594
595
                            xm.save(self.optimizer.state_dict(), os.path.join(output_dir, "optimizer.pt"))
                            xm.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, "scheduler.pt"))
                        elif self.is_world_process_zero():
                            torch.save(self.optimizer.state_dict(), os.path.join(output_dir, "optimizer.pt"))
                            torch.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, "scheduler.pt"))
Julien Chaumond's avatar
Julien Chaumond committed
596

597
                if self.args.max_steps > 0 and self.global_step > self.args.max_steps:
Julien Chaumond's avatar
Julien Chaumond committed
598
599
                    epoch_iterator.close()
                    break
600
            if self.args.max_steps > 0 and self.global_step > self.args.max_steps:
Julien Chaumond's avatar
Julien Chaumond committed
601
602
                train_iterator.close()
                break
603
            if self.args.tpu_metrics_debug or self.args.debug:
Lysandre Debut's avatar
Lysandre Debut committed
604
605
                # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
                xm.master_print(met.metrics_report())
Julien Chaumond's avatar
Julien Chaumond committed
606
607
608

        if self.tb_writer:
            self.tb_writer.close()
609
610
611
        if self.args.past_index and hasattr(self, "_past"):
            # Clean the state at the end of training
            delattr(self, "_past")
Julien Chaumond's avatar
Julien Chaumond committed
612
613

        logger.info("\n\nTraining completed. Do not forget to share your model on huggingface.co/models =)\n\n")
614
615
        return TrainOutput(self.global_step, tr_loss / self.global_step)

616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
    def log(self, logs: Dict[str, float], iterator: Optional[tqdm] = None) -> None:
        """
        Log :obj:`logs` on the various objects watching training.

        Subclass and override this method to inject custom behavior.

        Args:
            logs (:obj:`Dict[str, float]`):
                The values to log.
            iterator (:obj:`tqdm`, `optional`):
                A potential tqdm progress bar to write the logs on.
        """
        if hasattr(self, "_log"):
            warnings.warn(
                "The `_log` method is deprecated and won't be called in a future version, define `log` in your subclass.",
                FutureWarning,
            )
            return self._log(logs, iterator=iterator)

635
636
        if self.epoch is not None:
            logs["epoch"] = self.epoch
637
638
639
        if self.global_step is None:
            # when logging evaluation metrics without training
            self.global_step = 0
640
641
        if self.tb_writer:
            for k, v in logs.items():
642
643
644
645
646
647
648
649
650
651
652
653
                if isinstance(v, (int, float)):
                    self.tb_writer.add_scalar(k, v, self.global_step)
                else:
                    logger.warning(
                        "Trainer is attempting to log a value of "
                        '"%s" of type %s for key "%s" as a scalar. '
                        "This invocation of Tensorboard's writer.add_scalar() "
                        "is incorrect so we dropped this attribute.",
                        v,
                        type(v),
                        k,
                    )
654
            self.tb_writer.flush()
655
        if is_wandb_available():
656
            if self.is_world_process_zero():
657
                wandb.log(logs, step=self.global_step)
658
        output = {**logs, **{"step": self.global_step}}
659
660
661
        if iterator is not None:
            iterator.write(output)
        else:
662
            print(output)
Julien Chaumond's avatar
Julien Chaumond committed
663

664
665
666
667
668
669
670
    def _prepare_inputs(
        self, inputs: Dict[str, Union[torch.Tensor, Any]], model: nn.Module
    ) -> Dict[str, Union[torch.Tensor, Any]]:
        """
        Prepare :obj:`inputs` before feeding them to the model, converting them to tensors if they are not already and
        handling potential state.
        """
Julien Chaumond's avatar
Julien Chaumond committed
671
        for k, v in inputs.items():
672
673
            if isinstance(v, torch.Tensor):
                inputs[k] = v.to(self.args.device)
Julien Chaumond's avatar
Julien Chaumond committed
674

675
676
        if self.args.past_index >= 0 and self._past is not None:
            inputs["mems"] = self._past
677

678
679
        return inputs

680
    def training_step(self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]]) -> float:
681
        """
682
        Perform a training step on a batch of inputs.
683
684
685
686
687
688
689
690
691
692
693
694
695

        Subclass and override to inject custom behavior.

        Args:
            model (:obj:`nn.Module`):
                The model to train.
            inputs (:obj:`Dict[str, Union[torch.Tensor, Any]]`):
                The inputs and targets of the model.

                The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
                argument :obj:`labels`. Check your model's documentation for all accepted arguments.

        Return:
696
            :obj:`float`: The training loss on this batch.
697
698
699
700
701
702
        """
        if hasattr(self, "_training_step"):
            warnings.warn(
                "The `_training_step` method is deprecated and won't be called in a future version, define `training_step` in your subclass.",
                FutureWarning,
            )
703
            return self._training_step(model, inputs, self.optimizer)
704
705
706

        model.train()
        inputs = self._prepare_inputs(inputs, model)
707

708
709
710
711
712
713
714
715
        if self.args.fp16 and _use_native_amp:
            with autocast():
                outputs = model(**inputs)
                loss = outputs[0]
        else:
            outputs = model(**inputs)
            # We don't use .loss here since the model may return tuples instead of ModelOutput.
            loss = outputs[0]
Julien Chaumond's avatar
Julien Chaumond committed
716

717
718
719
        if self.args.past_index >= 0:
            self._past = outputs[self.args.past_index]

Julien Chaumond's avatar
Julien Chaumond committed
720
721
        if self.args.n_gpu > 1:
            loss = loss.mean()  # mean() to average on multi-gpu parallel training
722

Julien Chaumond's avatar
Julien Chaumond committed
723
724
725
        if self.args.gradient_accumulation_steps > 1:
            loss = loss / self.args.gradient_accumulation_steps

726
727
728
        if self.args.fp16 and _use_native_amp:
            self.scaler.scale(loss).backward()
        elif self.args.fp16 and _use_apex:
729
            with amp.scale_loss(loss, self.optimizer) as scaled_loss:
Julien Chaumond's avatar
Julien Chaumond committed
730
731
732
733
734
735
                scaled_loss.backward()
        else:
            loss.backward()

        return loss.item()

Lysandre Debut's avatar
Lysandre Debut committed
736
    def is_local_master(self) -> bool:
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
        """
        Whether or not this process is the local (e.g., on one machine if training in a distributed fashion on
        several machines) main process.

        .. warning::

            This method is deprecated, use :meth:`~transformers.Trainer.is_local_process_zero` instead.
        """
        warnings.warn("This method is deprecated, use `Trainer.is_local_process_zero()` instead.", FutureWarning)
        return self.is_local_process_zero()

    def is_local_process_zero(self) -> bool:
        """
        Whether or not this process is the local (e.g., on one machine if training in a distributed fashion on
        several machines) main process.
        """
753
        if is_torch_tpu_available():
Lysandre Debut's avatar
Lysandre Debut committed
754
755
756
757
            return xm.is_master_ordinal(local=True)
        else:
            return self.args.local_rank in [-1, 0]

Julien Chaumond's avatar
Julien Chaumond committed
758
759
    def is_world_master(self) -> bool:
        """
760
761
762
763
764
765
766
767
768
769
770
771
772
773
        Whether or not this process is the global main process (when training in a distributed fashion on
        several machines, this is only going to be :obj:`True` for one process).

        .. warning::

            This method is deprecated, use :meth:`~transformers.Trainer.is_world_process_zero` instead.
        """
        warnings.warn("This method is deprecated, use `Trainer.is_world_process_zero()` instead.", FutureWarning)
        return self.is_world_process_zero()

    def is_world_process_zero(self) -> bool:
        """
        Whether or not this process is the global main process (when training in a distributed fashion on
        several machines, this is only going to be :obj:`True` for one process).
Julien Chaumond's avatar
Julien Chaumond committed
774
        """
775
        if is_torch_tpu_available():
Lysandre Debut's avatar
Lysandre Debut committed
776
777
778
            return xm.is_master_ordinal(local=False)
        else:
            return self.args.local_rank == -1 or torch.distributed.get_rank() == 0
Julien Chaumond's avatar
Julien Chaumond committed
779
780
781

    def save_model(self, output_dir: Optional[str] = None):
        """
782
        Will save the model, so you can reload it using :obj:`from_pretrained()`.
Julien Chaumond's avatar
Julien Chaumond committed
783

784
        Will only save from the world_master process (unless in TPUs).
Julien Chaumond's avatar
Julien Chaumond committed
785
        """
786

787
        if is_torch_tpu_available():
788
            self._save_tpu(output_dir)
789
        elif self.is_world_process_zero():
Julien Chaumond's avatar
Julien Chaumond committed
790
791
            self._save(output_dir)

792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
    def _save_tpu(self, output_dir: Optional[str] = None):
        output_dir = output_dir if output_dir is not None else self.args.output_dir
        logger.info("Saving model checkpoint to %s", output_dir)

        if xm.is_master_ordinal():
            os.makedirs(output_dir, exist_ok=True)
            torch.save(self.args, os.path.join(output_dir, "training_args.bin"))

        # Save a trained model and configuration using `save_pretrained()`.
        # They can then be reloaded using `from_pretrained()`
        if not isinstance(self.model, PreTrainedModel):
            raise ValueError("Trainer.model appears to not be a PreTrainedModel")

        xm.rendezvous("saving_checkpoint")
        self.model.save_pretrained(output_dir)

Julien Chaumond's avatar
Julien Chaumond committed
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
    def _save(self, output_dir: Optional[str] = None):
        output_dir = output_dir if output_dir is not None else self.args.output_dir
        os.makedirs(output_dir, exist_ok=True)
        logger.info("Saving model checkpoint to %s", output_dir)
        # Save a trained model and configuration using `save_pretrained()`.
        # They can then be reloaded using `from_pretrained()`
        if not isinstance(self.model, PreTrainedModel):
            raise ValueError("Trainer.model appears to not be a PreTrainedModel")
        self.model.save_pretrained(output_dir)

        # Good practice: save your training arguments together with the trained model
        torch.save(self.args, os.path.join(output_dir, "training_args.bin"))

    def _sorted_checkpoints(self, checkpoint_prefix=PREFIX_CHECKPOINT_DIR, use_mtime=False) -> List[str]:
        ordering_and_checkpoint_path = []

824
        glob_checkpoints = [str(x) for x in Path(self.args.output_dir).glob(f"{checkpoint_prefix}-*")]
Julien Chaumond's avatar
Julien Chaumond committed
825
826
827
828
829

        for path in glob_checkpoints:
            if use_mtime:
                ordering_and_checkpoint_path.append((os.path.getmtime(path), path))
            else:
830
                regex_match = re.match(f".*{checkpoint_prefix}-([0-9]+)", path)
Julien Chaumond's avatar
Julien Chaumond committed
831
832
833
834
835
836
837
838
                if regex_match and regex_match.groups():
                    ordering_and_checkpoint_path.append((int(regex_match.groups()[0]), path))

        checkpoints_sorted = sorted(ordering_and_checkpoint_path)
        checkpoints_sorted = [checkpoint[1] for checkpoint in checkpoints_sorted]
        return checkpoints_sorted

    def _rotate_checkpoints(self, use_mtime=False) -> None:
839
        if self.args.save_total_limit is None or self.args.save_total_limit <= 0:
Julien Chaumond's avatar
Julien Chaumond committed
840
841
842
843
844
845
846
847
848
849
850
851
852
            return

        # Check if we should delete older checkpoint(s)
        checkpoints_sorted = self._sorted_checkpoints(use_mtime=use_mtime)
        if len(checkpoints_sorted) <= self.args.save_total_limit:
            return

        number_of_checkpoints_to_delete = max(0, len(checkpoints_sorted) - self.args.save_total_limit)
        checkpoints_to_be_deleted = checkpoints_sorted[:number_of_checkpoints_to_delete]
        for checkpoint in checkpoints_to_be_deleted:
            logger.info("Deleting older checkpoint [{}] due to args.save_total_limit".format(checkpoint))
            shutil.rmtree(checkpoint)

853
    def evaluate(self, eval_dataset: Optional[Dataset] = None) -> Dict[str, float]:
Julien Chaumond's avatar
Julien Chaumond committed
854
        """
855
        Run evaluation and returns metrics.
Julien Chaumond's avatar
Julien Chaumond committed
856
857

        The calling script will be responsible for providing a method to compute metrics, as they are
858
        task-dependent (pass it to the init :obj:`compute_metrics` argument).
Julien Chaumond's avatar
Julien Chaumond committed
859

860
861
        You can also subclass and override this method to inject custom behavior.

Julien Chaumond's avatar
Julien Chaumond committed
862
        Args:
863
864
            eval_dataset (:obj:`Dataset`, `optional`):
                Pass a dataset if you wish to override :obj:`self.eval_dataset`.
865

Julien Chaumond's avatar
Julien Chaumond committed
866
        Returns:
867
            A dictionary containing the evaluation loss and the potential metrics computed from the predictions.
Julien Chaumond's avatar
Julien Chaumond committed
868
869
870
        """
        eval_dataloader = self.get_eval_dataloader(eval_dataset)

871
        output = self.prediction_loop(eval_dataloader, description="Evaluation")
Lysandre Debut's avatar
Lysandre Debut committed
872

873
        self.log(output.metrics)
874

875
        if self.args.tpu_metrics_debug or self.args.debug:
Lysandre Debut's avatar
Lysandre Debut committed
876
877
878
            # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
            xm.master_print(met.metrics_report())

Julien Chaumond's avatar
Julien Chaumond committed
879
880
881
882
        return output.metrics

    def predict(self, test_dataset: Dataset) -> PredictionOutput:
        """
883
        Run prediction and returns predictions and potential metrics.
Julien Chaumond's avatar
Julien Chaumond committed
884
885

        Depending on the dataset and your use case, your test dataset may contain labels.
886
887
888
889
890
        In that case, this method will also return metrics, like in :obj:`evaluate()`.

        Args:
            test_dataset (:obj:`Dataset`):
                Dataset to run the predictions on.
891

892
893
894
895
896
897
898
899
        Returns:
            `NamedTuple`:
            predictions (:obj:`np.ndarray`):
                The predictions on :obj:`test_dataset`.
            label_ids (:obj:`np.ndarray`, `optional`):
                The labels (if the dataset contained some).
            metrics (:obj:`Dict[str, float]`, `optional`):
                The potential dictionary of metrics (if the dataset contained labels).
Julien Chaumond's avatar
Julien Chaumond committed
900
901
        """
        test_dataloader = self.get_test_dataloader(test_dataset)
902

903
        return self.prediction_loop(test_dataloader, description="Prediction")
Julien Chaumond's avatar
Julien Chaumond committed
904

905
    def prediction_loop(
Julien Chaumond's avatar
Julien Chaumond committed
906
907
908
        self, dataloader: DataLoader, description: str, prediction_loss_only: Optional[bool] = None
    ) -> PredictionOutput:
        """
909
        Prediction/evaluation loop, shared by :obj:`Trainer.evaluate()` and :obj:`Trainer.predict()`.
Julien Chaumond's avatar
Julien Chaumond committed
910
911
912

        Works both with or without labels.
        """
913
914
915
916
917
918
        if hasattr(self, "_prediction_loop"):
            warnings.warn(
                "The `_prediction_loop` method is deprecated and won't be called in a future version, define `prediction_loop` in your subclass.",
                FutureWarning,
            )
            return self._prediction_loop(dataloader, description, prediction_loss_only=prediction_loss_only)
Julien Chaumond's avatar
Julien Chaumond committed
919
920
921

        prediction_loss_only = prediction_loss_only if prediction_loss_only is not None else self.prediction_loss_only

922
        model = self.model
Julien Chaumond's avatar
Julien Chaumond committed
923
        # multi-gpu eval
924
925
        if self.args.n_gpu > 1:
            model = torch.nn.DataParallel(model)
Julien Chaumond's avatar
Julien Chaumond committed
926
927
        else:
            model = self.model
928
929
        # Note: in torch.distributed mode, there's no point in wrapping the model
        # inside a DistributedDataParallel as we'll be under `no_grad` anyways.
Julien Chaumond's avatar
Julien Chaumond committed
930

931
        batch_size = dataloader.batch_size
Julien Chaumond's avatar
Julien Chaumond committed
932
        logger.info("***** Running %s *****", description)
933
934
        logger.info("  Num examples = %d", self.num_examples(dataloader))
        logger.info("  Batch size = %d", batch_size)
Julien Chaumond's avatar
Julien Chaumond committed
935
        eval_losses: List[float] = []
936
937
        preds: torch.Tensor = None
        label_ids: torch.Tensor = None
Julien Chaumond's avatar
Julien Chaumond committed
938
939
        model.eval()

940
        if is_torch_tpu_available():
941
942
            dataloader = pl.ParallelLoader(dataloader, [self.args.device]).per_device_loader(self.args.device)

943
        if self.args.past_index >= 0:
944
            self._past = None
945

Julien Chaumond's avatar
Julien Chaumond committed
946
        for inputs in tqdm(dataloader, desc=description):
947
948
949
950
951
952
953
            loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only)
            if loss is not None:
                eval_losses.append(loss)
            if logits is not None:
                preds = logits if preds is None else torch.cat((preds, logits), dim=0)
            if labels is not None:
                label_ids = labels if label_ids is None else torch.cat((label_ids, labels), dim=0)
Julien Chaumond's avatar
Julien Chaumond committed
954

955
956
957
        if self.args.past_index and hasattr(self, "_past"):
            # Clean the state at the end of the evaluation loop
            delattr(self, "_past")
Julien Chaumond's avatar
Julien Chaumond committed
958

959
960
961
962
963
964
        if self.args.local_rank != -1:
            # In distributed mode, concatenate all results from all nodes:
            if preds is not None:
                preds = self.distributed_concat(preds, num_total_examples=self.num_examples(dataloader))
            if label_ids is not None:
                label_ids = self.distributed_concat(label_ids, num_total_examples=self.num_examples(dataloader))
965
        elif is_torch_tpu_available():
Lysandre Debut's avatar
Lysandre Debut committed
966
            # tpu-comment: Get all predictions and labels from all worker shards of eval dataset
967
968
969
970
971
972
973
974
975
976
            if preds is not None:
                preds = xm.mesh_reduce("eval_preds", preds, torch.cat)
            if label_ids is not None:
                label_ids = xm.mesh_reduce("eval_label_ids", label_ids, torch.cat)

        # Finally, turn the aggregated tensors into numpy arrays.
        if preds is not None:
            preds = preds.cpu().numpy()
        if label_ids is not None:
            label_ids = label_ids.cpu().numpy()
Lysandre Debut's avatar
Lysandre Debut committed
977

Julien Chaumond's avatar
Julien Chaumond committed
978
979
980
981
982
        if self.compute_metrics is not None and preds is not None and label_ids is not None:
            metrics = self.compute_metrics(EvalPrediction(predictions=preds, label_ids=label_ids))
        else:
            metrics = {}
        if len(eval_losses) > 0:
983
984
985
986
987
988
            metrics["eval_loss"] = np.mean(eval_losses)

        # Prefix all keys with eval_
        for key in list(metrics.keys()):
            if not key.startswith("eval_"):
                metrics[f"eval_{key}"] = metrics.pop(key)
Julien Chaumond's avatar
Julien Chaumond committed
989
990

        return PredictionOutput(predictions=preds, label_ids=label_ids, metrics=metrics)
991
992
993
994
995
996
997
998
999
1000
1001
1002

    def distributed_concat(self, tensor: torch.Tensor, num_total_examples: int) -> torch.Tensor:
        assert self.args.local_rank != -1

        output_tensors = [tensor.clone() for _ in range(torch.distributed.get_world_size())]
        torch.distributed.all_gather(output_tensors, tensor)

        concat = torch.cat(output_tensors, dim=0)

        # truncate the dummy elements added by SequentialDistributedSampler
        output = concat[:num_total_examples]
        return output
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048

    def prediction_step(
        self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]], prediction_loss_only: bool
    ) -> Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
        """
        Perform an evaluation step on :obj:`model` using obj:`inputs`.

        Subclass and override to inject custom behavior.

        Args:
            model (:obj:`nn.Module`):
                The model to evaluate.
            inputs (:obj:`Dict[str, Union[torch.Tensor, Any]]`):
                The inputs and targets of the model.

                The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
                argument :obj:`labels`. Check your model's documentation for all accepted arguments.
            prediction_loss_only (:obj:`bool`):
                Whether or not to return the loss only.

        Return:
            Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
            A tuple with the loss, logits and labels (each being optional).
        """
        has_labels = any(inputs.get(k) is not None for k in ["labels", "lm_labels", "masked_lm_labels"])

        inputs = self._prepare_inputs(inputs, model)

        with torch.no_grad():
            outputs = model(**inputs)
            if has_labels:
                loss, logits = outputs[:2]
                loss = loss.mean().item()
            else:
                loss = None
                logits = outputs[0]
            if self.args.past_index >= 0:
                self._past = outputs[self.args.past_index if has_labels else self.args.past_index - 1]

        if prediction_loss_only:
            return (loss, None, None)

        labels = inputs.get("labels")
        if labels is not None:
            labels = labels.detach()
        return (loss, logits.detach(), labels)