"tests/test_feature_extraction_utils.py" did not exist on "6a5272b2054622e27b4843e0c9722cc716d5458d"
trainer.py 115 KB
Newer Older
Sylvain Gugger's avatar
Sylvain Gugger committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
# coding=utf-8
# Copyright 2020-present 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.
"""
The Trainer class, to easily train a 馃 Transformers from scratch or finetune it on a new task.
"""

19
import collections
20
import inspect
21
import math
Julien Chaumond's avatar
Julien Chaumond committed
22
23
24
import os
import re
import shutil
25
import sys
Sylvain Gugger's avatar
Sylvain Gugger committed
26
import tempfile
27
import time
28
import warnings
29
from logging import StreamHandler
Julien Chaumond's avatar
Julien Chaumond committed
30
from pathlib import Path
31
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
Julien Chaumond's avatar
Julien Chaumond committed
32

33
34
from tqdm.auto import tqdm

Julien Chaumond's avatar
Julien Chaumond committed
35

36
37
# Integrations must be imported before ML frameworks:
from .integrations import (  # isort: split
38
    default_hp_search_backend,
39
    get_reporting_integration_callbacks,
40
    hp_params,
41
    is_fairscale_available,
42
    is_optuna_available,
43
    is_ray_tune_available,
44
45
    run_hp_search_optuna,
    run_hp_search_ray,
46
47
    deepspeed_init,
    is_deepspeed_zero3_enabled,
48
)
49
50
51
52
53
54

import numpy as np
import torch
from packaging import version
from torch import nn
from torch.utils.data.dataloader import DataLoader
55
from torch.utils.data.dataset import Dataset, IterableDataset
56
57
58
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data.sampler import RandomSampler, SequentialSampler

59
60
from . import __version__
from .configuration_utils import PretrainedConfig
61
from .data.data_collator import DataCollator, DataCollatorWithPadding, default_data_collator
62
from .dependency_versions_check import dep_version_check
Sylvain Gugger's avatar
Sylvain Gugger committed
63
from .file_utils import (
64
    CONFIG_NAME,
Sylvain Gugger's avatar
Sylvain Gugger committed
65
    WEIGHTS_NAME,
Sylvain Gugger's avatar
Sylvain Gugger committed
66
    PushToHubMixin,
Sylvain Gugger's avatar
Sylvain Gugger committed
67
68
69
    is_apex_available,
    is_datasets_available,
    is_in_notebook,
Sylvain Gugger's avatar
Sylvain Gugger committed
70
71
    is_sagemaker_dp_enabled,
    is_sagemaker_mp_enabled,
Sylvain Gugger's avatar
Sylvain Gugger committed
72
    is_torch_tpu_available,
73
    is_training_run_on_sagemaker,
Sylvain Gugger's avatar
Sylvain Gugger committed
74
)
75
from .modeling_utils import PreTrainedModel, unwrap_model
Sylvain Gugger's avatar
Sylvain Gugger committed
76
from .optimization import Adafactor, AdamW, get_scheduler
77
from .tokenization_utils_base import PreTrainedTokenizerBase
Sylvain Gugger's avatar
Sylvain Gugger committed
78
79
80
81
82
83
84
85
86
87
from .trainer_callback import (
    CallbackHandler,
    DefaultFlowCallback,
    PrinterCallback,
    ProgressCallback,
    TrainerCallback,
    TrainerControl,
    TrainerState,
)
from .trainer_pt_utils import (
88
    DistributedLengthGroupedSampler,
89
    DistributedSamplerWithLoop,
90
    DistributedTensorGatherer,
91
    IterableDatasetShard,
Sylvain Gugger's avatar
Sylvain Gugger committed
92
    LabelSmoother,
93
    LengthGroupedSampler,
Sylvain Gugger's avatar
Sylvain Gugger committed
94
    SequentialDistributedSampler,
95
    ShardSampler,
Sylvain Gugger's avatar
Sylvain Gugger committed
96
97
    distributed_broadcast_scalars,
    distributed_concat,
98
    find_batch_size,
99
    get_parameter_names,
Sylvain Gugger's avatar
Sylvain Gugger committed
100
101
102
    nested_concat,
    nested_detach,
    nested_numpify,
103
    nested_truncate,
Sylvain Gugger's avatar
Sylvain Gugger committed
104
105
106
    nested_xla_mesh_reduce,
    reissue_pt_warnings,
)
107
108
109
from .trainer_utils import (
    PREFIX_CHECKPOINT_DIR,
    BestRun,
110
    EvalLoopOutput,
111
112
113
    EvalPrediction,
    HPSearchBackend,
    PredictionOutput,
114
    ShardedDDPOption,
115
    TrainerMemoryTracker,
116
117
118
    TrainOutput,
    default_compute_objective,
    default_hp_space,
119
    denumpify_detensorize,
120
    get_last_checkpoint,
121
    set_seed,
122
    speed_metrics,
123
)
124
from .training_args import ParallelMode, TrainingArguments
Lysandre Debut's avatar
Lysandre Debut committed
125
from .utils import logging
126
from .utils.modeling_auto_mapping import MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES
Julien Chaumond's avatar
Julien Chaumond committed
127
128


129
_is_native_amp_available = False
130

Sylvain Gugger's avatar
Sylvain Gugger committed
131
DEFAULT_CALLBACKS = [DefaultFlowCallback]
132
DEFAULT_PROGRESS_CALLBACK = ProgressCallback
Sylvain Gugger's avatar
Sylvain Gugger committed
133

134
135
136
137
if is_in_notebook():
    from .utils.notebook import NotebookProgressCallback

    DEFAULT_PROGRESS_CALLBACK = NotebookProgressCallback
138

139
140
if is_apex_available():
    from apex import amp
141

142
if version.parse(torch.__version__) >= version.parse("1.6"):
143
    _is_native_amp_available = True
144
    from torch.cuda.amp import autocast
Julien Chaumond's avatar
Julien Chaumond committed
145

146
147
if is_datasets_available():
    import datasets
Julien Chaumond's avatar
Julien Chaumond committed
148

149
if is_torch_tpu_available():
Lysandre Debut's avatar
Lysandre Debut committed
150
151
152
153
    import torch_xla.core.xla_model as xm
    import torch_xla.debug.metrics as met
    import torch_xla.distributed.parallel_loader as pl

154
if is_fairscale_available():
155
    dep_version_check("fairscale")
156
    import fairscale
157
    from fairscale.nn.data_parallel import FullyShardedDataParallel as FullyShardedDDP
158
    from fairscale.nn.data_parallel import ShardedDataParallel as ShardedDDP
159
    from fairscale.nn.wrap import auto_wrap
160
161
162
    from fairscale.optim import OSS
    from fairscale.optim.grad_scaler import ShardedGradScaler

Sylvain Gugger's avatar
Sylvain Gugger committed
163
if is_sagemaker_dp_enabled():
Sylvain Gugger's avatar
Sylvain Gugger committed
164
165
166
167
    import smdistributed.dataparallel.torch.distributed as dist
    from smdistributed.dataparallel.torch.parallel.distributed import DistributedDataParallel as DDP
else:
    import torch.distributed as dist
168

Sylvain Gugger's avatar
Sylvain Gugger committed
169
170
171
172
173
if is_sagemaker_mp_enabled():
    import smdistributed.modelparallel.torch as smp

    from .trainer_pt_utils import smp_forward_backward, smp_forward_only, smp_gather, smp_nested_concat

174
175
176
177
if is_training_run_on_sagemaker():
    logging.add_handler(StreamHandler(sys.stdout))


178
179
180
if TYPE_CHECKING:
    import optuna

Lysandre Debut's avatar
Lysandre Debut committed
181
logger = logging.get_logger(__name__)
Julien Chaumond's avatar
Julien Chaumond committed
182
183
184
185


class Trainer:
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
186
    Trainer is a simple but feature-complete training and eval loop for PyTorch, optimized for 馃 Transformers.
187
188

    Args:
Sylvain Gugger's avatar
Sylvain Gugger committed
189
        model (:class:`~transformers.PreTrainedModel` or :obj:`torch.nn.Module`, `optional`):
190
            The model to train, evaluate or use for predictions. If not provided, a ``model_init`` must be passed.
Sylvain Gugger's avatar
Sylvain Gugger committed
191
192
193
194
195
196

            .. note::

                :class:`~transformers.Trainer` is optimized to work with the :class:`~transformers.PreTrainedModel`
                provided by the library. You can still use your own models defined as :obj:`torch.nn.Module` as long as
                they work the same way as the 馃 Transformers models.
197
        args (:class:`~transformers.TrainingArguments`, `optional`):
Sylvain Gugger's avatar
Sylvain Gugger committed
198
199
200
            The arguments to tweak for training. Will default to a basic instance of
            :class:`~transformers.TrainingArguments` with the ``output_dir`` set to a directory named `tmp_trainer` in
            the current directory if not provided.
201
        data_collator (:obj:`DataCollator`, `optional`):
Sylvain Gugger's avatar
Sylvain Gugger committed
202
203
204
            The function to use to form a batch from a list of elements of :obj:`train_dataset` or :obj:`eval_dataset`.
            Will default to :func:`~transformers.default_data_collator` if no ``tokenizer`` is provided, an instance of
            :func:`~transformers.DataCollatorWithPadding` otherwise.
205
        train_dataset (:obj:`torch.utils.data.dataset.Dataset` or :obj:`torch.utils.data.dataset.IterableDataset`, `optional`):
206
            The dataset to use for training. If it is an :obj:`datasets.Dataset`, columns not accepted by the
207
            ``model.forward()`` method are automatically removed.
208
209
210
211
212
213

            Note that if it's a :obj:`torch.utils.data.dataset.IterableDataset` with some randomization and you are
            training in a distributed fashion, your iterable dataset should either use a internal attribute
            :obj:`generator` that is a :obj:`torch.Generator` for the randomization that must be identic on all
            processes (and the Trainer will manually set the seed of this :obj:`generator` at each epoch) or have a
            :obj:`set_epoch()` method that internally sets the seed of the RNGs used.
Sylvain Gugger's avatar
Sylvain Gugger committed
214
        eval_dataset (:obj:`torch.utils.data.dataset.Dataset`, `optional`):
215
             The dataset to use for evaluation. If it is an :obj:`datasets.Dataset`, columns not accepted by the
Sylvain Gugger's avatar
Sylvain Gugger committed
216
             ``model.forward()`` method are automatically removed.
217
218
219
220
        tokenizer (:class:`PreTrainedTokenizerBase`, `optional`):
            The tokenizer used to preprocess the data. If provided, will be used to automatically pad the inputs the
            maximum length when batching inputs, and it will be saved along the model to make it easier to rerun an
            interrupted training or reuse the fine-tuned model.
221
222
223
        model_init (:obj:`Callable[[], PreTrainedModel]`, `optional`):
            A function that instantiates the model to be used. If provided, each call to
            :meth:`~transformers.Trainer.train` will start from a new instance of the model as given by this function.
224

Sylvain Gugger's avatar
Sylvain Gugger committed
225
226
            The function may have zero argument, or a single one containing the optuna/Ray Tune trial object, to be
            able to choose different architectures according to hyper parameters (such as layer count, sizes of inner
Sylvain Gugger's avatar
Sylvain Gugger committed
227
228
            layers, dropout probabilities etc).
        compute_metrics (:obj:`Callable[[EvalPrediction], Dict]`, `optional`):
229
            The function that will be used to compute metrics at evaluation. Must take a
Sylvain Gugger's avatar
Sylvain Gugger committed
230
231
232
233
            :class:`~transformers.EvalPrediction` and return a dictionary string to metric values.
        callbacks (List of :obj:`~transformers.TrainerCallback`, `optional`):
            A list of callbacks to customize the training loop. Will add those to the list of default callbacks
            detailed in :doc:`here <callback>`.
Sylvain Gugger's avatar
Sylvain Gugger committed
234
235

            If you want to remove one of the default callbacks used, use the :meth:`Trainer.remove_callback` method.
Sylvain Gugger's avatar
Sylvain Gugger committed
236
        optimizers (:obj:`Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR`, `optional`): A tuple
Sylvain Gugger's avatar
Sylvain Gugger committed
237
            containing the optimizer and the scheduler to use. Will default to an instance of
238
            :class:`~transformers.AdamW` on your model and a scheduler given by
Sylvain Gugger's avatar
Sylvain Gugger committed
239
            :func:`~transformers.get_linear_schedule_with_warmup` controlled by :obj:`args`.
240

241
242
243
244
245
246
247
248
249
250
    Important attributes:

        - **model** -- Always points to the core model. If using a transformers model, it will be a
          :class:`~transformers.PreTrainedModel` subclass.
        - **model_wrapped** -- Always points to the most external model in case one or more other modules wrap the
          original model. This is the model that should be used for the forward pass. For example, under ``DeepSpeed``,
          the inner model is wrapped in ``DeepSpeed`` and then again in ``torch.nn.DistributedDataParallel``. If the
          inner model hasn't been wrapped, then ``self.model_wrapped`` is the same as ``self.model``.
        - **is_model_parallel** -- Whether or not a model has been switched to a model parallel mode (different from
          data parallelism, this means some of the model layers are split on different GPUs).
251
252
253
        - **place_model_on_device** -- Whether or not to automatically place the model on the device - it will be set
          to :obj:`False` if model parallel or deepspeed is used, or if the default
          ``TrainingArguments.place_model_on_device`` is overridden to return :obj:`False` .
254
255
        - **is_in_train** -- Whether or not a model is currently running ``train`` (e.g. when ``evaluate`` is called
          while in ``train``)
256

Julien Chaumond's avatar
Julien Chaumond committed
257
258
    """

259
    from .trainer_pt_utils import _get_learning_rate, log_metrics, metrics_format, save_metrics, save_state
260

Julien Chaumond's avatar
Julien Chaumond committed
261
262
    def __init__(
        self,
Sylvain Gugger's avatar
Sylvain Gugger committed
263
        model: Union[PreTrainedModel, torch.nn.Module] = None,
264
        args: TrainingArguments = None,
Julien Chaumond's avatar
Julien Chaumond committed
265
266
267
        data_collator: Optional[DataCollator] = None,
        train_dataset: Optional[Dataset] = None,
        eval_dataset: Optional[Dataset] = None,
268
        tokenizer: Optional[PreTrainedTokenizerBase] = None,
269
        model_init: Callable[[], PreTrainedModel] = None,
Julien Chaumond's avatar
Julien Chaumond committed
270
        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
Sylvain Gugger's avatar
Sylvain Gugger committed
271
        callbacks: Optional[List[TrainerCallback]] = None,
272
        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
Julien Chaumond's avatar
Julien Chaumond committed
273
    ):
Sylvain Gugger's avatar
Sylvain Gugger committed
274
        if args is None:
275
276
277
            output_dir = "tmp_trainer"
            logger.info(f"No `TrainingArguments` passed, using `output_dir={output_dir}`.")
            args = TrainingArguments(output_dir=output_dir)
Sylvain Gugger's avatar
Sylvain Gugger committed
278
279
280
        self.args = args
        # Seed must be set before instantiating the model when using model
        set_seed(self.args.seed)
281
        self.hp_name = None
282
        self.deepspeed = None
283
        self.is_in_train = False
284

285
286
287
288
        # memory metrics - must set up as early as possible
        self._memory_tracker = TrainerMemoryTracker(self.args.skip_memory_metrics)
        self._memory_tracker.start()

289
290
291
        # force device and distributed setup init explicitly
        args._setup_devices

292
293
294
295
296
297
298
299
300
301
302
303
304
305
        if model is None:
            if model_init is not None:
                self.model_init = model_init
                model = self.call_model_init()
            else:
                raise RuntimeError("`Trainer` requires either a `model` or `model_init` argument")
        else:
            if model_init is not None:
                warnings.warn(
                    "`Trainer` requires either a `model` or `model_init` argument, but not both. "
                    "`model_init` will overwrite your model when calling the `train` method. This will become a fatal error in the next release.",
                    FutureWarning,
                )
            self.model_init = model_init
306

307
308
309
310
311
        if hasattr(model, "is_parallelizable") and model.is_parallelizable and model.model_parallel:
            self.is_model_parallel = True
        else:
            self.is_model_parallel = False

312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
        # Setup Sharded DDP training
        self.sharded_ddp = None
        if len(args.sharded_ddp) > 0:
            if args.deepspeed:
                raise ValueError(
                    "Using --sharded_ddp xxx together with --deepspeed is not possible, deactivate one of those flags."
                )

            if args.local_rank == -1:
                raise ValueError("Using sharded DDP only works in distributed training.")
            elif not is_fairscale_available():
                raise ImportError("Sharded DDP training requires fairscale: `pip install fairscale`.")
            elif ShardedDDPOption.SIMPLE not in args.sharded_ddp and FullyShardedDDP is None:
                raise ImportError(
                    "Sharded DDP in a mode other than simple training requires fairscale version >= 0.3, found "
                    f"{fairscale.__version__}. Upgrade your fairscale library: `pip install --upgrade fairscale`."
                )
            elif ShardedDDPOption.SIMPLE in args.sharded_ddp:
                self.sharded_ddp = ShardedDDPOption.SIMPLE
            elif ShardedDDPOption.ZERO_DP_2 in args.sharded_ddp:
                self.sharded_ddp = ShardedDDPOption.ZERO_DP_2
            elif ShardedDDPOption.ZERO_DP_3 in args.sharded_ddp:
                self.sharded_ddp = ShardedDDPOption.ZERO_DP_3

336
        # one place to sort out whether to place the model on device or not
337
338
339
340
341
342
        # postpone switching model to cuda when:
        # 1. MP - since we are trying to fit a much bigger than 1 gpu model
        # 2. fp16-enabled DeepSpeed loads the model in half the size and it doesn't need .to() anyway,
        #    and we only use deepspeed for training at the moment
        # 3. full fp16 eval - since the model needs to be half'ed first
        # 4. Sharded DDP - same as MP
343
        self.place_model_on_device = args.place_model_on_device
344
345
        if (
            self.is_model_parallel
346
            or args.deepspeed
347
348
349
            or (args.fp16_full_eval and not args.do_train)
            or (self.sharded_ddp in [ShardedDDPOption.ZERO_DP_2, ShardedDDPOption.ZERO_DP_3])
        ):
350
351
            self.place_model_on_device = False

352
353
        default_collator = default_data_collator if tokenizer is None else DataCollatorWithPadding(tokenizer)
        self.data_collator = data_collator if data_collator is not None else default_collator
Julien Chaumond's avatar
Julien Chaumond committed
354
355
        self.train_dataset = train_dataset
        self.eval_dataset = eval_dataset
356
        self.tokenizer = tokenizer
357

358
        if self.place_model_on_device:
359
            model = model.to(args.device)
Stas Bekman's avatar
Stas Bekman committed
360
361
362

        # Force n_gpu to 1 to avoid DataParallel as MP will manage the GPUs
        if self.is_model_parallel:
363
            self.args._n_gpu = 1
364
365
366
367
368

        # later use `self.model is self.model_wrapped` to check if it's wrapped or not
        self.model_wrapped = model
        self.model = model

Julien Chaumond's avatar
Julien Chaumond committed
369
        self.compute_metrics = compute_metrics
370
        self.optimizer, self.lr_scheduler = optimizers
Sylvain Gugger's avatar
Sylvain Gugger committed
371
372
373
374
375
        if model_init is not None and (self.optimizer is not None or self.lr_scheduler is not None):
            raise RuntimeError(
                "Passing a `model_init` is incompatible with providing the `optimizers` argument."
                "You should subclass `Trainer` and override the `create_optimizer_and_scheduler` method."
            )
376
377
        default_callbacks = DEFAULT_CALLBACKS + get_reporting_integration_callbacks(self.args.report_to)
        callbacks = default_callbacks if callbacks is None else default_callbacks + callbacks
378
379
380
        self.callback_handler = CallbackHandler(
            callbacks, self.model, self.tokenizer, self.optimizer, self.lr_scheduler
        )
381
        self.add_callback(PrinterCallback if self.args.disable_tqdm else DEFAULT_PROGRESS_CALLBACK)
Sylvain Gugger's avatar
Sylvain Gugger committed
382

383
384
385
        # Will be set to True by `self._setup_loggers()` on first call to `self.log()`.
        self._loggers_initialized = False

Julien Chaumond's avatar
Julien Chaumond committed
386
        # Create output directory if needed
387
        if self.is_world_process_zero():
Julien Chaumond's avatar
Julien Chaumond committed
388
            os.makedirs(self.args.output_dir, exist_ok=True)
389
        if not callable(self.data_collator) and callable(getattr(self.data_collator, "collate_batch", None)):
Sylvain Gugger's avatar
Sylvain Gugger committed
390
            raise ValueError("The `data_collator` should be a simple callable (function, class with `__call__`).")
391

392
393
394
395
396
397
        if args.max_steps > 0:
            logger.info("max_steps is given, it will override any value given in num_train_epochs")

        if train_dataset is not None and not isinstance(train_dataset, collections.abc.Sized) and args.max_steps <= 0:
            raise ValueError("train_dataset does not implement __len__, max_steps has to be specified")

398
        self._signature_columns = None
399

400
401
402
        # Mixed precision setup
        self.use_apex = False
        self.use_amp = False
403
404
        self.fp16_backend = None

405
406
        if args.fp16:
            if args.fp16_backend == "auto":
407
                self.fp16_backend = "amp" if _is_native_amp_available else "apex"
408
            else:
409
410
                self.fp16_backend = args.fp16_backend
            logger.info(f"Using {self.fp16_backend} fp16 backend")
411

412
413
        if args.fp16 and not args.deepspeed:  # deepspeed manages its own fp16
            if self.fp16_backend == "amp":
414
                self.use_amp = True
415
                self.scaler = ShardedGradScaler() if self.sharded_ddp is not None else torch.cuda.amp.GradScaler()
416
417
418
419
420
421
422
            else:
                if not is_apex_available():
                    raise ImportError(
                        "Using FP16 with APEX but APEX is not installed, please refer to https://www.github.com/nvidia/apex."
                    )
                self.use_apex = True

Sylvain Gugger's avatar
Sylvain Gugger committed
423
424
425
426
427
428
        # Label smoothing
        if self.args.label_smoothing_factor != 0:
            self.label_smoother = LabelSmoother(epsilon=self.args.label_smoothing_factor)
        else:
            self.label_smoother = None

429
        self.state = TrainerState()
Sylvain Gugger's avatar
Sylvain Gugger committed
430
        self.control = TrainerControl()
431
432
433
        # Internal variable to count flos in each process, will be accumulated in `self.state.total_flos` then
        # returned to 0 every time flos need to be logged
        self.current_flos = 0
434
        self.hp_search_backend = None
435
        self.use_tune_checkpoints = False
436
        default_label_names = (
437
            ["start_positions", "end_positions"]
438
            if type(self.model).__name__ in MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES.values()
439
440
441
            else ["labels"]
        )
        self.label_names = default_label_names if self.args.label_names is None else self.args.label_names
Sylvain Gugger's avatar
Sylvain Gugger committed
442
443
        self.control = self.callback_handler.on_init_end(self.args, self.state, self.control)

444
445
446
        # very last
        self._memory_tracker.stop_and_update_metrics()

Sylvain Gugger's avatar
Sylvain Gugger committed
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
    def add_callback(self, callback):
        """
        Add a callback to the current list of :class:`~transformer.TrainerCallback`.

        Args:
           callback (:obj:`type` or :class:`~transformer.TrainerCallback`):
               A :class:`~transformer.TrainerCallback` class or an instance of a :class:`~transformer.TrainerCallback`.
               In the first case, will instantiate a member of that class.
        """
        self.callback_handler.add_callback(callback)

    def pop_callback(self, callback):
        """
        Remove a callback from the current list of :class:`~transformer.TrainerCallback` and returns it.

        If the callback is not found, returns :obj:`None` (and no error is raised).

        Args:
           callback (:obj:`type` or :class:`~transformer.TrainerCallback`):
               A :class:`~transformer.TrainerCallback` class or an instance of a :class:`~transformer.TrainerCallback`.
               In the first case, will pop the first member of that class found in the list of callbacks.

        Returns:
            :class:`~transformer.TrainerCallback`: The callback removed, if found.
        """
        return self.callback_handler.pop_callback(callback)

    def remove_callback(self, callback):
        """
        Remove a callback from the current list of :class:`~transformer.TrainerCallback`.

        Args:
           callback (:obj:`type` or :class:`~transformer.TrainerCallback`):
               A :class:`~transformer.TrainerCallback` class or an instance of a :class:`~transformer.TrainerCallback`.
               In the first case, will remove the first member of that class found in the list of callbacks.
        """
        self.callback_handler.remove_callback(callback)
Julien Chaumond's avatar
Julien Chaumond committed
484

485
    def _remove_unused_columns(self, dataset: "datasets.Dataset", description: Optional[str] = None):
486
        if not self.args.remove_unused_columns:
487
            return dataset
488
489
490
491
492
493
494
495
        if self._signature_columns is None:
            # Inspect model forward signature to keep only the arguments it accepts.
            signature = inspect.signature(self.model.forward)
            self._signature_columns = list(signature.parameters.keys())
            # Labels may be named label or label_ids, the default data collator handles that.
            self._signature_columns += ["label", "label_ids"]
        columns = [k for k in self._signature_columns if k in dataset.column_names]
        ignored_columns = list(set(dataset.column_names) - set(self._signature_columns))
496
497
498
499
500
501
        if len(ignored_columns) > 0:
            dset_description = "" if description is None else f"in the {description} set "
            logger.info(
                f"The following columns {dset_description} don't have a corresponding argument in "
                f"`{self.model.__class__.__name__}.forward` and have been ignored: {', '.join(ignored_columns)}."
            )
502

503
504
505
506
507
508
509
        if version.parse(datasets.__version__) < version.parse("1.4.0"):
            dataset.set_format(
                type=dataset.format["type"], columns=columns, format_kwargs=dataset.format["format_kwargs"]
            )
            return dataset
        else:
            return dataset.remove_columns(ignored_columns)
510

511
    def _get_train_sampler(self) -> Optional[torch.utils.data.sampler.Sampler]:
512
        if not isinstance(self.train_dataset, collections.abc.Sized):
513
            return None
514
515
516

        # Build the sampler.
        if self.args.group_by_length:
517
518
519
520
521
522
523
524
            if is_datasets_available() and isinstance(self.train_dataset, datasets.Dataset):
                lengths = (
                    self.train_dataset[self.args.length_column_name]
                    if self.args.length_column_name in self.train_dataset.column_names
                    else None
                )
            else:
                lengths = None
525
            model_input_name = self.tokenizer.model_input_names[0] if self.tokenizer is not None else None
526
            if self.args.world_size <= 1:
527
                return LengthGroupedSampler(
528
                    self.train_dataset, self.args.train_batch_size, lengths=lengths, model_input_name=model_input_name
529
                )
530
531
            else:
                return DistributedLengthGroupedSampler(
532
533
                    self.train_dataset,
                    self.args.train_batch_size,
534
535
                    num_replicas=self.args.world_size,
                    rank=self.args.process_index,
536
                    lengths=lengths,
537
                    model_input_name=model_input_name,
538
539
540
                )

        else:
541
            if self.args.world_size <= 1:
542
                return RandomSampler(self.train_dataset)
Sylvain Gugger's avatar
Sylvain Gugger committed
543
544
545
546
            elif (
                self.args.parallel_mode in [ParallelMode.TPU, ParallelMode.SAGEMAKER_MODEL_PARALLEL]
                and not self.args.dataloader_drop_last
            ):
547
548
549
550
551
552
553
                # Use a loop for TPUs when drop_last is False to have all batches have the same size.
                return DistributedSamplerWithLoop(
                    self.train_dataset,
                    batch_size=self.args.per_device_train_batch_size,
                    num_replicas=self.args.world_size,
                    rank=self.args.process_index,
                )
554
            else:
555
556
557
                return DistributedSampler(
                    self.train_dataset, num_replicas=self.args.world_size, rank=self.args.process_index
                )
558
559
560
561
562

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

Sylvain Gugger's avatar
Sylvain Gugger committed
563
564
        Will use no sampler if :obj:`self.train_dataset` does not implement :obj:`__len__`, a random sampler (adapted
        to distributed training if necessary) otherwise.
565
566
567
568
569

        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.")
570

571
572
573
574
575
        train_dataset = self.train_dataset
        if is_datasets_available() and isinstance(train_dataset, datasets.Dataset):
            train_dataset = self._remove_unused_columns(train_dataset, description="training")

        if isinstance(train_dataset, torch.utils.data.dataset.IterableDataset):
576
577
            if self.args.world_size > 1:
                train_dataset = IterableDatasetShard(
578
                    train_dataset,
579
580
581
582
583
                    batch_size=self.args.train_batch_size,
                    drop_last=self.args.dataloader_drop_last,
                    num_processes=self.args.world_size,
                    process_index=self.args.process_index,
                )
584

585
586
587
588
589
590
591
592
            return DataLoader(
                train_dataset,
                batch_size=self.args.train_batch_size,
                collate_fn=self.data_collator,
                num_workers=self.args.dataloader_num_workers,
                pin_memory=self.args.dataloader_pin_memory,
            )

593
594
595
        train_sampler = self._get_train_sampler()

        return DataLoader(
596
            train_dataset,
Julien Chaumond's avatar
Julien Chaumond committed
597
598
            batch_size=self.args.train_batch_size,
            sampler=train_sampler,
599
            collate_fn=self.data_collator,
Setu Shah's avatar
Setu Shah committed
600
            drop_last=self.args.dataloader_drop_last,
Chady Kamar's avatar
Chady Kamar committed
601
            num_workers=self.args.dataloader_num_workers,
602
            pin_memory=self.args.dataloader_pin_memory,
Julien Chaumond's avatar
Julien Chaumond committed
603
604
        )

605
    def _get_eval_sampler(self, eval_dataset: Dataset) -> Optional[torch.utils.data.sampler.Sampler]:
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
        # Deprecated code
        if self.args.use_legacy_prediction_loop:
            if is_torch_tpu_available():
                return SequentialDistributedSampler(
                    eval_dataset, num_replicas=xm.xrt_world_size(), rank=xm.get_ordinal()
                )
            elif is_sagemaker_mp_enabled():
                return SequentialDistributedSampler(
                    eval_dataset,
                    num_replicas=smp.dp_size(),
                    rank=smp.dp_rank(),
                    batch_size=self.args.per_device_eval_batch_size,
                )
            elif self.args.local_rank != -1:
                return SequentialDistributedSampler(eval_dataset)
            else:
                return SequentialSampler(eval_dataset)

        if self.args.world_size <= 1:
            return SequentialSampler(eval_dataset)
        else:
            return ShardSampler(
Sylvain Gugger's avatar
Sylvain Gugger committed
628
629
                eval_dataset,
                batch_size=self.args.per_device_eval_batch_size,
630
631
                num_processes=self.args.world_size,
                process_index=self.args.process_index,
Sylvain Gugger's avatar
Sylvain Gugger committed
632
            )
Lysandre Debut's avatar
Lysandre Debut committed
633

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

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

640
        Args:
641
            eval_dataset (:obj:`torch.utils.data.dataset.Dataset`, `optional`):
642
                If provided, will override :obj:`self.eval_dataset`. If it is an :obj:`datasets.Dataset`, columns not
643
                accepted by the ``model.forward()`` method are automatically removed. It must implement :obj:`__len__`.
644
        """
Julien Chaumond's avatar
Julien Chaumond committed
645
646
        if eval_dataset is None and self.eval_dataset is None:
            raise ValueError("Trainer: evaluation requires an eval_dataset.")
647
        eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset
648

649
650
651
        if is_datasets_available() and isinstance(eval_dataset, datasets.Dataset):
            eval_dataset = self._remove_unused_columns(eval_dataset, description="evaluation")

652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
        if isinstance(eval_dataset, torch.utils.data.dataset.IterableDataset):
            if self.args.world_size > 1:
                eval_dataset = IterableDatasetShard(
                    eval_dataset,
                    batch_size=self.args.eval_batch_size,
                    drop_last=self.args.dataloader_drop_last,
                    num_processes=self.args.world_size,
                    process_index=self.args.process_index,
                )
            return DataLoader(
                eval_dataset,
                batch_size=self.args.eval_batch_size,
                collate_fn=self.data_collator,
                num_workers=self.args.dataloader_num_workers,
                pin_memory=self.args.dataloader_pin_memory,
            )

669
        eval_sampler = self._get_eval_sampler(eval_dataset)
670

671
        return DataLoader(
672
            eval_dataset,
673
            sampler=eval_sampler,
Julien Chaumond's avatar
Julien Chaumond committed
674
            batch_size=self.args.eval_batch_size,
675
            collate_fn=self.data_collator,
Setu Shah's avatar
Setu Shah committed
676
            drop_last=self.args.dataloader_drop_last,
Chady Kamar's avatar
Chady Kamar committed
677
            num_workers=self.args.dataloader_num_workers,
678
            pin_memory=self.args.dataloader_pin_memory,
Julien Chaumond's avatar
Julien Chaumond committed
679
680
681
        )

    def get_test_dataloader(self, test_dataset: Dataset) -> DataLoader:
682
683
684
        """
        Returns the test :class:`~torch.utils.data.DataLoader`.

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

687
        Args:
688
            test_dataset (:obj:`torch.utils.data.dataset.Dataset`, `optional`):
689
                The test dataset to use. If it is an :obj:`datasets.Dataset`, columns not accepted by the
690
                ``model.forward()`` method are automatically removed. It must implement :obj:`__len__`.
691
        """
692
        if is_datasets_available() and isinstance(test_dataset, datasets.Dataset):
693
            test_dataset = self._remove_unused_columns(test_dataset, description="test")
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711

        if isinstance(test_dataset, torch.utils.data.dataset.IterableDataset):
            if self.args.world_size > 1:
                test_dataset = IterableDatasetShard(
                    test_dataset,
                    batch_size=self.args.eval_batch_size,
                    drop_last=self.args.dataloader_drop_last,
                    num_processes=self.args.world_size,
                    process_index=self.args.process_index,
                )
            return DataLoader(
                test_dataset,
                batch_size=self.args.eval_batch_size,
                collate_fn=self.data_collator,
                num_workers=self.args.dataloader_num_workers,
                pin_memory=self.args.dataloader_pin_memory,
            )

712
        test_sampler = self._get_eval_sampler(test_dataset)
Lysandre Debut's avatar
Lysandre Debut committed
713

714
715
        # We use the same batch_size as for eval.
        return DataLoader(
Julien Chaumond's avatar
Julien Chaumond committed
716
            test_dataset,
717
            sampler=test_sampler,
Julien Chaumond's avatar
Julien Chaumond committed
718
            batch_size=self.args.eval_batch_size,
719
            collate_fn=self.data_collator,
720
            drop_last=self.args.dataloader_drop_last,
721
            pin_memory=self.args.dataloader_pin_memory,
Julien Chaumond's avatar
Julien Chaumond committed
722
        )
Lysandre Debut's avatar
Lysandre Debut committed
723

724
    def create_optimizer_and_scheduler(self, num_training_steps: int):
725
726
727
        """
        Setup the optimizer and the learning rate scheduler.

728
729
730
731
732
733
734
735
736
737
738
        We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the
        Trainer's init through :obj:`optimizers`, or subclass and override this method (or :obj:`create_optimizer`
        and/or :obj:`create_scheduler`) in a subclass.
        """
        self.create_optimizer()
        self.create_scheduler(num_training_steps)

    def create_optimizer(self):
        """
        Setup the optimizer.

739
        We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the
740
        Trainer's init through :obj:`optimizers`, or subclass and override this method in a subclass.
741
        """
742
        if self.optimizer is None:
743
744
            decay_parameters = get_parameter_names(self.model, [torch.nn.LayerNorm])
            decay_parameters = [name for name in decay_parameters if "bias" not in name]
745
746
            optimizer_grouped_parameters = [
                {
747
                    "params": [p for n, p in self.model.named_parameters() if n in decay_parameters],
748
749
750
                    "weight_decay": self.args.weight_decay,
                },
                {
751
                    "params": [p for n, p in self.model.named_parameters() if n not in decay_parameters],
752
753
754
                    "weight_decay": 0.0,
                },
            ]
Sylvain Gugger's avatar
Sylvain Gugger committed
755
756
757
758
759
760
761
762
763
764
765
            optimizer_cls = Adafactor if self.args.adafactor else AdamW
            if self.args.adafactor:
                optimizer_cls = Adafactor
                optimizer_kwargs = {"scale_parameter": False, "relative_step": False}
            else:
                optimizer_cls = AdamW
                optimizer_kwargs = {
                    "betas": (self.args.adam_beta1, self.args.adam_beta2),
                    "eps": self.args.adam_epsilon,
                }
            optimizer_kwargs["lr"] = self.args.learning_rate
766
            if self.sharded_ddp == ShardedDDPOption.SIMPLE:
767
768
                self.optimizer = OSS(
                    params=optimizer_grouped_parameters,
Sylvain Gugger's avatar
Sylvain Gugger committed
769
770
                    optim=optimizer_cls,
                    **optimizer_kwargs,
771
772
                )
            else:
Sylvain Gugger's avatar
Sylvain Gugger committed
773
774
                self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)

Sylvain Gugger's avatar
Sylvain Gugger committed
775
776
777
        if is_sagemaker_mp_enabled():
            self.optimizer = smp.DistributedOptimizer(self.optimizer)

778
779
780
781
782
783
784
    def create_scheduler(self, num_training_steps: int):
        """
        Setup the scheduler. The optimizer of the trainer must have been set up before this method is called.

        Args:
            num_training_steps (int): The number of training steps to do.
        """
785
        if self.lr_scheduler is None:
786
787
788
789
790
791
            warmup_steps = (
                self.args.warmup_steps
                if self.args.warmup_steps > 0
                else math.ceil(num_training_steps * self.args.warmup_ratio)
            )

Sylvain Gugger's avatar
Sylvain Gugger committed
792
793
794
            self.lr_scheduler = get_scheduler(
                self.args.lr_scheduler_type,
                self.optimizer,
795
                num_warmup_steps=warmup_steps,
Sylvain Gugger's avatar
Sylvain Gugger committed
796
                num_training_steps=num_training_steps,
797
            )
Julien Chaumond's avatar
Julien Chaumond committed
798

799
    def num_examples(self, dataloader: DataLoader) -> int:
800
        """
801
        Helper to get number of samples in a :class:`~torch.utils.data.DataLoader` by accessing its dataset.
802

803
        Will raise an exception if the underlying dataset does not implement method :obj:`__len__`
804
        """
805
        return len(dataloader.dataset)
806

807
    def _hp_search_setup(self, trial: Union["optuna.Trial", Dict[str, Any]]):
Patrick von Platen's avatar
Patrick von Platen committed
808
        """HP search setup code"""
809
810
        self._trial = trial

811
812
        if self.hp_search_backend is None or trial is None:
            return
813
814
815
816
817
        if self.hp_search_backend == HPSearchBackend.OPTUNA:
            params = self.hp_space(trial)
        elif self.hp_search_backend == HPSearchBackend.RAY:
            params = trial
            params.pop("wandb", None)
818

819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
        for key, value in params.items():
            if not hasattr(self.args, key):
                raise AttributeError(
                    f"Trying to set {key} in the hyperparameter search but there is no corresponding field in `TrainingArguments`."
                )
            old_attr = getattr(self.args, key, None)
            # Casting value to the proper type
            if old_attr is not None:
                value = type(old_attr)(value)
            setattr(self.args, key, value)
        if self.hp_search_backend == HPSearchBackend.OPTUNA:
            logger.info("Trial:", trial.params)

    def _report_to_hp_search(
        self, trial: Union["optuna.Trial", Dict[str, Any]], epoch: int, metrics: Dict[str, float]
    ):
        if self.hp_search_backend is None or trial is None:
            return
837
        self.objective = self.compute_objective(metrics.copy())
838
        if self.hp_search_backend == HPSearchBackend.OPTUNA:
839
840
            import optuna

841
842
843
844
            trial.report(self.objective, epoch)
            if trial.should_prune():
                raise optuna.TrialPruned()
        elif self.hp_search_backend == HPSearchBackend.RAY:
845
846
            from ray import tune

847
            if self.control.should_save:
848
                self._tune_save_checkpoint()
849
850
            tune.report(objective=self.objective, **metrics)

851
    def _tune_save_checkpoint(self):
852
853
        from ray import tune

854
855
        if not self.use_tune_checkpoints:
            return
856
        with tune.checkpoint_dir(step=self.state.global_step) as checkpoint_dir:
857
            output_dir = os.path.join(checkpoint_dir, f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}")
858
            self.save_model(output_dir)
859
            if self.is_world_process_zero():
860
                self.state.save_to_json(os.path.join(output_dir, "trainer_state.json"))
861
862
863
                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"))

864
865
866
867
868
869
870
    def call_model_init(self, trial=None):
        model_init_argcount = len(inspect.signature(self.model_init).parameters)
        if model_init_argcount == 0:
            model = self.model_init()
        elif model_init_argcount == 1:
            model = self.model_init(trial)
        else:
871
872
873
874
            raise RuntimeError("model_init should have 0 or 1 argument.")

        if model is None:
            raise RuntimeError("model_init should not return None.")
875
876
877

        return model

878
    def _wrap_model(self, model, training=True):
Sylvain Gugger's avatar
Sylvain Gugger committed
879
880
881
882
883
884
        if is_sagemaker_mp_enabled():
            # Wrapping the base model twice in a DistributedModel will raise an error.
            if isinstance(self.model_wrapped, smp.model.DistributedModel):
                return self.model_wrapped
            return smp.DistributedModel(model, backward_passes_per_step=self.args.gradient_accumulation_steps)

885
886
        # already initialized its own DDP and AMP
        if self.deepspeed:
887
            return self.deepspeed
888

889
890
891
892
        # train/eval could be run multiple-times - if already wrapped, don't re-wrap it again
        if unwrap_model(model) is not model:
            return model

893
894
895
896
897
898
899
900
901
902
903
904
905
906
        # Mixed precision training with apex (torch < 1.6)
        if self.use_apex and training:
            model, self.optimizer = amp.initialize(model, self.optimizer, opt_level=self.args.fp16_opt_level)

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

        # Note: in torch.distributed mode, there's no point in wrapping the model
        # inside a DistributedDataParallel as we'll be under `no_grad` anyways.
        if not training:
            return model

        # Distributed training (should be after apex fp16 initialization)
907
908
909
910
911
912
913
914
915
        if self.sharded_ddp is not None:
            # Sharded DDP!
            if self.sharded_ddp == ShardedDDPOption.SIMPLE:
                model = ShardedDDP(model, self.optimizer)
            else:
                mixed_precision = self.args.fp16
                cpu_offload = ShardedDDPOption.OFFLOAD in self.args.sharded_ddp
                zero_3 = self.sharded_ddp == ShardedDDPOption.ZERO_DP_3
                # XXX: Breaking the self.model convention but I see no way around it for now.
916
917
                if ShardedDDPOption.AUTO_WRAP in self.args.sharded_ddp:
                    model = auto_wrap(model)
918
                self.model = model = FullyShardedDDP(
919
920
921
922
                    model,
                    mixed_precision=mixed_precision,
                    reshard_after_forward=zero_3,
                    cpu_offload=cpu_offload,
923
924
                ).to(self.args.device)

Sylvain Gugger's avatar
Sylvain Gugger committed
925
        elif is_sagemaker_dp_enabled():
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
            model = DDP(model, device_ids=[dist.get_local_rank()], broadcast_buffers=False)
        elif self.args.local_rank != -1:
            if self.args.ddp_find_unused_parameters is not None:
                find_unused_parameters = self.args.ddp_find_unused_parameters
            elif isinstance(model, PreTrainedModel):
                # find_unused_parameters breaks checkpointing as per
                # https://github.com/huggingface/transformers/pull/4659#issuecomment-643356021
                find_unused_parameters = not getattr(model.config, "gradient_checkpointing", False)
            else:
                find_unused_parameters = True
            model = torch.nn.parallel.DistributedDataParallel(
                model,
                device_ids=[self.args.local_rank],
                output_device=self.args.local_rank,
                find_unused_parameters=find_unused_parameters,
            )

        return model

945
946
    def train(
        self,
947
        resume_from_checkpoint: Optional[Union[str, bool]] = None,
948
        trial: Union["optuna.Trial", Dict[str, Any]] = None,
949
        **kwargs,
950
    ):
Julien Chaumond's avatar
Julien Chaumond committed
951
952
953
954
        """
        Main training entry point.

        Args:
955
956
957
958
959
            resume_from_checkpoint (:obj:`str` or :obj:`bool`, `optional`):
                If a :obj:`str`, local path to a saved checkpoint as saved by a previous instance of
                :class:`~transformers.Trainer`. If a :obj:`bool` and equals `True`, load the last checkpoint in
                `args.output_dir` as saved by a previous instance of :class:`~transformers.Trainer`. If present,
                training will resume from the model/optimizer/scheduler states loaded here.
960
961
            trial (:obj:`optuna.Trial` or :obj:`Dict[str, Any]`, `optional`):
                The trial run or the hyperparameter dictionary for hyperparameter search.
962
963
            kwargs:
                Additional keyword arguments used to hide deprecated arguments
Julien Chaumond's avatar
Julien Chaumond committed
964
        """
965
966
967
968

        # memory metrics - must set up as early as possible
        self._memory_tracker.start()

969
970
        args = self.args

971
972
        self.is_in_train = True

973
974
975
976
977
        # do_train is not a reliable argument, as it might not be set and .train() still called, so
        # the following is a workaround:
        if args.fp16_full_eval and not args.do_train:
            self.model = self.model.to(args.device)

978
979
980
981
982
983
984
985
986
        if "model_path" in kwargs:
            resume_from_checkpoint = kwargs.pop("model_path")
            warnings.warn(
                "`model_path` is deprecated and will be removed in a future version. Use `resume_from_checkpoint` "
                "instead.",
                FutureWarning,
            )
        if len(kwargs) > 0:
            raise TypeError(f"train() received got unexpected keyword arguments: {', '.join(list(kwargs.keys()))}.")
Sylvain Gugger's avatar
Sylvain Gugger committed
987
988
989
        # This might change the seed so needs to run first.
        self._hp_search_setup(trial)

990
        # Model re-init
991
        model_reloaded = False
992
        if self.model_init is not None:
Sylvain Gugger's avatar
Sylvain Gugger committed
993
            # Seed must be set before instantiating the model when using model_init.
994
            set_seed(args.seed)
995
996
            self.model = self.call_model_init(trial)
            model_reloaded = True
Sylvain Gugger's avatar
Sylvain Gugger committed
997
998
            # Reinitializes optimizer and scheduler
            self.optimizer, self.lr_scheduler = None, None
999

1000
        # Load potential model checkpoint
1001
        if isinstance(resume_from_checkpoint, bool) and resume_from_checkpoint:
1002
            resume_from_checkpoint = get_last_checkpoint(args.output_dir)
1003
            if resume_from_checkpoint is None:
1004
                raise ValueError(f"No valid checkpoint found in output directory ({args.output_dir})")
1005

1006
1007
1008
1009
        if resume_from_checkpoint is not None:
            if not os.path.isfile(os.path.join(resume_from_checkpoint, WEIGHTS_NAME)):
                raise ValueError(f"Can't find a valid checkpoint at {resume_from_checkpoint}")

1010
            logger.info(f"Loading model from {resume_from_checkpoint}).")
1011

1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
            if os.path.isfile(os.path.join(resume_from_checkpoint, CONFIG_NAME)):
                config = PretrainedConfig.from_json_file(os.path.join(resume_from_checkpoint, CONFIG_NAME))
                checkpoint_version = config.transformers_version
                if checkpoint_version is not None and checkpoint_version != __version__:
                    logger.warn(
                        f"You are resuming training from a checkpoint trained with {checkpoint_version} of "
                        f"Transformers but your current version is {__version__}. This is not recommended and could "
                        "yield to errors or unwanted behaviors."
                    )

1022
            if args.deepspeed:
1023
                # will be resumed in deepspeed_init
1024
                pass
1025
            else:
1026
1027
1028
                # We load the model state dict on the CPU to avoid an OOM error.
                state_dict = torch.load(os.path.join(resume_from_checkpoint, WEIGHTS_NAME), map_location="cpu")
                # If the model is on the GPU, it still works!
1029
1030
1031
1032
                self.model.load_state_dict(state_dict)

        # If model was re-initialized, put it on the right device and update self.model_wrapped
        if model_reloaded:
1033
            if self.place_model_on_device:
1034
                self.model = self.model.to(args.device)
1035
1036
            self.model_wrapped = self.model

1037
1038
1039
        # Keeping track whether we can can len() on the dataset or not
        train_dataset_is_sized = isinstance(self.train_dataset, collections.abc.Sized)

1040
        # Data loader and number of training steps
Julien Chaumond's avatar
Julien Chaumond committed
1041
        train_dataloader = self.get_train_dataloader()
1042
1043
1044
1045
1046
1047

        # Setting up training control variables:
        # number of training epochs: num_train_epochs
        # number of training steps per epoch: num_update_steps_per_epoch
        # total number of training steps to execute: max_steps
        if train_dataset_is_sized:
1048
            num_update_steps_per_epoch = len(train_dataloader) // args.gradient_accumulation_steps
1049
            num_update_steps_per_epoch = max(num_update_steps_per_epoch, 1)
1050
1051
1052
1053
            if args.max_steps > 0:
                max_steps = args.max_steps
                num_train_epochs = args.max_steps // num_update_steps_per_epoch + int(
                    args.max_steps % num_update_steps_per_epoch > 0
1054
1055
                )
            else:
1056
1057
                max_steps = math.ceil(args.num_train_epochs * num_update_steps_per_epoch)
                num_train_epochs = math.ceil(args.num_train_epochs)
Julien Chaumond's avatar
Julien Chaumond committed
1058
        else:
1059
            # see __init__. max_steps is set when the dataset has no __len__
1060
1061
            max_steps = args.max_steps
            num_train_epochs = int(args.num_train_epochs)
1062
            num_update_steps_per_epoch = max_steps
Julien Chaumond's avatar
Julien Chaumond committed
1063

1064
        delay_optimizer_creation = self.sharded_ddp is not None and self.sharded_ddp != ShardedDDPOption.SIMPLE
1065
        if args.deepspeed:
1066
            deepspeed_engine, optimizer, lr_scheduler = deepspeed_init(
1067
1068
                self, num_training_steps=max_steps, resume_from_checkpoint=resume_from_checkpoint
            )
1069
1070
1071
            self.model = deepspeed_engine.module
            self.model_wrapped = deepspeed_engine
            self.deepspeed = deepspeed_engine
1072
1073
            self.optimizer = optimizer
            self.lr_scheduler = lr_scheduler
1074
        elif not delay_optimizer_creation:
1075
1076
            self.create_optimizer_and_scheduler(num_training_steps=max_steps)

1077
        self.state = TrainerState()
1078
        self.state.is_hyper_param_search = trial is not None
Julien Chaumond's avatar
Julien Chaumond committed
1079

1080
        model = self._wrap_model(self.model_wrapped)
Julien Chaumond's avatar
Julien Chaumond committed
1081

1082
1083
1084
1085
        # for the rest of this function `model` is the outside model, whether it was wrapped or not
        if model is not self.model:
            self.model_wrapped = model

1086
1087
1088
        if delay_optimizer_creation:
            self.create_optimizer_and_scheduler(num_training_steps=max_steps)

1089
1090
1091
        # Check if saved optimizer or scheduler states exist
        self._load_optimizer_and_scheduler(resume_from_checkpoint)

1092
1093
        # important: at this point:
        # self.model         is the Transformers Model
1094
        # self.model_wrapped is DDP(Transformers Model), Deepspeed(Transformers Model), etc.
1095

Julien Chaumond's avatar
Julien Chaumond committed
1096
        # Train!
1097
        if is_torch_tpu_available():
Sylvain Gugger's avatar
Sylvain Gugger committed
1098
            world_size = xm.xrt_world_size()
1099
        elif args.local_rank != -1:
Sylvain Gugger's avatar
Sylvain Gugger committed
1100
            world_size = dist.get_world_size()
Lysandre Debut's avatar
Lysandre Debut committed
1101
        else:
Sylvain Gugger's avatar
Sylvain Gugger committed
1102
            world_size = 1
1103

1104
        total_train_batch_size = args.train_batch_size * args.gradient_accumulation_steps * world_size
1105
        num_examples = (
1106
            self.num_examples(train_dataloader) if train_dataset_is_sized else total_train_batch_size * args.max_steps
1107
1108
        )

Julien Chaumond's avatar
Julien Chaumond committed
1109
        logger.info("***** Running training *****")
1110
1111
        logger.info(f"  Num examples = {num_examples}")
        logger.info(f"  Num Epochs = {num_train_epochs}")
1112
        logger.info(f"  Instantaneous batch size per device = {args.per_device_train_batch_size}")
1113
        logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_train_batch_size}")
1114
        logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
1115
        logger.info(f"  Total optimization steps = {max_steps}")
Julien Chaumond's avatar
Julien Chaumond committed
1116

1117
        self.state.epoch = 0
1118
        start_time = time.time()
Julien Chaumond's avatar
Julien Chaumond committed
1119
1120
        epochs_trained = 0
        steps_trained_in_current_epoch = 0
1121
        steps_trained_progress_bar = None
1122

Julien Chaumond's avatar
Julien Chaumond committed
1123
        # Check if continuing training from a checkpoint
1124
1125
1126
1127
        if resume_from_checkpoint is not None and os.path.isfile(
            os.path.join(resume_from_checkpoint, "trainer_state.json")
        ):
            self.state = TrainerState.load_from_json(os.path.join(resume_from_checkpoint, "trainer_state.json"))
1128
            epochs_trained = self.state.global_step // num_update_steps_per_epoch
1129
            if not args.ignore_data_skip:
1130
                steps_trained_in_current_epoch = self.state.global_step % (num_update_steps_per_epoch)
1131
                steps_trained_in_current_epoch *= args.gradient_accumulation_steps
1132
1133
            else:
                steps_trained_in_current_epoch = 0
1134
1135

            logger.info("  Continuing training from checkpoint, will skip to saved global_step")
1136
1137
            logger.info(f"  Continuing training from epoch {epochs_trained}")
            logger.info(f"  Continuing training from global step {self.state.global_step}")
1138
            if not args.ignore_data_skip:
1139
1140
                logger.info(
                    f"  Will skip the first {epochs_trained} epochs then the first {steps_trained_in_current_epoch} "
1141
1142
                    "batches in the first epoch. If this takes a lot of time, you can add the `--ignore_data_skip` "
                    "flag to your launch command, but you will resume the training on data already seen by your model."
1143
                )
1144
1145
1146
                if self.is_local_process_zero() and not args.disable_tqdm:
                    steps_trained_progress_bar = tqdm(total=steps_trained_in_current_epoch)
                    steps_trained_progress_bar.set_description("Skipping the first batches")
1147

Sylvain Gugger's avatar
Sylvain Gugger committed
1148
1149
1150
1151
1152
        # Update the references
        self.callback_handler.model = self.model
        self.callback_handler.optimizer = self.optimizer
        self.callback_handler.lr_scheduler = self.lr_scheduler
        self.callback_handler.train_dataloader = train_dataloader
1153
1154
        self.state.trial_name = self.hp_name(trial) if self.hp_name is not None else None
        self.state.trial_params = hp_params(trial) if trial is not None else None
1155
1156
1157
1158
        # This should be the same if the state has been saved but in case the training arguments changed, it's safer
        # to set this after the load.
        self.state.max_steps = max_steps
        self.state.num_train_epochs = num_train_epochs
Sylvain Gugger's avatar
Sylvain Gugger committed
1159
1160
        self.state.is_local_process_zero = self.is_local_process_zero()
        self.state.is_world_process_zero = self.is_world_process_zero()
Julien Chaumond's avatar
Julien Chaumond committed
1161

1162
        # tr_loss is a tensor to avoid synchronization of TPUs through .item()
1163
        tr_loss = torch.tensor(0.0).to(args.device)
1164
1165
        # _total_loss_scalar is updated everytime .item() has to be called on tr_loss and stores the sum of all losses
        self._total_loss_scalar = 0.0
1166
        self._globalstep_last_logged = self.state.global_step
Julien Chaumond's avatar
Julien Chaumond committed
1167
        model.zero_grad()
Sylvain Gugger's avatar
Sylvain Gugger committed
1168

1169
        self.control = self.callback_handler.on_train_begin(args, self.state, self.control)
Sylvain Gugger's avatar
Sylvain Gugger committed
1170

1171
        # Skip the first epochs_trained epochs to get the random state of the dataloader at the right point.
1172
        if not args.ignore_data_skip:
1173
1174
1175
1176
1177
            for epoch in range(epochs_trained):
                # We just need to begin an iteration to create the randomization of the sampler.
                for _ in train_dataloader:
                    break

1178
        for epoch in range(epochs_trained, num_train_epochs):
1179
1180
            if isinstance(train_dataloader, DataLoader) and isinstance(train_dataloader.sampler, DistributedSampler):
                train_dataloader.sampler.set_epoch(epoch)
1181
1182
            elif isinstance(train_dataloader.dataset, IterableDatasetShard):
                train_dataloader.dataset.set_epoch(epoch)
1183

1184
            if is_torch_tpu_available():
1185
                parallel_loader = pl.ParallelLoader(train_dataloader, [args.device]).per_device_loader(args.device)
1186
                epoch_iterator = parallel_loader
1187
            else:
1188
                epoch_iterator = train_dataloader
1189

1190
            # Reset the past mems state at the beginning of each epoch if necessary.
1191
            if args.past_index >= 0:
1192
1193
                self._past = None

1194
            steps_in_epoch = (
1195
                len(epoch_iterator) if train_dataset_is_sized else args.max_steps * args.gradient_accumulation_steps
1196
            )
1197
            self.control = self.callback_handler.on_epoch_begin(args, self.state, self.control)
Sylvain Gugger's avatar
Sylvain Gugger committed
1198

Julien Chaumond's avatar
Julien Chaumond committed
1199
1200
1201
1202
1203
            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
1204
1205
                    if steps_trained_progress_bar is not None:
                        steps_trained_progress_bar.update(1)
Julien Chaumond's avatar
Julien Chaumond committed
1206
                    continue
1207
1208
1209
                elif steps_trained_progress_bar is not None:
                    steps_trained_progress_bar.close()
                    steps_trained_progress_bar = None
Julien Chaumond's avatar
Julien Chaumond committed
1210

1211
1212
                if step % args.gradient_accumulation_steps == 0:
                    self.control = self.callback_handler.on_step_begin(args, self.state, self.control)
Sylvain Gugger's avatar
Sylvain Gugger committed
1213

1214
                if (
1215
1216
1217
                    ((step + 1) % args.gradient_accumulation_steps != 0)
                    and args.local_rank != -1
                    and args._no_sync_in_gradient_accumulation
1218
                ):
1219
                    # Avoid unnecessary DDP synchronization since there will be no backward pass on this example.
1220
1221
1222
1223
                    with model.no_sync():
                        tr_loss += self.training_step(model, inputs)
                else:
                    tr_loss += self.training_step(model, inputs)
1224
                self.current_flos += float(self.floating_point_ops(inputs))
Julien Chaumond's avatar
Julien Chaumond committed
1225

1226
1227
1228
1229
                # Optimizer step for deepspeed must be called on every step regardless of the value of gradient_accumulation_steps
                if self.deepspeed:
                    self.deepspeed.step()

1230
                if (step + 1) % args.gradient_accumulation_steps == 0 or (
Julien Chaumond's avatar
Julien Chaumond committed
1231
                    # last step in epoch but step is always smaller than gradient_accumulation_steps
1232
                    steps_in_epoch <= args.gradient_accumulation_steps
1233
                    and (step + 1) == steps_in_epoch
Julien Chaumond's avatar
Julien Chaumond committed
1234
                ):
1235
                    # Gradient clipping
1236
                    if args.max_grad_norm is not None and args.max_grad_norm > 0 and not self.deepspeed:
1237
1238
                        # deepspeed does its own clipping

1239
1240
1241
1242
1243
1244
                        if self.use_amp:
                            # AMP: gradients need unscaling
                            self.scaler.unscale_(self.optimizer)

                        if hasattr(self.optimizer, "clip_grad_norm"):
                            # Some optimizers (like the sharded optimizer) have a specific way to do gradient clipping
1245
                            self.optimizer.clip_grad_norm(args.max_grad_norm)
1246
1247
                        elif hasattr(model, "clip_grad_norm_"):
                            # Some models (like FullyShardedDDP) have a specific way to do gradient clipping
1248
                            model.clip_grad_norm_(args.max_grad_norm)
1249
1250
1251
1252
                        else:
                            # Revert to normal clipping otherwise, handling Apex or full precision
                            torch.nn.utils.clip_grad_norm_(
                                amp.master_params(self.optimizer) if self.use_apex else model.parameters(),
1253
                                args.max_grad_norm,
1254
1255
1256
                            )

                    # Optimizer step
1257
                    optimizer_was_run = True
Stas Bekman's avatar
Stas Bekman committed
1258
                    if self.deepspeed:
1259
                        pass  # called outside the loop
Stas Bekman's avatar
Stas Bekman committed
1260
                    elif is_torch_tpu_available():
1261
                        xm.optimizer_step(self.optimizer)
1262
                    elif self.use_amp:
1263
                        scale_before = self.scaler.get_scale()
1264
                        self.scaler.step(self.optimizer)
1265
                        self.scaler.update()
1266
1267
                        scale_after = self.scaler.get_scale()
                        optimizer_was_run = scale_before <= scale_after
Lysandre Debut's avatar
Lysandre Debut committed
1268
                    else:
1269
                        self.optimizer.step()
Lysandre Debut's avatar
Lysandre Debut committed
1270

1271
                    if optimizer_was_run and not self.deepspeed:
1272
1273
                        self.lr_scheduler.step()

Julien Chaumond's avatar
Julien Chaumond committed
1274
                    model.zero_grad()
1275
                    self.state.global_step += 1
1276
                    self.state.epoch = epoch + (step + 1) / steps_in_epoch
1277
                    self.control = self.callback_handler.on_step_end(args, self.state, self.control)
Sylvain Gugger's avatar
Sylvain Gugger committed
1278

1279
                    self._maybe_log_save_evaluate(tr_loss, model, trial, epoch)
Julien Chaumond's avatar
Julien Chaumond committed
1280

Sylvain Gugger's avatar
Sylvain Gugger committed
1281
                if self.control.should_epoch_stop or self.control.should_training_stop:
Julien Chaumond's avatar
Julien Chaumond committed
1282
                    break
1283

1284
            self.control = self.callback_handler.on_epoch_end(args, self.state, self.control)
1285
            self._maybe_log_save_evaluate(tr_loss, model, trial, epoch)
1286

1287
            if args.tpu_metrics_debug or args.debug:
1288
1289
1290
1291
1292
1293
1294
1295
                if is_torch_tpu_available():
                    # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
                    xm.master_print(met.metrics_report())
                else:
                    logger.warning(
                        "You enabled PyTorch/XLA debug metrics but you don't have a TPU "
                        "configured. Check your training configuration if this is unexpected."
                    )
Sylvain Gugger's avatar
Sylvain Gugger committed
1296
            if self.control.should_training_stop:
1297
                break
Julien Chaumond's avatar
Julien Chaumond committed
1298

1299
        if args.past_index and hasattr(self, "_past"):
1300
1301
            # Clean the state at the end of training
            delattr(self, "_past")
Julien Chaumond's avatar
Julien Chaumond committed
1302
1303

        logger.info("\n\nTraining completed. Do not forget to share your model on huggingface.co/models =)\n\n")
1304
        if args.load_best_model_at_end and self.state.best_model_checkpoint is not None:
1305
1306
1307
            # Wait for everyone to get here so we are sur the model has been saved by process 0.
            if is_torch_tpu_available():
                xm.rendezvous("load_best_model_at_end")
1308
            elif args.local_rank != -1:
1309
1310
                dist.barrier()

1311
1312
1313
            logger.info(
                f"Loading best model from {self.state.best_model_checkpoint} (score: {self.state.best_metric})."
            )
1314
1315
1316
1317
            # We load the model state dict on the CPU to avoid an OOM error.
            state_dict = torch.load(os.path.join(self.state.best_model_checkpoint, WEIGHTS_NAME), map_location="cpu")
            # If the model is on the GPU, it still works!
            self.model.load_state_dict(state_dict)
1318

1319
1320
1321
1322
1323
            if self.deepspeed:
                self.deepspeed.load_checkpoint(
                    self.state.best_model_checkpoint, load_optimizer_states=False, load_lr_scheduler_states=False
                )

1324
        metrics = speed_metrics("train", start_time, self.state.max_steps)
1325
1326
        self.store_flos()
        metrics["total_flos"] = self.state.total_flos
1327
        self.log(metrics)
1328

1329
        self.control = self.callback_handler.on_train_end(args, self.state, self.control)
1330
1331
        # add remaining tr_loss
        self._total_loss_scalar += tr_loss.item()
Sylvain Gugger's avatar
Sylvain Gugger committed
1332

1333
        self.is_in_train = False
1334

1335
1336
        self._memory_tracker.stop_and_update_metrics(metrics)

1337
        return TrainOutput(self.state.global_step, self._total_loss_scalar / self.state.global_step, metrics)
1338

1339
    def _maybe_log_save_evaluate(self, tr_loss, model, trial, epoch):
Sylvain Gugger's avatar
Sylvain Gugger committed
1340
1341
1342
        if self.control.should_log:
            logs: Dict[str, float] = {}
            tr_loss_scalar = tr_loss.item()
1343
1344
1345
            # reset tr_loss to zero
            tr_loss -= tr_loss

1346
            logs["loss"] = round(tr_loss_scalar / (self.state.global_step - self._globalstep_last_logged), 4)
1347
            logs["learning_rate"] = self._get_learning_rate()
1348

1349
            self._total_loss_scalar += tr_loss_scalar
1350
            self._globalstep_last_logged = self.state.global_step
Sylvain Gugger's avatar
Sylvain Gugger committed
1351
1352
1353
1354
1355
1356
1357

            self.log(logs)

        metrics = None
        if self.control.should_evaluate:
            metrics = self.evaluate()
            self._report_to_hp_search(trial, epoch, metrics)
1358

Sylvain Gugger's avatar
Sylvain Gugger committed
1359
1360
1361
1362
1363
        if self.control.should_save:
            self._save_checkpoint(model, trial, metrics=metrics)
            self.control = self.callback_handler.on_save(self.args, self.state, self.control)

    def _save_checkpoint(self, model, trial, metrics=None):
1364
        # In all cases, including ddp/dp/deepspeed, self.model is always a reference to the model we
1365
        # want to save except FullyShardedDDP.
1366
        # assert unwrap_model(model) is self.model, "internal model should be a reference to self.model"
1367

1368
        # Save model checkpoint
1369
        checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}"
1370

1371
        if self.hp_search_backend is not None and trial is not None:
1372
1373
1374
1375
1376
1377
            if self.hp_search_backend == HPSearchBackend.OPTUNA:
                run_id = trial.number
            else:
                from ray import tune

                run_id = tune.get_trial_id()
1378
            run_name = self.hp_name(trial) if self.hp_name is not None else f"run-{run_id}"
1379
            run_dir = os.path.join(self.args.output_dir, run_name)
1380
        else:
1381
            run_dir = self.args.output_dir
1382
            self.store_flos()
1383

1384
        output_dir = os.path.join(run_dir, checkpoint_folder)
1385
        self.save_model(output_dir)
1386
        if self.deepspeed:
1387
1388
            # under zero3 model file itself doesn't get saved since it's bogus! Unless deepspeed
            # config `stage3_gather_fp16_weights_on_model_save` is True
1389
            self.deepspeed.save_checkpoint(output_dir)
1390
1391

        # Save optimizer and scheduler
1392
        if self.sharded_ddp == ShardedDDPOption.SIMPLE:
1393
            self.optimizer.consolidate_state_dict()
1394

1395
1396
1397
1398
1399
1400
        if is_torch_tpu_available():
            xm.rendezvous("saving_optimizer_states")
            xm.save(self.optimizer.state_dict(), os.path.join(output_dir, "optimizer.pt"))
            with warnings.catch_warnings(record=True) as caught_warnings:
                xm.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, "scheduler.pt"))
                reissue_pt_warnings(caught_warnings)
Sylvain Gugger's avatar
Sylvain Gugger committed
1401
1402
1403
1404
1405
1406
1407
1408
1409
        elif is_sagemaker_mp_enabled():
            # Consolidate the state dict on all processed of dp_rank 0
            opt_state_dict = self.optimizer.state_dict()
            # Save it and the scheduler on the main process
            if self.is_world_process_zero():
                torch.save(opt_state_dict, os.path.join(output_dir, "optimizer.pt"))
                with warnings.catch_warnings(record=True) as caught_warnings:
                    torch.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, "scheduler.pt"))
                reissue_pt_warnings(caught_warnings)
1410
1411
        elif self.is_world_process_zero() and not self.deepspeed:
            # deepspeed.save_checkpoint above saves model/optim/sched
1412
1413
1414
1415
1416
1417
            torch.save(self.optimizer.state_dict(), os.path.join(output_dir, "optimizer.pt"))
            with warnings.catch_warnings(record=True) as caught_warnings:
                torch.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, "scheduler.pt"))
            reissue_pt_warnings(caught_warnings)

        # Determine the new best metric / best model checkpoint
Sylvain Gugger's avatar
Sylvain Gugger committed
1418
        if metrics is not None and self.args.metric_for_best_model is not None:
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
            metric_to_check = self.args.metric_for_best_model
            if not metric_to_check.startswith("eval_"):
                metric_to_check = f"eval_{metric_to_check}"
            metric_value = metrics[metric_to_check]

            operator = np.greater if self.args.greater_is_better else np.less
            if (
                self.state.best_metric is None
                or self.state.best_model_checkpoint is None
                or operator(metric_value, self.state.best_metric)
            ):
                self.state.best_metric = metric_value
                self.state.best_model_checkpoint = output_dir

        # Save the Trainer state
        if self.is_world_process_zero():
            self.state.save_to_json(os.path.join(output_dir, "trainer_state.json"))

        # Maybe delete some older checkpoints.
        if self.is_world_process_zero():
1439
            self._rotate_checkpoints(use_mtime=True, output_dir=run_dir)
1440

1441
    def _load_optimizer_and_scheduler(self, checkpoint):
Sylvain Gugger's avatar
Sylvain Gugger committed
1442
        """If optimizer and scheduler states exist, load them."""
1443
        if checkpoint is None:
1444
1445
            return

1446
        if self.deepspeed:
1447
            # deepspeed loads optimizer/lr_scheduler together with the model in deepspeed_init
1448
1449
            return

1450
1451
        if os.path.isfile(os.path.join(checkpoint, "optimizer.pt")) and os.path.isfile(
            os.path.join(checkpoint, "scheduler.pt")
Sylvain Gugger's avatar
Sylvain Gugger committed
1452
1453
1454
1455
        ):
            # Load in optimizer and scheduler states
            if is_torch_tpu_available():
                # On TPU we have to take some extra precautions to properly load the states on the right device.
1456
                optimizer_state = torch.load(os.path.join(checkpoint, "optimizer.pt"), map_location="cpu")
Sylvain Gugger's avatar
Sylvain Gugger committed
1457
                with warnings.catch_warnings(record=True) as caught_warnings:
1458
                    lr_scheduler_state = torch.load(os.path.join(checkpoint, "scheduler.pt"), map_location="cpu")
Sylvain Gugger's avatar
Sylvain Gugger committed
1459
1460
1461
1462
1463
1464
1465
1466
                reissue_pt_warnings(caught_warnings)

                xm.send_cpu_data_to_device(optimizer_state, self.args.device)
                xm.send_cpu_data_to_device(lr_scheduler_state, self.args.device)

                self.optimizer.load_state_dict(optimizer_state)
                self.lr_scheduler.load_state_dict(lr_scheduler_state)
            else:
Sylvain Gugger's avatar
Sylvain Gugger committed
1467
                map_location = "cpu" if is_sagemaker_mp_enabled() else self.args.device
Sylvain Gugger's avatar
Sylvain Gugger committed
1468
                self.optimizer.load_state_dict(
Sylvain Gugger's avatar
Sylvain Gugger committed
1469
                    torch.load(os.path.join(checkpoint, "optimizer.pt"), map_location=map_location)
Sylvain Gugger's avatar
Sylvain Gugger committed
1470
1471
                )
                with warnings.catch_warnings(record=True) as caught_warnings:
1472
                    self.lr_scheduler.load_state_dict(torch.load(os.path.join(checkpoint, "scheduler.pt")))
Sylvain Gugger's avatar
Sylvain Gugger committed
1473
1474
                reissue_pt_warnings(caught_warnings)

1475
1476
1477
1478
1479
1480
1481
    def hyperparameter_search(
        self,
        hp_space: Optional[Callable[["optuna.Trial"], Dict[str, float]]] = None,
        compute_objective: Optional[Callable[[Dict[str, float]], float]] = None,
        n_trials: int = 20,
        direction: str = "minimize",
        backend: Optional[Union["str", HPSearchBackend]] = None,
1482
        hp_name: Optional[Callable[["optuna.Trial"], str]] = None,
1483
        **kwargs,
1484
1485
    ) -> BestRun:
        """
1486
        Launch an hyperparameter search using ``optuna`` or ``Ray Tune``. The optimized quantity is determined by
Sylvain Gugger's avatar
Sylvain Gugger committed
1487
1488
        :obj:`compute_objective`, which defaults to a function returning the evaluation loss when no metric is
        provided, the sum of all metrics otherwise.
1489

Sylvain Gugger's avatar
Sylvain Gugger committed
1490
1491
1492
1493
1494
1495
1496
        .. warning::

            To use this method, you need to have provided a ``model_init`` when initializing your
            :class:`~transformers.Trainer`: we need to reinitialize the model at each new run. This is incompatible
            with the ``optimizers`` argument, so you need to subclass :class:`~transformers.Trainer` and override the
            method :meth:`~transformers.Trainer.create_optimizer_and_scheduler` for custom optimizer/scheduler.

1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
        Args:
            hp_space (:obj:`Callable[["optuna.Trial"], Dict[str, float]]`, `optional`):
                A function that defines the hyperparameter search space. Will default to
                :func:`~transformers.trainer_utils.default_hp_space_optuna` or
                :func:`~transformers.trainer_utils.default_hp_space_ray` depending on your backend.
            compute_objective (:obj:`Callable[[Dict[str, float]], float]`, `optional`):
                A function computing the objective to minimize or maximize from the metrics returned by the
                :obj:`evaluate` method. Will default to :func:`~transformers.trainer_utils.default_compute_objective`.
            n_trials (:obj:`int`, `optional`, defaults to 100):
                The number of trial runs to test.
            direction(:obj:`str`, `optional`, defaults to :obj:`"minimize"`):
                Whether to optimize greater or lower objects. Can be :obj:`"minimize"` or :obj:`"maximize"`, you should
                pick :obj:`"minimize"` when optimizing the validation loss, :obj:`"maximize"` when optimizing one or
                several metrics.
            backend(:obj:`str` or :class:`~transformers.training_utils.HPSearchBackend`, `optional`):
                The backend to use for hyperparameter search. Will default to optuna or Ray Tune, depending on which
                one is installed. If both are installed, will default to optuna.
            kwargs:
                Additional keyword arguments passed along to :obj:`optuna.create_study` or :obj:`ray.tune.run`. For
                more information see:

Sylvain Gugger's avatar
Sylvain Gugger committed
1518
                - the documentation of `optuna.create_study
1519
                  <https://optuna.readthedocs.io/en/stable/reference/generated/optuna.study.create_study.html>`__
Sylvain Gugger's avatar
Sylvain Gugger committed
1520
1521
                - the documentation of `tune.run
                  <https://docs.ray.io/en/latest/tune/api_docs/execution.html#tune-run>`__
1522
1523

        Returns:
Tiger's avatar
Tiger committed
1524
            :class:`transformers.trainer_utils.BestRun`: All the information about the best run.
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
        """
        if backend is None:
            backend = default_hp_search_backend()
            if backend is None:
                raise RuntimeError(
                    "At least one of optuna or ray should be installed. "
                    "To install optuna run `pip install optuna`."
                    "To install ray run `pip install ray[tune]`."
                )
        backend = HPSearchBackend(backend)
        if backend == HPSearchBackend.OPTUNA and not is_optuna_available():
Sylvain Gugger's avatar
Sylvain Gugger committed
1536
            raise RuntimeError("You picked the optuna backend, but it is not installed. Use `pip install optuna`.")
1537
        if backend == HPSearchBackend.RAY and not is_ray_tune_available():
1538
            raise RuntimeError(
Sylvain Gugger's avatar
Sylvain Gugger committed
1539
                "You picked the Ray Tune backend, but it is not installed. Use `pip install 'ray[tune]'`."
1540
1541
            )
        self.hp_search_backend = backend
Sylvain Gugger's avatar
Sylvain Gugger committed
1542
1543
1544
1545
1546
        if self.model_init is None:
            raise RuntimeError(
                "To use hyperparameter search, you need to pass your model through a model_init function."
            )

1547
        self.hp_space = default_hp_space[backend] if hp_space is None else hp_space
1548
        self.hp_name = hp_name
1549
1550
        self.compute_objective = default_compute_objective if compute_objective is None else compute_objective

1551
1552
        run_hp_search = run_hp_search_optuna if backend == HPSearchBackend.OPTUNA else run_hp_search_ray
        best_run = run_hp_search(self, n_trials, direction, **kwargs)
1553
1554
1555
1556

        self.hp_search_backend = None
        return best_run

Sylvain Gugger's avatar
Sylvain Gugger committed
1557
    def log(self, logs: Dict[str, float]) -> None:
1558
1559
1560
1561
1562
1563
1564
1565
1566
        """
        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.
        """
1567
        if self.state.epoch is not None:
1568
            logs["epoch"] = round(self.state.epoch, 2)
1569

1570
1571
        output = {**logs, **{"step": self.state.global_step}}
        self.state.log_history.append(output)
1572
        self.control = self.callback_handler.on_log(self.args, self.state, self.control, logs)
Julien Chaumond's avatar
Julien Chaumond committed
1573

sgugger's avatar
Fix CI  
sgugger committed
1574
    def _prepare_inputs(self, inputs: Dict[str, Union[torch.Tensor, Any]]) -> Dict[str, Union[torch.Tensor, Any]]:
1575
1576
1577
1578
        """
        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
1579
        for k, v in inputs.items():
1580
1581
            if isinstance(v, torch.Tensor):
                inputs[k] = v.to(self.args.device)
Julien Chaumond's avatar
Julien Chaumond committed
1582

1583
1584
        if self.args.past_index >= 0 and self._past is not None:
            inputs["mems"] = self._past
1585

1586
1587
        return inputs

1588
    def training_step(self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]]) -> torch.Tensor:
1589
        """
1590
        Perform a training step on a batch of inputs.
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603

        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:
1604
            :obj:`torch.Tensor`: The tensor with training loss on this batch.
1605
1606
        """
        model.train()
1607
        inputs = self._prepare_inputs(inputs)
1608

Sylvain Gugger's avatar
Sylvain Gugger committed
1609
1610
1611
1612
        if is_sagemaker_mp_enabled():
            loss_mb = smp_forward_backward(model, inputs, self.args.gradient_accumulation_steps)
            return loss_mb.reduce_mean().detach().to(self.args.device)

1613
        if self.use_amp:
1614
            with autocast():
Sylvain Gugger's avatar
Sylvain Gugger committed
1615
                loss = self.compute_loss(model, inputs)
1616
        else:
Sylvain Gugger's avatar
Sylvain Gugger committed
1617
            loss = self.compute_loss(model, inputs)
1618

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

1622
1623
        if self.args.gradient_accumulation_steps > 1 and not self.deepspeed:
            # deepspeed handles loss scaling by gradient_accumulation_steps in its `backward`
Julien Chaumond's avatar
Julien Chaumond committed
1624
1625
            loss = loss / self.args.gradient_accumulation_steps

1626
        if self.use_amp:
1627
            self.scaler.scale(loss).backward()
1628
        elif self.use_apex:
1629
            with amp.scale_loss(loss, self.optimizer) as scaled_loss:
Julien Chaumond's avatar
Julien Chaumond committed
1630
                scaled_loss.backward()
1631
        elif self.deepspeed:
1632
1633
            # loss gets scaled under gradient_accumulation_steps in deepspeed
            loss = self.deepspeed.backward(loss)
Julien Chaumond's avatar
Julien Chaumond committed
1634
1635
1636
        else:
            loss.backward()

1637
        return loss.detach()
Julien Chaumond's avatar
Julien Chaumond committed
1638

1639
    def compute_loss(self, model, inputs, return_outputs=False):
Sylvain Gugger's avatar
Sylvain Gugger committed
1640
1641
1642
1643
1644
        """
        How the loss is computed by Trainer. By default, all models return the loss in the first element.

        Subclass and override for custom behavior.
        """
1645
1646
1647
1648
        if self.label_smoother is not None and "labels" in inputs:
            labels = inputs.pop("labels")
        else:
            labels = None
Sylvain Gugger's avatar
Sylvain Gugger committed
1649
1650
        outputs = model(**inputs)
        # Save past state if it exists
1651
        # TODO: this needs to be fixed and made cleaner later.
Sylvain Gugger's avatar
Sylvain Gugger committed
1652
1653
        if self.args.past_index >= 0:
            self._past = outputs[self.args.past_index]
Sylvain Gugger's avatar
Sylvain Gugger committed
1654

1655
        if labels is not None:
1656
            loss = self.label_smoother(outputs, labels)
Sylvain Gugger's avatar
Sylvain Gugger committed
1657
1658
        else:
            # We don't use .loss here since the model may return tuples instead of ModelOutput.
1659
1660
1661
            loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0]

        return (loss, outputs) if return_outputs else loss
Sylvain Gugger's avatar
Sylvain Gugger committed
1662

1663
1664
    def is_local_process_zero(self) -> bool:
        """
Sylvain Gugger's avatar
Sylvain Gugger committed
1665
1666
        Whether or not this process is the local (e.g., on one machine if training in a distributed fashion on several
        machines) main process.
1667
        """
1668
        if is_torch_tpu_available():
Lysandre Debut's avatar
Lysandre Debut committed
1669
            return xm.is_master_ordinal(local=True)
Sylvain Gugger's avatar
Sylvain Gugger committed
1670
1671
        elif is_sagemaker_mp_enabled():
            return smp.local_rank() == 0
Lysandre Debut's avatar
Lysandre Debut committed
1672
1673
1674
        else:
            return self.args.local_rank in [-1, 0]

1675
1676
    def is_world_process_zero(self) -> bool:
        """
Sylvain Gugger's avatar
Sylvain Gugger committed
1677
1678
        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
1679
        """
1680
        if is_torch_tpu_available():
Lysandre Debut's avatar
Lysandre Debut committed
1681
            return xm.is_master_ordinal(local=False)
Sylvain Gugger's avatar
Sylvain Gugger committed
1682
1683
        elif is_sagemaker_mp_enabled():
            return smp.rank() == 0
Lysandre Debut's avatar
Lysandre Debut committed
1684
        else:
Sylvain Gugger's avatar
Sylvain Gugger committed
1685
            return self.args.process_index == 0
Julien Chaumond's avatar
Julien Chaumond committed
1686
1687
1688

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

1691
        Will only save from the main process.
Julien Chaumond's avatar
Julien Chaumond committed
1692
        """
1693
1694
1695
1696

        if output_dir is None:
            output_dir = self.args.output_dir

1697
        if is_torch_tpu_available():
1698
            self._save_tpu(output_dir)
Sylvain Gugger's avatar
Sylvain Gugger committed
1699
1700
1701
1702
1703
        elif is_sagemaker_mp_enabled():
            # Calling the state_dict needs to be done on the wrapped model and on all processes.
            state_dict = self.model_wrapped.state_dict()
            if self.is_world_process_zero():
                self._save(output_dir, state_dict=state_dict)
1704
1705
1706
1707
        elif (
            ShardedDDPOption.ZERO_DP_2 in self.args.sharded_ddp or ShardedDDPOption.ZERO_DP_3 in self.args.sharded_ddp
        ):
            state_dict = self.model.state_dict()
1708

1709
            if self.is_world_process_zero():
1710
                self._save(output_dir, state_dict=state_dict)
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
        elif self.deepspeed:

            # this takes care of everything as long as we aren't under zero3
            if self.is_world_process_zero():
                self._save(output_dir)

            if is_deepspeed_zero3_enabled():
                # It's too complicated to try to override different places where the weights dump gets
                # saved, so since under zero3 the file is bogus, simply delete it. The user should
                # either user deepspeed checkpoint to resume or to recover full weights use
                # zero_to_fp32.py stored in the checkpoint.
                if self.is_world_process_zero():
                    file = os.path.join(output_dir, WEIGHTS_NAME)
                    if os.path.isfile(file):
                        # logger.info(f"deepspeed zero3: removing {file}, see zero_to_fp32.py to recover weights")
                        os.remove(file)

                # now save the real model if stage3_gather_fp16_weights_on_model_save=True
                # if false it will not be saved.
                # This must be called on all ranks
                self.deepspeed.save_fp16_model(output_dir, WEIGHTS_NAME)

1733
1734
        elif self.is_world_process_zero():
            self._save(output_dir)
Julien Chaumond's avatar
Julien Chaumond committed
1735

1736
1737
    def _save_tpu(self, output_dir: Optional[str] = None):
        output_dir = output_dir if output_dir is not None else self.args.output_dir
1738
        logger.info(f"Saving model checkpoint to {output_dir}")
1739
1740
1741
1742
1743
1744
1745
1746

        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()`
        xm.rendezvous("saving_checkpoint")
1747
        if not isinstance(self.model, PreTrainedModel):
1748
1749
1750
1751
1752
1753
1754
            if isinstance(unwrap_model(self.model), PreTrainedModel):
                unwrap_model(self.model).save_pretrained(
                    output_dir,
                    save_config=self.is_world_process_zero(),
                    state_dict=self.model.state_dict(),
                    save_function=xm.save,
                )
1755
1756
            else:
                logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.")
1757
1758
                state_dict = self.model.state_dict()
                xm.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME))
1759
        else:
1760
            self.model.save_pretrained(output_dir, save_config=self.is_world_process_zero(), save_function=xm.save)
Sylvain Gugger's avatar
Sylvain Gugger committed
1761
        if self.tokenizer is not None and self.is_world_process_zero():
1762
            self.tokenizer.save_pretrained(output_dir)
1763

1764
    def _save(self, output_dir: Optional[str] = None, state_dict=None):
1765
        # If we are executing this function, we are the process zero, so we don't check for that.
Julien Chaumond's avatar
Julien Chaumond committed
1766
1767
        output_dir = output_dir if output_dir is not None else self.args.output_dir
        os.makedirs(output_dir, exist_ok=True)
1768
        logger.info(f"Saving model checkpoint to {output_dir}")
Julien Chaumond's avatar
Julien Chaumond committed
1769
1770
1771
        # Save a trained model and configuration using `save_pretrained()`.
        # They can then be reloaded using `from_pretrained()`
        if not isinstance(self.model, PreTrainedModel):
1772
            if isinstance(unwrap_model(self.model), PreTrainedModel):
1773
1774
1775
                if state_dict is None:
                    state_dict = self.model.state_dict()
                unwrap_model(self.model).save_pretrained(output_dir, state_dict=state_dict)
1776
1777
            else:
                logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.")
1778
1779
                if state_dict is None:
                    state_dict = self.model.state_dict()
1780
                torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME))
1781
        else:
1782
            self.model.save_pretrained(output_dir, state_dict=state_dict)
1783
        if self.tokenizer is not None:
1784
            self.tokenizer.save_pretrained(output_dir)
Julien Chaumond's avatar
Julien Chaumond committed
1785
1786
1787

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

1789
    def store_flos(self):
1790
        # Storing the number of floating-point operations that went into the model
1791
1792
1793
1794
1795
1796
        if self.args.local_rank != -1:
            self.state.total_flos += distributed_broadcast_scalars([self.current_flos]).sum().item()
            self.current_flos = 0
        else:
            self.state.total_flos = self.current_flos
            self.current_flos = 0
Julien Chaumond's avatar
Julien Chaumond committed
1797

1798
1799
1800
    def _sorted_checkpoints(
        self, output_dir=None, checkpoint_prefix=PREFIX_CHECKPOINT_DIR, use_mtime=False
    ) -> List[str]:
Julien Chaumond's avatar
Julien Chaumond committed
1801
1802
        ordering_and_checkpoint_path = []

1803
        glob_checkpoints = [str(x) for x in Path(output_dir).glob(f"{checkpoint_prefix}-*")]
Julien Chaumond's avatar
Julien Chaumond committed
1804
1805
1806
1807
1808

        for path in glob_checkpoints:
            if use_mtime:
                ordering_and_checkpoint_path.append((os.path.getmtime(path), path))
            else:
1809
                regex_match = re.match(f".*{checkpoint_prefix}-([0-9]+)", path)
Julien Chaumond's avatar
Julien Chaumond committed
1810
1811
1812
1813
1814
                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]
1815
1816
        # Make sure we don't delete the best model.
        if self.state.best_model_checkpoint is not None:
1817
            best_model_index = checkpoints_sorted.index(str(Path(self.state.best_model_checkpoint)))
1818
            checkpoints_sorted[best_model_index], checkpoints_sorted[-1] = (
1819
1820
1821
                checkpoints_sorted[-1],
                checkpoints_sorted[best_model_index],
            )
Julien Chaumond's avatar
Julien Chaumond committed
1822
1823
        return checkpoints_sorted

1824
    def _rotate_checkpoints(self, use_mtime=False, output_dir=None) -> None:
1825
        if self.args.save_total_limit is None or self.args.save_total_limit <= 0:
Julien Chaumond's avatar
Julien Chaumond committed
1826
1827
1828
            return

        # Check if we should delete older checkpoint(s)
1829
        checkpoints_sorted = self._sorted_checkpoints(use_mtime=use_mtime, output_dir=output_dir)
Julien Chaumond's avatar
Julien Chaumond committed
1830
1831
1832
1833
1834
1835
        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:
1836
            logger.info(f"Deleting older checkpoint [{checkpoint}] due to args.save_total_limit")
Julien Chaumond's avatar
Julien Chaumond committed
1837
1838
            shutil.rmtree(checkpoint)

1839
    def evaluate(
1840
1841
1842
1843
        self,
        eval_dataset: Optional[Dataset] = None,
        ignore_keys: Optional[List[str]] = None,
        metric_key_prefix: str = "eval",
1844
    ) -> Dict[str, float]:
Julien Chaumond's avatar
Julien Chaumond committed
1845
        """
1846
        Run evaluation and returns metrics.
Julien Chaumond's avatar
Julien Chaumond committed
1847

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

1851
1852
        You can also subclass and override this method to inject custom behavior.

Julien Chaumond's avatar
Julien Chaumond committed
1853
        Args:
1854
            eval_dataset (:obj:`Dataset`, `optional`):
1855
                Pass a dataset if you wish to override :obj:`self.eval_dataset`. If it is an :obj:`datasets.Dataset`,
Sylvain Gugger's avatar
Sylvain Gugger committed
1856
1857
                columns not accepted by the ``model.forward()`` method are automatically removed. It must implement the
                :obj:`__len__` method.
1858
1859
1860
            ignore_keys (:obj:`Lst[str]`, `optional`):
                A list of keys in the output of your model (if it is a dictionary) that should be ignored when
                gathering predictions.
1861
1862
1863
            metric_key_prefix (:obj:`str`, `optional`, defaults to :obj:`"eval"`):
                An optional prefix to be used as the metrics key prefix. For example the metrics "bleu" will be named
                "eval_bleu" if the prefix is "eval" (default)
1864

Julien Chaumond's avatar
Julien Chaumond committed
1865
        Returns:
1866
1867
            A dictionary containing the evaluation loss and the potential metrics computed from the predictions. The
            dictionary also contains the epoch number which comes from the training state.
Julien Chaumond's avatar
Julien Chaumond committed
1868
        """
1869
1870
1871
        # memory metrics - must set up as early as possible
        self._memory_tracker.start()

Julien Chaumond's avatar
Julien Chaumond committed
1872
        eval_dataloader = self.get_eval_dataloader(eval_dataset)
1873
        start_time = time.time()
Julien Chaumond's avatar
Julien Chaumond committed
1874

1875
1876
        eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
        output = eval_loop(
1877
1878
1879
1880
1881
            eval_dataloader,
            description="Evaluation",
            # No point gathering the predictions if there are no metrics, otherwise we defer to
            # self.args.prediction_loss_only
            prediction_loss_only=True if self.compute_metrics is None else None,
1882
            ignore_keys=ignore_keys,
1883
            metric_key_prefix=metric_key_prefix,
1884
        )
Lysandre Debut's avatar
Lysandre Debut committed
1885

1886
        output.metrics.update(speed_metrics(metric_key_prefix, start_time, output.num_samples))
1887

1888
        self.log(output.metrics)
1889

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

Sylvain Gugger's avatar
Sylvain Gugger committed
1894
        self.control = self.callback_handler.on_evaluate(self.args, self.state, self.control, output.metrics)
1895
1896
1897

        self._memory_tracker.stop_and_update_metrics(output.metrics)

Julien Chaumond's avatar
Julien Chaumond committed
1898
1899
        return output.metrics

1900
    def predict(
Bhadresh Savani's avatar
Bhadresh Savani committed
1901
        self, test_dataset: Dataset, ignore_keys: Optional[List[str]] = None, metric_key_prefix: str = "test"
1902
    ) -> PredictionOutput:
Julien Chaumond's avatar
Julien Chaumond committed
1903
        """
1904
        Run prediction and returns predictions and potential metrics.
Julien Chaumond's avatar
Julien Chaumond committed
1905

Sylvain Gugger's avatar
Sylvain Gugger committed
1906
1907
        Depending on the dataset and your use case, your test dataset may contain labels. In that case, this method
        will also return metrics, like in :obj:`evaluate()`.
1908
1909
1910

        Args:
            test_dataset (:obj:`Dataset`):
1911
                Dataset to run the predictions on. If it is an :obj:`datasets.Dataset`, columns not accepted by the
1912
                ``model.forward()`` method are automatically removed. Has to implement the method :obj:`__len__`
1913
1914
1915
            ignore_keys (:obj:`Lst[str]`, `optional`):
                A list of keys in the output of your model (if it is a dictionary) that should be ignored when
                gathering predictions.
Bhadresh Savani's avatar
Bhadresh Savani committed
1916
            metric_key_prefix (:obj:`str`, `optional`, defaults to :obj:`"test"`):
1917
                An optional prefix to be used as the metrics key prefix. For example the metrics "bleu" will be named
Bhadresh Savani's avatar
Bhadresh Savani committed
1918
                "test_bleu" if the prefix is "test" (default)
1919

1920
1921
1922
1923
1924
1925
        .. note::

            If your predictions or labels have different sequence length (for instance because you're doing dynamic
            padding in a token classification task) the predictions will be padded (on the right) to allow for
            concatenation into one array. The padding index is -100.

Sylvain Gugger's avatar
Sylvain Gugger committed
1926
1927
1928
1929
1930
1931
        Returns: `NamedTuple` A namedtuple with the following keys:

            - 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
1932
        """
1933
1934
1935
        # memory metrics - must set up as early as possible
        self._memory_tracker.start()

Julien Chaumond's avatar
Julien Chaumond committed
1936
        test_dataloader = self.get_test_dataloader(test_dataset)
1937
        start_time = time.time()
1938

1939
1940
        eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
        output = eval_loop(
1941
1942
            test_dataloader, description="Prediction", ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix
        )
1943
        output.metrics.update(speed_metrics(metric_key_prefix, start_time, output.num_samples))
1944
1945
1946

        self._memory_tracker.stop_and_update_metrics(output.metrics)

1947
        return PredictionOutput(predictions=output.predictions, label_ids=output.label_ids, metrics=output.metrics)
Julien Chaumond's avatar
Julien Chaumond committed
1948

1949
    def evaluation_loop(
1950
1951
1952
1953
1954
        self,
        dataloader: DataLoader,
        description: str,
        prediction_loss_only: Optional[bool] = None,
        ignore_keys: Optional[List[str]] = None,
1955
        metric_key_prefix: str = "eval",
1956
    ) -> EvalLoopOutput:
Julien Chaumond's avatar
Julien Chaumond committed
1957
        """
1958
        Prediction/evaluation loop, shared by :obj:`Trainer.evaluate()` and :obj:`Trainer.predict()`.
Julien Chaumond's avatar
Julien Chaumond committed
1959
1960
1961

        Works both with or without labels.
        """
1962
1963
1964
        prediction_loss_only = (
            prediction_loss_only if prediction_loss_only is not None else self.args.prediction_loss_only
        )
Julien Chaumond's avatar
Julien Chaumond committed
1965

1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
        # if eval is called w/o train init deepspeed here
        if self.args.deepspeed and not self.deepspeed:

            # XXX: eval doesn't have `resume_from_checkpoint` arg but we should be able to do eval
            # from the checkpoint eventually
            deepspeed_engine, _, _ = deepspeed_init(self, num_training_steps=0, resume_from_checkpoint=None)
            self.model = deepspeed_engine.module
            self.model_wrapped = deepspeed_engine
            self.deepspeed = deepspeed_engine
            # XXX: we don't need optim/sched for inference, but this needs to be sorted out, since
            # for example the Z3-optimizer is a must for zero3 to work even for inference - what we
            # don't need is the deepspeed basic optimizer which is self.optimizer.optimizer
            deepspeed_engine.optimizer.optimizer = None
            deepspeed_engine.lr_scheduler = None
1980

1981
        model = self._wrap_model(self.model, training=False)
Julien Chaumond's avatar
Julien Chaumond committed
1982

1983
        # if full fp16 is wanted on eval and this ``evaluation`` or ``predict`` isn't called while
1984
        # ``train`` is running, halve it first and then put on device
1985
1986
1987
        if not self.is_in_train and self.args.fp16_full_eval:
            model = model.half().to(self.args.device)

1988
        batch_size = dataloader.batch_size
1989

1990
        logger.info(f"***** Running {description} *****")
1991
1992
1993
1994
        if isinstance(dataloader.dataset, collections.abc.Sized):
            logger.info(f"  Num examples = {self.num_examples(dataloader)}")
        else:
            logger.info("  Num examples: Unknown")
1995
        logger.info(f"  Batch size = {batch_size}")
1996

Julien Chaumond's avatar
Julien Chaumond committed
1997
1998
        model.eval()

1999
2000
2001
2002
        self.callback_handler.eval_dataloader = dataloader
        # Do this before wrapping.
        eval_dataset = dataloader.dataset

2003
        if is_torch_tpu_available():
2004
2005
            dataloader = pl.ParallelLoader(dataloader, [self.args.device]).per_device_loader(self.args.device)

2006
        if self.args.past_index >= 0:
2007
            self._past = None
2008

2009
2010
2011
2012
2013
2014
2015
2016
2017
2018
2019
2020
2021
        # Initialize containers
        # losses/preds/labels on GPU/TPU (accumulated for eval_accumulation_steps)
        losses_host = None
        preds_host = None
        labels_host = None
        # losses/preds/labels on CPU (final containers)
        all_losses = None
        all_preds = None
        all_labels = None
        # Will be useful when we have an iterable dataset so don't know its length.

        observed_num_examples = 0
        # Main evaluation loop
2022
        for step, inputs in enumerate(dataloader):
2023
2024
2025
2026
2027
2028
            # Update the observed num examples
            observed_batch_size = find_batch_size(inputs)
            if observed_batch_size is not None:
                observed_num_examples += observed_batch_size

            # Prediction step
2029
            loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)
2030
2031

            # Update containers on host
2032
            if loss is not None:
2033
                losses = self._nested_gather(loss.repeat(batch_size))
2034
                losses_host = losses if losses_host is None else torch.cat((losses_host, losses), dim=0)
2035
            if logits is not None:
2036
2037
                logits = self._pad_across_processes(logits)
                logits = self._nested_gather(logits)
2038
                preds_host = logits if preds_host is None else nested_concat(preds_host, logits, padding_index=-100)
2039
            if labels is not None:
2040
2041
                labels = self._pad_across_processes(labels)
                labels = self._nested_gather(labels)
2042
                labels_host = labels if labels_host is None else nested_concat(labels_host, labels, padding_index=-100)
Sylvain Gugger's avatar
Sylvain Gugger committed
2043
            self.control = self.callback_handler.on_prediction_step(self.args, self.state, self.control)
Julien Chaumond's avatar
Julien Chaumond committed
2044

2045
2046
            # Gather all tensors and put them back on the CPU if we have done enough accumulation steps.
            if self.args.eval_accumulation_steps is not None and (step + 1) % self.args.eval_accumulation_steps == 0:
2047
2048
2049
2050
2051
2052
2053
2054
2055
2056
2057
                if losses_host is not None:
                    losses = nested_numpify(losses_host)
                    all_losses = losses if all_losses is None else np.concatenate((all_losses, losses), axis=0)
                if preds_host is not None:
                    logits = nested_numpify(preds_host)
                    all_preds = logits if all_preds is None else nested_concat(all_preds, logits, padding_index=-100)
                if labels_host is not None:
                    labels = nested_numpify(labels_host)
                    all_labels = (
                        labels if all_labels is None else nested_concat(all_labels, labels, padding_index=-100)
                    )
2058
2059
2060
2061

                # Set back to None to begin a new accumulation
                losses_host, preds_host, labels_host = None, None, None

2062
2063
2064
        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
2065

2066
        # Gather all remaining tensors and put them back on the CPU
2067
2068
2069
2070
2071
2072
2073
2074
2075
2076
2077
2078
2079
2080
2081
2082
2083
2084
2085
2086
2087
2088
2089
2090
2091
2092
2093
2094
2095
2096
        if losses_host is not None:
            losses = nested_numpify(losses_host)
            all_losses = losses if all_losses is None else np.concatenate((all_losses, losses), axis=0)
        if preds_host is not None:
            logits = nested_numpify(preds_host)
            all_preds = logits if all_preds is None else nested_concat(all_preds, logits, padding_index=-100)
        if labels_host is not None:
            labels = nested_numpify(labels_host)
            all_labels = labels if all_labels is None else nested_concat(all_labels, labels, padding_index=-100)

        # Number of samples
        if not isinstance(eval_dataset, IterableDataset):
            num_samples = len(eval_dataset)
        elif isinstance(eval_dataset, IterableDatasetShard):
            num_samples = eval_dataset.num_examples
        else:
            num_samples = observed_num_examples

        # Number of losses has been rounded to a multiple of batch_size and in a distributed training, the number of
        # samplers has been rounded to a multiple of batch_size, so we truncate.
        if all_losses is not None:
            all_losses = all_losses[:num_samples]
        if all_preds is not None:
            all_preds = nested_truncate(all_preds, num_samples)
        if all_labels is not None:
            all_labels = nested_truncate(all_labels, num_samples)

        # Metrics!
        if self.compute_metrics is not None and all_preds is not None and all_labels is not None:
            metrics = self.compute_metrics(EvalPrediction(predictions=all_preds, label_ids=all_labels))
Julien Chaumond's avatar
Julien Chaumond committed
2097
2098
        else:
            metrics = {}
2099

2100
2101
2102
        # To be JSON-serializable, we need to remove numpy types or zero-d tensors
        metrics = denumpify_detensorize(metrics)

2103
2104
        if all_losses is not None:
            metrics[f"{metric_key_prefix}_loss"] = all_losses.mean().item()
2105

2106
        # Prefix all keys with metric_key_prefix + '_'
2107
        for key in list(metrics.keys()):
2108
2109
            if not key.startswith(f"{metric_key_prefix}_"):
                metrics[f"{metric_key_prefix}_{key}"] = metrics.pop(key)
Julien Chaumond's avatar
Julien Chaumond committed
2110

2111
        return EvalLoopOutput(predictions=all_preds, label_ids=all_labels, metrics=metrics, num_samples=num_samples)
2112

2113
    def _nested_gather(self, tensors, name=None):
2114
2115
2116
2117
2118
2119
2120
        """
        Gather value of `tensors` (tensor or list/tuple of nested tensors) and convert them to numpy before
        concatenating them to `gathered`
        """
        if tensors is None:
            return
        if is_torch_tpu_available():
2121
2122
            if name is None:
                name = "nested_gather"
2123
            tensors = nested_xla_mesh_reduce(tensors, name)
Sylvain Gugger's avatar
Sylvain Gugger committed
2124
2125
        elif is_sagemaker_mp_enabled():
            tensors = smp_gather(tensors)
2126
2127
        elif self.args.local_rank != -1:
            tensors = distributed_concat(tensors)
2128
        return tensors
2129

2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
2142
2143
2144
2145
2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
    # Copied from Accelerate.
    def _pad_across_processes(self, tensor, pad_index=-100):
        """
        Recursively pad the tensors in a nested list/tuple/dictionary of tensors from all devices to the same size so
        they can safely be gathered.
        """
        if isinstance(tensor, (list, tuple)):
            return type(tensor)(self._pad_across_processes(t, pad_index=pad_index) for t in tensor)
        elif isinstance(tensor, dict):
            return type(tensor)({k: self._pad_across_processes(v, pad_index=pad_index) for k, v in tensor.items()})
        elif not isinstance(tensor, torch.Tensor):
            raise TypeError(
                f"Can't pad the values of type {type(tensor)}, only of nested list/tuple/dicts of tensors."
            )

        if len(tensor.shape) < 2:
            return tensor
        # Gather all sizes
        size = torch.tensor(tensor.shape, device=tensor.device)[None]
        sizes = self._nested_gather(size).cpu()

        max_size = max(s[1] for s in sizes)
        if tensor.shape[1] == max_size:
            return tensor

        # Then pad to the maximum size
        old_size = tensor.shape
        new_size = list(old_size)
        new_size[1] = max_size
        new_tensor = tensor.new_zeros(tuple(new_size)) + pad_index
        new_tensor[:, : old_size[1]] = tensor
        return new_tensor
2162

2163
    def prediction_step(
2164
2165
2166
2167
2168
        self,
        model: nn.Module,
        inputs: Dict[str, Union[torch.Tensor, Any]],
        prediction_loss_only: bool,
        ignore_keys: Optional[List[str]] = None,
2169
    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
2170
2171
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
        """
        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.
2185
2186
2187
            ignore_keys (:obj:`Lst[str]`, `optional`):
                A list of keys in the output of your model (if it is a dictionary) that should be ignored when
                gathering predictions.
2188
2189

        Return:
2190
2191
            Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: A tuple with the loss,
            logits and labels (each being optional).
2192
        """
2193
        has_labels = all(inputs.get(k) is not None for k in self.label_names)
2194
        inputs = self._prepare_inputs(inputs)
2195
2196
2197
2198
2199
        if ignore_keys is None:
            if hasattr(self.model, "config"):
                ignore_keys = getattr(self.model.config, "keys_to_ignore_at_inference", [])
            else:
                ignore_keys = []
2200

2201
2202
2203
2204
2205
2206
2207
2208
        # labels may be popped when computing the loss (label smoothing for instance) so we grab them first.
        if has_labels:
            labels = nested_detach(tuple(inputs.get(name) for name in self.label_names))
            if len(labels) == 1:
                labels = labels[0]
        else:
            labels = None

2209
        with torch.no_grad():
Sylvain Gugger's avatar
Sylvain Gugger committed
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
            if is_sagemaker_mp_enabled():
                raw_outputs = smp_forward_only(model, inputs)
                if has_labels:
                    if isinstance(raw_outputs, dict):
                        loss_mb = raw_outputs["loss"]
                        logits_mb = tuple(v for k, v in raw_outputs.items() if k not in ignore_keys + ["loss"])
                    else:
                        loss_mb = raw_outputs[0]
                        logits_mb = raw_outputs[1:]

                    loss = loss_mb.reduce_mean().detach().cpu()
                    logits = smp_nested_concat(logits_mb)
2222
                else:
Sylvain Gugger's avatar
Sylvain Gugger committed
2223
2224
2225
2226
2227
2228
                    loss = None
                    if isinstance(raw_outputs, dict):
                        logits_mb = tuple(v for k, v in raw_outputs.items() if k not in ignore_keys)
                    else:
                        logits_mb = raw_outputs
                    logits = smp_nested_concat(logits_mb)
2229
            else:
Sylvain Gugger's avatar
Sylvain Gugger committed
2230
2231
2232
2233
2234
2235
2236
                if has_labels:
                    loss, outputs = self.compute_loss(model, inputs, return_outputs=True)
                    loss = loss.mean().detach()
                    if isinstance(outputs, dict):
                        logits = tuple(v for k, v in outputs.items() if k not in ignore_keys + ["loss"])
                    else:
                        logits = outputs[1:]
2237
                else:
Sylvain Gugger's avatar
Sylvain Gugger committed
2238
2239
2240
2241
2242
2243
2244
2245
2246
2247
2248
2249
2250
                    loss = None
                    if self.use_amp:
                        with autocast():
                            outputs = model(**inputs)
                    else:
                        outputs = model(**inputs)
                    if isinstance(outputs, dict):
                        logits = tuple(v for k, v in outputs.items() if k not in ignore_keys)
                    else:
                        logits = outputs
                    # TODO: this needs to be fixed and made cleaner later.
                    if self.args.past_index >= 0:
                        self._past = outputs[self.args.past_index - 1]
2251
2252
2253
2254

        if prediction_loss_only:
            return (loss, None, None)

2255
        logits = nested_detach(logits)
Sylvain Gugger's avatar
Sylvain Gugger committed
2256
2257
2258
2259
        if len(logits) == 1:
            logits = logits[0]

        return (loss, logits, labels)
2260
2261
2262

    def floating_point_ops(self, inputs: Dict[str, Union[torch.Tensor, Any]]):
        """
Sylvain Gugger's avatar
Sylvain Gugger committed
2263
2264
2265
        For models that inherit from :class:`~transformers.PreTrainedModel`, uses that method to compute the number of
        floating point operations for every backward + forward pass. If using another model, either implement such a
        method in the model or subclass and override this method.
2266
2267
2268
2269
2270
2271
2272
2273

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

        Returns:
            :obj:`int`: The number of floating-point operations.
        """
2274
2275
        if hasattr(self.model, "floating_point_ops"):
            return self.model.floating_point_ops(inputs)
2276
2277
        else:
            return 0
2278

Sylvain Gugger's avatar
Sylvain Gugger committed
2279
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
2308
2309
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327
2328
2329
2330
2331
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
2343
    def push_to_hub(
        self,
        save_directory: Optional[str] = None,
        repo_name: Optional[str] = None,
        repo_url: Optional[str] = None,
        commit_message: Optional[str] = "add model",
        organization: Optional[str] = None,
        private: bool = None,
        use_auth_token: Optional[Union[bool, str]] = None,
    ):
        """
        Upload `self.model` to the 馃 model hub.

        Parameters:
            save_directory (:obj:`str` or :obj:`os.PathLike`):
                Folder containing the model weights and config. Will default to :obj:`self.args.output_dir`.
            repo_name (:obj:`str`, `optional`):
                Repository name for your model or tokenizer in the hub. If not specified, the repository name will be
                the stem of :obj:`save_directory`.
            repo_url (:obj:`str`, `optional`):
                Specify this in case you want to push to an existing repository in the hub. If unspecified, a new
                repository will be created in your namespace (unless you specify an :obj:`organization`) with
                :obj:`repo_name`.
            commit_message (:obj:`str`, `optional`, defaults to :obj:`"add model"`):
                Message to commit while pushing.
            organization (:obj:`str`, `optional`):
                Organization in which you want to push your model or tokenizer (you must be a member of this
                organization).
            private (:obj:`bool`, `optional`):
                Whether or not the repository created should be private (requires a paying subscription).
            use_auth_token (:obj:`bool` or :obj:`str`, `optional`):
                The token to use as HTTP bearer authorization for remote files. If :obj:`True`, will use the token
                generated when running :obj:`transformers-cli login` (stored in :obj:`~/.huggingface`). Will default to
                :obj:`True` if :obj:`repo_url` is not specified.

        Returns:
            The url of the commit of your model in the given repository.
        """
        if not self.is_world_process_zero():
            return

        if not isinstance(unwrap_model(self.model), PushToHubMixin):
            raise ValueError(
                "The `upload_model_to_hub` method only works for models that inherit from `PushToHubMixin` models."
            )
        if save_directory is None:
            save_directory = self.args.output_dir

        # To avoid pushing all checkpoints, we just copy all the files in save_directory in a tmp dir.
        with tempfile.TemporaryDirectory() as tmp_dir:
            for f in os.listdir(save_directory):
                fname = os.path.join(save_directory, f)
                if os.path.isfile(fname):
                    shutil.copy(fname, os.path.join(tmp_dir, f))

            return unwrap_model(self.model)._push_to_hub(
                save_directory=tmp_dir,
                repo_name=repo_name,
                repo_url=repo_url,
                commit_message=commit_message,
                organization=organization,
                private=private,
                use_auth_token=use_auth_token,
            )

2344
2345
2346
2347
2348
2349
2350
2351
2352
2353
2354
2355
2356
2357
2358
2359
2360
2361
2362
2363
2364
2365
2366
2367
2368
2369
2370
2371
2372
2373
2374
2375
2376
2377
2378
2379
2380
2381
2382
2383
2384
    #
    # Deprecated code
    #

    def prediction_loop(
        self,
        dataloader: DataLoader,
        description: str,
        prediction_loss_only: Optional[bool] = None,
        ignore_keys: Optional[List[str]] = None,
        metric_key_prefix: str = "eval",
    ) -> PredictionOutput:
        """
        Prediction/evaluation loop, shared by :obj:`Trainer.evaluate()` and :obj:`Trainer.predict()`.

        Works both with or without labels.
        """
        if not isinstance(dataloader.dataset, collections.abc.Sized):
            raise ValueError("dataset must implement __len__")
        prediction_loss_only = (
            prediction_loss_only if prediction_loss_only is not None else self.args.prediction_loss_only
        )

        # if eval is called w/o train init deepspeed here
        if self.args.deepspeed and not self.deepspeed:

            # XXX: eval doesn't have `resume_from_checkpoint` arg but we should be able to do eval
            # from the checkpoint eventually
            deepspeed_engine, _, _ = deepspeed_init(self, num_training_steps=0, resume_from_checkpoint=None)
            self.model = deepspeed_engine.module
            self.model_wrapped = deepspeed_engine
            self.deepspeed = deepspeed_engine
            # XXX: we don't need optim/sched for inference, but this needs to be sorted out, since
            # for example the Z3-optimizer is a must for zero3 to work even for inference - what we
            # don't need is the deepspeed basic optimizer which is self.optimizer.optimizer
            deepspeed_engine.optimizer.optimizer = None
            deepspeed_engine.lr_scheduler = None

        model = self._wrap_model(self.model, training=False)

        # if full fp16 is wanted on eval and this ``evaluation`` or ``predict`` isn't called while
2385
        # ``train`` is running, halve it first and then put on device
2386
2387
2388
2389
2390
2391
2392
2393
2394
2395
2396
2397
2398
2399
2400
2401
2402
2403
2404
2405
2406
2407
2408
2409
2410
2411
2412
2413
2414
2415
2416
2417
2418
2419
2420
2421
2422
2423
2424
2425
2426
2427
2428
2429
2430
2431
2432
2433
2434
2435
2436
2437
2438
2439
2440
2441
2442
2443
2444
2445
2446
2447
2448
2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
2461
2462
2463
2464
2465
2466
2467
2468
2469
2470
2471
2472
2473
2474
2475
2476
2477
2478
2479
2480
2481
2482
2483
2484
2485
2486
2487
        if not self.is_in_train and self.args.fp16_full_eval:
            model = model.half().to(self.args.device)

        batch_size = dataloader.batch_size
        num_examples = self.num_examples(dataloader)
        logger.info(f"***** Running {description} *****")
        logger.info(f"  Num examples = {num_examples}")
        logger.info(f"  Batch size = {batch_size}")
        losses_host: torch.Tensor = None
        preds_host: Union[torch.Tensor, List[torch.Tensor]] = None
        labels_host: Union[torch.Tensor, List[torch.Tensor]] = None

        world_size = max(1, self.args.world_size)

        eval_losses_gatherer = DistributedTensorGatherer(world_size, num_examples, make_multiple_of=batch_size)
        if not prediction_loss_only:
            # The actual number of eval_sample can be greater than num_examples in distributed settings (when we pass
            # a batch size to the sampler)
            make_multiple_of = None
            if hasattr(dataloader, "sampler") and isinstance(dataloader.sampler, SequentialDistributedSampler):
                make_multiple_of = dataloader.sampler.batch_size
            preds_gatherer = DistributedTensorGatherer(world_size, num_examples, make_multiple_of=make_multiple_of)
            labels_gatherer = DistributedTensorGatherer(world_size, num_examples, make_multiple_of=make_multiple_of)

        model.eval()

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

        if self.args.past_index >= 0:
            self._past = None

        self.callback_handler.eval_dataloader = dataloader

        for step, inputs in enumerate(dataloader):
            loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)
            if loss is not None:
                losses = loss.repeat(batch_size)
                losses_host = losses if losses_host is None else torch.cat((losses_host, losses), dim=0)
            if logits is not None:
                preds_host = logits if preds_host is None else nested_concat(preds_host, logits, padding_index=-100)
            if labels is not None:
                labels_host = labels if labels_host is None else nested_concat(labels_host, labels, padding_index=-100)
            self.control = self.callback_handler.on_prediction_step(self.args, self.state, self.control)

            # Gather all tensors and put them back on the CPU if we have done enough accumulation steps.
            if self.args.eval_accumulation_steps is not None and (step + 1) % self.args.eval_accumulation_steps == 0:
                eval_losses_gatherer.add_arrays(self._gather_and_numpify(losses_host, "eval_losses"))
                if not prediction_loss_only:
                    preds_gatherer.add_arrays(self._gather_and_numpify(preds_host, "eval_preds"))
                    labels_gatherer.add_arrays(self._gather_and_numpify(labels_host, "eval_label_ids"))

                # Set back to None to begin a new accumulation
                losses_host, preds_host, labels_host = None, None, None

        if self.args.past_index and hasattr(self, "_past"):
            # Clean the state at the end of the evaluation loop
            delattr(self, "_past")

        # Gather all remaining tensors and put them back on the CPU
        eval_losses_gatherer.add_arrays(self._gather_and_numpify(losses_host, "eval_losses"))
        if not prediction_loss_only:
            preds_gatherer.add_arrays(self._gather_and_numpify(preds_host, "eval_preds"))
            labels_gatherer.add_arrays(self._gather_and_numpify(labels_host, "eval_label_ids"))

        eval_loss = eval_losses_gatherer.finalize()
        preds = preds_gatherer.finalize() if not prediction_loss_only else None
        label_ids = labels_gatherer.finalize() if not prediction_loss_only else None

        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 = {}

        # To be JSON-serializable, we need to remove numpy types or zero-d tensors
        metrics = denumpify_detensorize(metrics)

        if eval_loss is not None:
            metrics[f"{metric_key_prefix}_loss"] = eval_loss.mean().item()

        # Prefix all keys with metric_key_prefix + '_'
        for key in list(metrics.keys()):
            if not key.startswith(f"{metric_key_prefix}_"):
                metrics[f"{metric_key_prefix}_{key}"] = metrics.pop(key)

        return PredictionOutput(predictions=preds, label_ids=label_ids, metrics=metrics)

    def _gather_and_numpify(self, tensors, name):
        """
        Gather value of `tensors` (tensor or list/tuple of nested tensors) and convert them to numpy before
        concatenating them to `gathered`
        """
        if tensors is None:
            return
        if is_torch_tpu_available():
            tensors = nested_xla_mesh_reduce(tensors, name)
        elif is_sagemaker_mp_enabled():
            tensors = smp_gather(tensors)
        elif self.args.local_rank != -1:
            tensors = distributed_concat(tensors)

        return nested_numpify(tensors)