"tests/utils/test_tokenization_utils.py" did not exist on "1b652295c5940a45079c9056864009418c5a8054"
trainer.py 123 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
import os
23
import random
Julien Chaumond's avatar
Julien Chaumond committed
24
25
import re
import shutil
26
import sys
Sylvain Gugger's avatar
Sylvain Gugger committed
27
import tempfile
28
import time
29
import warnings
30
from logging import StreamHandler
Julien Chaumond's avatar
Julien Chaumond committed
31
from pathlib import Path
32
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
Julien Chaumond's avatar
Julien Chaumond committed
33

34
35
from tqdm.auto import tqdm

Julien Chaumond's avatar
Julien Chaumond committed
36

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

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

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


132
_is_torch_generator_available = False
133
_is_native_amp_available = False
134

Sylvain Gugger's avatar
Sylvain Gugger committed
135
DEFAULT_CALLBACKS = [DefaultFlowCallback]
136
DEFAULT_PROGRESS_CALLBACK = ProgressCallback
Sylvain Gugger's avatar
Sylvain Gugger committed
137

138
139
140
141
if is_in_notebook():
    from .utils.notebook import NotebookProgressCallback

    DEFAULT_PROGRESS_CALLBACK = NotebookProgressCallback
142

143
144
if is_apex_available():
    from apex import amp
145

146
if version.parse(torch.__version__) >= version.parse("1.6"):
147
    _is_torch_generator_available = True
148
    _is_native_amp_available = True
149
    from torch.cuda.amp import autocast
Julien Chaumond's avatar
Julien Chaumond committed
150

151
152
if is_datasets_available():
    import datasets
Julien Chaumond's avatar
Julien Chaumond committed
153

154
if is_torch_tpu_available():
Lysandre Debut's avatar
Lysandre Debut committed
155
156
157
158
    import torch_xla.core.xla_model as xm
    import torch_xla.debug.metrics as met
    import torch_xla.distributed.parallel_loader as pl

159
if is_fairscale_available():
160
    dep_version_check("fairscale")
161
    import fairscale
162
    from fairscale.nn.data_parallel import FullyShardedDataParallel as FullyShardedDDP
163
    from fairscale.nn.data_parallel import ShardedDataParallel as ShardedDDP
164
    from fairscale.nn.wrap import auto_wrap
165
166
167
    from fairscale.optim import OSS
    from fairscale.optim.grad_scaler import ShardedGradScaler

Sylvain Gugger's avatar
Sylvain Gugger committed
168
if is_sagemaker_dp_enabled():
Sylvain Gugger's avatar
Sylvain Gugger committed
169
170
171
172
    import smdistributed.dataparallel.torch.distributed as dist
    from smdistributed.dataparallel.torch.parallel.distributed import DistributedDataParallel as DDP
else:
    import torch.distributed as dist
173

Sylvain Gugger's avatar
Sylvain Gugger committed
174
175
176
177
178
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

179
180
181
182
if is_training_run_on_sagemaker():
    logging.add_handler(StreamHandler(sys.stdout))


183
184
185
if TYPE_CHECKING:
    import optuna

Lysandre Debut's avatar
Lysandre Debut committed
186
logger = logging.get_logger(__name__)
Julien Chaumond's avatar
Julien Chaumond committed
187
188
189
190


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

    Args:
Sylvain Gugger's avatar
Sylvain Gugger committed
194
        model (:class:`~transformers.PreTrainedModel` or :obj:`torch.nn.Module`, `optional`):
195
            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
196
197
198
199
200
201

            .. 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.
202
        args (:class:`~transformers.TrainingArguments`, `optional`):
Sylvain Gugger's avatar
Sylvain Gugger committed
203
204
205
            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.
206
        data_collator (:obj:`DataCollator`, `optional`):
Sylvain Gugger's avatar
Sylvain Gugger committed
207
208
209
            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.
210
        train_dataset (:obj:`torch.utils.data.dataset.Dataset` or :obj:`torch.utils.data.dataset.IterableDataset`, `optional`):
211
            The dataset to use for training. If it is an :obj:`datasets.Dataset`, columns not accepted by the
212
            ``model.forward()`` method are automatically removed.
213
214
215

            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
216
            :obj:`generator` that is a :obj:`torch.Generator` for the randomization that must be identical on all
217
218
            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
219
        eval_dataset (:obj:`torch.utils.data.dataset.Dataset`, `optional`):
220
             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
221
             ``model.forward()`` method are automatically removed.
222
223
224
225
        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.
226
227
228
        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.
229

Sylvain Gugger's avatar
Sylvain Gugger committed
230
231
            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
232
233
            layers, dropout probabilities etc).
        compute_metrics (:obj:`Callable[[EvalPrediction], Dict]`, `optional`):
234
            The function that will be used to compute metrics at evaluation. Must take a
Sylvain Gugger's avatar
Sylvain Gugger committed
235
236
237
238
            :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
239
240

            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
241
        optimizers (:obj:`Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR`, `optional`): A tuple
Sylvain Gugger's avatar
Sylvain Gugger committed
242
            containing the optimizer and the scheduler to use. Will default to an instance of
243
            :class:`~transformers.AdamW` on your model and a scheduler given by
Sylvain Gugger's avatar
Sylvain Gugger committed
244
            :func:`~transformers.get_linear_schedule_with_warmup` controlled by :obj:`args`.
245

246
247
248
249
250
251
252
253
254
255
    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).
256
257
258
        - **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` .
259
260
        - **is_in_train** -- Whether or not a model is currently running ``train`` (e.g. when ``evaluate`` is called
          while in ``train``)
261

Julien Chaumond's avatar
Julien Chaumond committed
262
263
    """

264
    from .trainer_pt_utils import _get_learning_rate, log_metrics, metrics_format, save_metrics, save_state
265

Julien Chaumond's avatar
Julien Chaumond committed
266
267
    def __init__(
        self,
Sylvain Gugger's avatar
Sylvain Gugger committed
268
        model: Union[PreTrainedModel, torch.nn.Module] = None,
269
        args: TrainingArguments = None,
Julien Chaumond's avatar
Julien Chaumond committed
270
271
272
        data_collator: Optional[DataCollator] = None,
        train_dataset: Optional[Dataset] = None,
        eval_dataset: Optional[Dataset] = None,
273
        tokenizer: Optional[PreTrainedTokenizerBase] = None,
274
        model_init: Callable[[], PreTrainedModel] = None,
Julien Chaumond's avatar
Julien Chaumond committed
275
        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
Sylvain Gugger's avatar
Sylvain Gugger committed
276
        callbacks: Optional[List[TrainerCallback]] = None,
277
        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
Julien Chaumond's avatar
Julien Chaumond committed
278
    ):
Sylvain Gugger's avatar
Sylvain Gugger committed
279
        if args is None:
280
281
282
            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
283
284
285
        self.args = args
        # Seed must be set before instantiating the model when using model
        set_seed(self.args.seed)
286
        self.hp_name = None
287
        self.deepspeed = None
288
        self.is_in_train = False
289

290
291
292
293
        # memory metrics - must set up as early as possible
        self._memory_tracker = TrainerMemoryTracker(self.args.skip_memory_metrics)
        self._memory_tracker.start()

294
295
296
        # force device and distributed setup init explicitly
        args._setup_devices

297
298
299
300
301
302
303
304
305
306
307
308
309
310
        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
311

312
313
314
315
316
        if hasattr(model, "is_parallelizable") and model.is_parallelizable and model.model_parallel:
            self.is_model_parallel = True
        else:
            self.is_model_parallel = False

317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
        # 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

341
        # one place to sort out whether to place the model on device or not
342
343
344
345
346
347
        # 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
348
        self.place_model_on_device = args.place_model_on_device
349
350
        if (
            self.is_model_parallel
351
            or args.deepspeed
352
353
354
            or (args.fp16_full_eval and not args.do_train)
            or (self.sharded_ddp in [ShardedDDPOption.ZERO_DP_2, ShardedDDPOption.ZERO_DP_3])
        ):
355
356
            self.place_model_on_device = False

357
358
        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
359
360
        self.train_dataset = train_dataset
        self.eval_dataset = eval_dataset
361
        self.tokenizer = tokenizer
362

363
        if self.place_model_on_device:
364
            model = model.to(args.device)
Stas Bekman's avatar
Stas Bekman committed
365
366
367

        # Force n_gpu to 1 to avoid DataParallel as MP will manage the GPUs
        if self.is_model_parallel:
368
            self.args._n_gpu = 1
369
370
371
372
373

        # 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
374
        self.compute_metrics = compute_metrics
375
        self.optimizer, self.lr_scheduler = optimizers
Sylvain Gugger's avatar
Sylvain Gugger committed
376
377
378
379
380
        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."
            )
381
382
        default_callbacks = DEFAULT_CALLBACKS + get_reporting_integration_callbacks(self.args.report_to)
        callbacks = default_callbacks if callbacks is None else default_callbacks + callbacks
383
384
385
        self.callback_handler = CallbackHandler(
            callbacks, self.model, self.tokenizer, self.optimizer, self.lr_scheduler
        )
386
        self.add_callback(PrinterCallback if self.args.disable_tqdm else DEFAULT_PROGRESS_CALLBACK)
Sylvain Gugger's avatar
Sylvain Gugger committed
387

388
389
390
        # 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
391
        # Create output directory if needed
392
        if self.is_world_process_zero():
Julien Chaumond's avatar
Julien Chaumond committed
393
            os.makedirs(self.args.output_dir, exist_ok=True)
394
        if not callable(self.data_collator) and callable(getattr(self.data_collator, "collate_batch", None)):
Sylvain Gugger's avatar
Sylvain Gugger committed
395
            raise ValueError("The `data_collator` should be a simple callable (function, class with `__call__`).")
396

397
398
399
400
401
402
        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")

403
        self._signature_columns = None
404

405
406
407
        # Mixed precision setup
        self.use_apex = False
        self.use_amp = False
408
409
        self.fp16_backend = None

410
411
        if args.fp16:
            if args.fp16_backend == "auto":
412
                self.fp16_backend = "amp" if _is_native_amp_available else "apex"
413
            else:
414
415
                self.fp16_backend = args.fp16_backend
            logger.info(f"Using {self.fp16_backend} fp16 backend")
416

417
418
        if args.fp16 and not args.deepspeed:  # deepspeed manages its own fp16
            if self.fp16_backend == "amp":
419
                self.use_amp = True
420
421
422
423
424
425
                if is_sagemaker_mp_enabled():
                    self.scaler = smp.amp.GradScaler()
                elif self.sharded_ddp is not None:
                    self.scaler = ShardedGradScaler()
                else:
                    self.scaler = torch.cuda.amp.GradScaler()
426
427
428
429
430
431
432
            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

433
434
435
436
437
438
439
        # FP16 + model parallelism in SageMaker: gradient clipping does not work for now so we raise a helpful error.
        if is_sagemaker_mp_enabled() and self.use_amp and args.max_grad_norm is not None and args.max_grad_norm > 0:
            raise ValueError(
                "SageMaker Model Parallelism in mixed precision mode does not support gradient clipping yet. Pass "
                "along 'max_grad_norm': 0 in your hyperparameters."
            )

Sylvain Gugger's avatar
Sylvain Gugger committed
440
441
442
443
444
445
        # Label smoothing
        if self.args.label_smoothing_factor != 0:
            self.label_smoother = LabelSmoother(epsilon=self.args.label_smoothing_factor)
        else:
            self.label_smoother = None

446
        self.state = TrainerState()
Sylvain Gugger's avatar
Sylvain Gugger committed
447
        self.control = TrainerControl()
448
449
450
        # 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
451
        self.hp_search_backend = None
452
        self.use_tune_checkpoints = False
453
        default_label_names = (
454
            ["start_positions", "end_positions"]
455
            if type(self.model).__name__ in MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES.values()
456
457
458
            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
459
460
        self.control = self.callback_handler.on_init_end(self.args, self.state, self.control)

461
462
463
        # very last
        self._memory_tracker.stop_and_update_metrics()

Sylvain Gugger's avatar
Sylvain Gugger committed
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
    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
501

502
    def _remove_unused_columns(self, dataset: "datasets.Dataset", description: Optional[str] = None):
503
        if not self.args.remove_unused_columns:
504
            return dataset
505
506
507
508
509
510
511
512
        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))
513
514
515
516
517
518
        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)}."
            )
519

520
521
522
523
524
525
526
        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)
527

528
    def _get_train_sampler(self) -> Optional[torch.utils.data.sampler.Sampler]:
529
        if not isinstance(self.train_dataset, collections.abc.Sized):
530
            return None
531

532
533
534
535
536
        generator = None
        if self.args.world_size <= 1 and _is_torch_generator_available:
            generator = torch.Generator()
            generator.manual_seed(int(torch.empty((), dtype=torch.int64).random_().item()))

537
538
        # Build the sampler.
        if self.args.group_by_length:
539
540
541
542
543
544
545
546
            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
547
            model_input_name = self.tokenizer.model_input_names[0] if self.tokenizer is not None else None
548
            if self.args.world_size <= 1:
549
                return LengthGroupedSampler(
550
551
552
553
554
                    self.train_dataset,
                    self.args.train_batch_size,
                    lengths=lengths,
                    model_input_name=model_input_name,
                    generator=generator,
555
                )
556
557
            else:
                return DistributedLengthGroupedSampler(
558
559
                    self.train_dataset,
                    self.args.train_batch_size,
560
561
                    num_replicas=self.args.world_size,
                    rank=self.args.process_index,
562
                    lengths=lengths,
563
                    model_input_name=model_input_name,
564
                    seed=self.args.seed,
565
566
567
                )

        else:
568
            if self.args.world_size <= 1:
569
570
                if _is_torch_generator_available:
                    return RandomSampler(self.train_dataset, generator=generator)
571
                return RandomSampler(self.train_dataset)
Sylvain Gugger's avatar
Sylvain Gugger committed
572
573
574
575
            elif (
                self.args.parallel_mode in [ParallelMode.TPU, ParallelMode.SAGEMAKER_MODEL_PARALLEL]
                and not self.args.dataloader_drop_last
            ):
576
577
578
579
580
581
                # 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,
582
                    seed=self.args.seed,
583
                )
584
            else:
585
                return DistributedSampler(
586
587
588
589
                    self.train_dataset,
                    num_replicas=self.args.world_size,
                    rank=self.args.process_index,
                    seed=self.args.seed,
590
                )
591
592
593
594
595

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

Sylvain Gugger's avatar
Sylvain Gugger committed
596
597
        Will use no sampler if :obj:`self.train_dataset` does not implement :obj:`__len__`, a random sampler (adapted
        to distributed training if necessary) otherwise.
598
599
600
601
602

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

604
605
606
607
608
        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):
609
610
            if self.args.world_size > 1:
                train_dataset = IterableDatasetShard(
611
                    train_dataset,
612
613
614
615
616
                    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,
                )
617

618
619
620
621
622
623
624
625
            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,
            )

626
627
628
        train_sampler = self._get_train_sampler()

        return DataLoader(
629
            train_dataset,
Julien Chaumond's avatar
Julien Chaumond committed
630
631
            batch_size=self.args.train_batch_size,
            sampler=train_sampler,
632
            collate_fn=self.data_collator,
Setu Shah's avatar
Setu Shah committed
633
            drop_last=self.args.dataloader_drop_last,
Chady Kamar's avatar
Chady Kamar committed
634
            num_workers=self.args.dataloader_num_workers,
635
            pin_memory=self.args.dataloader_pin_memory,
Julien Chaumond's avatar
Julien Chaumond committed
636
637
        )

638
    def _get_eval_sampler(self, eval_dataset: Dataset) -> Optional[torch.utils.data.sampler.Sampler]:
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
        # 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
661
662
                eval_dataset,
                batch_size=self.args.per_device_eval_batch_size,
663
664
                num_processes=self.args.world_size,
                process_index=self.args.process_index,
Sylvain Gugger's avatar
Sylvain Gugger committed
665
            )
Lysandre Debut's avatar
Lysandre Debut committed
666

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

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

673
        Args:
674
            eval_dataset (:obj:`torch.utils.data.dataset.Dataset`, `optional`):
675
                If provided, will override :obj:`self.eval_dataset`. If it is an :obj:`datasets.Dataset`, columns not
676
                accepted by the ``model.forward()`` method are automatically removed. It must implement :obj:`__len__`.
677
        """
Julien Chaumond's avatar
Julien Chaumond committed
678
679
        if eval_dataset is None and self.eval_dataset is None:
            raise ValueError("Trainer: evaluation requires an eval_dataset.")
680
        eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset
681

682
683
684
        if is_datasets_available() and isinstance(eval_dataset, datasets.Dataset):
            eval_dataset = self._remove_unused_columns(eval_dataset, description="evaluation")

685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
        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,
            )

702
        eval_sampler = self._get_eval_sampler(eval_dataset)
703

704
        return DataLoader(
705
            eval_dataset,
706
            sampler=eval_sampler,
Julien Chaumond's avatar
Julien Chaumond committed
707
            batch_size=self.args.eval_batch_size,
708
            collate_fn=self.data_collator,
Setu Shah's avatar
Setu Shah committed
709
            drop_last=self.args.dataloader_drop_last,
Chady Kamar's avatar
Chady Kamar committed
710
            num_workers=self.args.dataloader_num_workers,
711
            pin_memory=self.args.dataloader_pin_memory,
Julien Chaumond's avatar
Julien Chaumond committed
712
713
714
        )

    def get_test_dataloader(self, test_dataset: Dataset) -> DataLoader:
715
716
717
        """
        Returns the test :class:`~torch.utils.data.DataLoader`.

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

720
        Args:
721
            test_dataset (:obj:`torch.utils.data.dataset.Dataset`, `optional`):
722
                The test dataset to use. If it is an :obj:`datasets.Dataset`, columns not accepted by the
723
                ``model.forward()`` method are automatically removed. It must implement :obj:`__len__`.
724
        """
725
        if is_datasets_available() and isinstance(test_dataset, datasets.Dataset):
726
            test_dataset = self._remove_unused_columns(test_dataset, description="test")
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744

        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,
            )

745
        test_sampler = self._get_eval_sampler(test_dataset)
Lysandre Debut's avatar
Lysandre Debut committed
746

747
748
        # We use the same batch_size as for eval.
        return DataLoader(
Julien Chaumond's avatar
Julien Chaumond committed
749
            test_dataset,
750
            sampler=test_sampler,
Julien Chaumond's avatar
Julien Chaumond committed
751
            batch_size=self.args.eval_batch_size,
752
            collate_fn=self.data_collator,
753
            drop_last=self.args.dataloader_drop_last,
754
            pin_memory=self.args.dataloader_pin_memory,
Julien Chaumond's avatar
Julien Chaumond committed
755
        )
Lysandre Debut's avatar
Lysandre Debut committed
756

757
    def create_optimizer_and_scheduler(self, num_training_steps: int):
758
759
760
        """
        Setup the optimizer and the learning rate scheduler.

761
762
763
764
765
766
767
768
769
770
771
        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.

772
        We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the
773
        Trainer's init through :obj:`optimizers`, or subclass and override this method in a subclass.
774
        """
775
        if self.optimizer is None:
776
777
            decay_parameters = get_parameter_names(self.model, [torch.nn.LayerNorm])
            decay_parameters = [name for name in decay_parameters if "bias" not in name]
778
779
            optimizer_grouped_parameters = [
                {
780
                    "params": [p for n, p in self.model.named_parameters() if n in decay_parameters],
781
782
783
                    "weight_decay": self.args.weight_decay,
                },
                {
784
                    "params": [p for n, p in self.model.named_parameters() if n not in decay_parameters],
785
786
787
                    "weight_decay": 0.0,
                },
            ]
Sylvain Gugger's avatar
Sylvain Gugger committed
788
789
790
791
792
793
794
795
796
797
798
            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
799
            if self.sharded_ddp == ShardedDDPOption.SIMPLE:
800
801
                self.optimizer = OSS(
                    params=optimizer_grouped_parameters,
Sylvain Gugger's avatar
Sylvain Gugger committed
802
803
                    optim=optimizer_cls,
                    **optimizer_kwargs,
804
805
                )
            else:
Sylvain Gugger's avatar
Sylvain Gugger committed
806
807
                self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)

Sylvain Gugger's avatar
Sylvain Gugger committed
808
809
810
        if is_sagemaker_mp_enabled():
            self.optimizer = smp.DistributedOptimizer(self.optimizer)

811
812
813
814
815
816
817
    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.
        """
818
        if self.lr_scheduler is None:
819
820
821
822
823
824
            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
825
826
827
            self.lr_scheduler = get_scheduler(
                self.args.lr_scheduler_type,
                self.optimizer,
828
                num_warmup_steps=warmup_steps,
Sylvain Gugger's avatar
Sylvain Gugger committed
829
                num_training_steps=num_training_steps,
830
            )
Julien Chaumond's avatar
Julien Chaumond committed
831

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

836
        Will raise an exception if the underlying dataset does not implement method :obj:`__len__`
837
        """
838
        return len(dataloader.dataset)
839

840
    def _hp_search_setup(self, trial: Union["optuna.Trial", Dict[str, Any]]):
Patrick von Platen's avatar
Patrick von Platen committed
841
        """HP search setup code"""
842
843
        self._trial = trial

844
845
        if self.hp_search_backend is None or trial is None:
            return
846
847
848
849
850
        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)
851

852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
        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
870
        self.objective = self.compute_objective(metrics.copy())
871
        if self.hp_search_backend == HPSearchBackend.OPTUNA:
872
873
            import optuna

874
875
876
877
            trial.report(self.objective, epoch)
            if trial.should_prune():
                raise optuna.TrialPruned()
        elif self.hp_search_backend == HPSearchBackend.RAY:
878
879
            from ray import tune

880
            if self.control.should_save:
881
                self._tune_save_checkpoint()
882
883
            tune.report(objective=self.objective, **metrics)

884
    def _tune_save_checkpoint(self):
885
886
        from ray import tune

887
888
        if not self.use_tune_checkpoints:
            return
889
        with tune.checkpoint_dir(step=self.state.global_step) as checkpoint_dir:
890
            output_dir = os.path.join(checkpoint_dir, f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}")
891
            self.save_model(output_dir)
892
            if self.is_world_process_zero():
893
                self.state.save_to_json(os.path.join(output_dir, "trainer_state.json"))
894
895
896
                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"))

897
898
899
900
901
902
903
    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:
904
905
906
907
            raise RuntimeError("model_init should have 0 or 1 argument.")

        if model is None:
            raise RuntimeError("model_init should not return None.")
908
909
910

        return model

911
    def _wrap_model(self, model, training=True):
Sylvain Gugger's avatar
Sylvain Gugger committed
912
913
914
915
916
917
        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)

918
919
        # already initialized its own DDP and AMP
        if self.deepspeed:
920
            return self.deepspeed
921

922
923
924
925
        # 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

926
927
928
929
930
931
932
933
934
935
936
937
938
939
        # 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)
940
941
942
943
944
945
946
947
948
        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.
949
950
                if ShardedDDPOption.AUTO_WRAP in self.args.sharded_ddp:
                    model = auto_wrap(model)
951
                self.model = model = FullyShardedDDP(
952
953
954
955
                    model,
                    mixed_precision=mixed_precision,
                    reshard_after_forward=zero_3,
                    cpu_offload=cpu_offload,
956
957
                ).to(self.args.device)

Sylvain Gugger's avatar
Sylvain Gugger committed
958
        elif is_sagemaker_dp_enabled():
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
            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

978
979
    def train(
        self,
980
        resume_from_checkpoint: Optional[Union[str, bool]] = None,
981
        trial: Union["optuna.Trial", Dict[str, Any]] = None,
982
        **kwargs,
983
    ):
Julien Chaumond's avatar
Julien Chaumond committed
984
985
986
987
        """
        Main training entry point.

        Args:
988
989
990
991
992
            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.
993
994
            trial (:obj:`optuna.Trial` or :obj:`Dict[str, Any]`, `optional`):
                The trial run or the hyperparameter dictionary for hyperparameter search.
995
996
            kwargs:
                Additional keyword arguments used to hide deprecated arguments
Julien Chaumond's avatar
Julien Chaumond committed
997
        """
998
999
1000
1001

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

1002
1003
        args = self.args

1004
1005
        self.is_in_train = True

1006
1007
1008
1009
1010
        # 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)

1011
1012
1013
1014
1015
1016
1017
1018
1019
        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
1020
1021
1022
        # This might change the seed so needs to run first.
        self._hp_search_setup(trial)

1023
        # Model re-init
1024
        model_reloaded = False
1025
        if self.model_init is not None:
Sylvain Gugger's avatar
Sylvain Gugger committed
1026
            # Seed must be set before instantiating the model when using model_init.
1027
            set_seed(args.seed)
1028
1029
            self.model = self.call_model_init(trial)
            model_reloaded = True
Sylvain Gugger's avatar
Sylvain Gugger committed
1030
1031
            # Reinitializes optimizer and scheduler
            self.optimizer, self.lr_scheduler = None, None
1032

1033
        # Load potential model checkpoint
1034
        if isinstance(resume_from_checkpoint, bool) and resume_from_checkpoint:
1035
            resume_from_checkpoint = get_last_checkpoint(args.output_dir)
1036
            if resume_from_checkpoint is None:
1037
                raise ValueError(f"No valid checkpoint found in output directory ({args.output_dir})")
1038

1039
1040
1041
1042
        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}")

1043
            logger.info(f"Loading model from {resume_from_checkpoint}).")
1044

1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
            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."
                    )

1055
            if args.deepspeed:
1056
                # will be resumed in deepspeed_init
1057
                pass
1058
            else:
1059
1060
1061
                # 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!
1062
                self._load_state_dict_in_model(state_dict)
1063
1064
1065

        # If model was re-initialized, put it on the right device and update self.model_wrapped
        if model_reloaded:
1066
            if self.place_model_on_device:
1067
                self.model = self.model.to(args.device)
1068
1069
            self.model_wrapped = self.model

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

1073
        # Data loader and number of training steps
Julien Chaumond's avatar
Julien Chaumond committed
1074
        train_dataloader = self.get_train_dataloader()
1075
1076
1077
1078
1079

        # 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
1080
        total_train_batch_size = args.train_batch_size * args.gradient_accumulation_steps * args.world_size
1081
        if train_dataset_is_sized:
1082
            num_update_steps_per_epoch = len(train_dataloader) // args.gradient_accumulation_steps
1083
            num_update_steps_per_epoch = max(num_update_steps_per_epoch, 1)
1084
1085
1086
1087
            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
1088
                )
1089
1090
1091
                # May be slightly incorrect if the last batch in the training datalaoder has a smaller size but it's
                # the best we can do.
                num_train_samples = args.max_steps * total_train_batch_size
1092
            else:
1093
1094
                max_steps = math.ceil(args.num_train_epochs * num_update_steps_per_epoch)
                num_train_epochs = math.ceil(args.num_train_epochs)
1095
                num_train_samples = len(self.train_dataset) * args.num_train_epochs
Julien Chaumond's avatar
Julien Chaumond committed
1096
        else:
1097
            # see __init__. max_steps is set when the dataset has no __len__
1098
1099
            max_steps = args.max_steps
            num_train_epochs = int(args.num_train_epochs)
1100
            num_update_steps_per_epoch = max_steps
1101
            num_train_samples = args.max_steps * total_train_batch_size
Julien Chaumond's avatar
Julien Chaumond committed
1102

1103
1104
1105
        if DebugOption.UNDERFLOW_OVERFLOW in self.args.debug:
            debug_overflow = DebugUnderflowOverflow(self.model)  # noqa

1106
        delay_optimizer_creation = self.sharded_ddp is not None and self.sharded_ddp != ShardedDDPOption.SIMPLE
1107
        if args.deepspeed:
1108
            deepspeed_engine, optimizer, lr_scheduler = deepspeed_init(
1109
1110
                self, num_training_steps=max_steps, resume_from_checkpoint=resume_from_checkpoint
            )
1111
1112
1113
            self.model = deepspeed_engine.module
            self.model_wrapped = deepspeed_engine
            self.deepspeed = deepspeed_engine
1114
1115
            self.optimizer = optimizer
            self.lr_scheduler = lr_scheduler
1116
        elif not delay_optimizer_creation:
1117
1118
            self.create_optimizer_and_scheduler(num_training_steps=max_steps)

1119
        self.state = TrainerState()
1120
        self.state.is_hyper_param_search = trial is not None
Julien Chaumond's avatar
Julien Chaumond committed
1121

1122
        model = self._wrap_model(self.model_wrapped)
Julien Chaumond's avatar
Julien Chaumond committed
1123

1124
1125
1126
1127
        # 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

1128
1129
1130
        if delay_optimizer_creation:
            self.create_optimizer_and_scheduler(num_training_steps=max_steps)

1131
1132
1133
        # Check if saved optimizer or scheduler states exist
        self._load_optimizer_and_scheduler(resume_from_checkpoint)

1134
1135
        # important: at this point:
        # self.model         is the Transformers Model
1136
        # self.model_wrapped is DDP(Transformers Model), Deepspeed(Transformers Model), etc.
1137

Julien Chaumond's avatar
Julien Chaumond committed
1138
        # Train!
1139
        num_examples = (
1140
            self.num_examples(train_dataloader) if train_dataset_is_sized else total_train_batch_size * args.max_steps
1141
1142
        )

Julien Chaumond's avatar
Julien Chaumond committed
1143
        logger.info("***** Running training *****")
1144
1145
        logger.info(f"  Num examples = {num_examples}")
        logger.info(f"  Num Epochs = {num_train_epochs}")
1146
        logger.info(f"  Instantaneous batch size per device = {args.per_device_train_batch_size}")
1147
        logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_train_batch_size}")
1148
        logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
1149
        logger.info(f"  Total optimization steps = {max_steps}")
Julien Chaumond's avatar
Julien Chaumond committed
1150

1151
        self.state.epoch = 0
1152
        start_time = time.time()
Julien Chaumond's avatar
Julien Chaumond committed
1153
1154
        epochs_trained = 0
        steps_trained_in_current_epoch = 0
1155
        steps_trained_progress_bar = None
1156

Julien Chaumond's avatar
Julien Chaumond committed
1157
        # Check if continuing training from a checkpoint
1158
1159
1160
1161
        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"))
1162
            epochs_trained = self.state.global_step // num_update_steps_per_epoch
1163
            if not args.ignore_data_skip:
1164
                steps_trained_in_current_epoch = self.state.global_step % (num_update_steps_per_epoch)
1165
                steps_trained_in_current_epoch *= args.gradient_accumulation_steps
1166
1167
            else:
                steps_trained_in_current_epoch = 0
1168
1169

            logger.info("  Continuing training from checkpoint, will skip to saved global_step")
1170
1171
            logger.info(f"  Continuing training from epoch {epochs_trained}")
            logger.info(f"  Continuing training from global step {self.state.global_step}")
1172
            if not args.ignore_data_skip:
1173
1174
                logger.info(
                    f"  Will skip the first {epochs_trained} epochs then the first {steps_trained_in_current_epoch} "
1175
1176
                    "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."
1177
                )
1178
1179
1180
                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")
1181

Sylvain Gugger's avatar
Sylvain Gugger committed
1182
1183
1184
1185
1186
        # 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
1187
1188
        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
1189
1190
1191
1192
        # 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
1193
1194
        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
1195

1196
        # tr_loss is a tensor to avoid synchronization of TPUs through .item()
1197
        tr_loss = torch.tensor(0.0).to(args.device)
1198
1199
        # _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
1200
        self._globalstep_last_logged = self.state.global_step
Julien Chaumond's avatar
Julien Chaumond committed
1201
        model.zero_grad()
Sylvain Gugger's avatar
Sylvain Gugger committed
1202

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

1205
        # Skip the first epochs_trained epochs to get the random state of the dataloader at the right point.
1206
        if not args.ignore_data_skip:
1207
1208
1209
1210
1211
            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

1212
        for epoch in range(epochs_trained, num_train_epochs):
1213
1214
            if isinstance(train_dataloader, DataLoader) and isinstance(train_dataloader.sampler, DistributedSampler):
                train_dataloader.sampler.set_epoch(epoch)
1215
1216
            elif isinstance(train_dataloader.dataset, IterableDatasetShard):
                train_dataloader.dataset.set_epoch(epoch)
1217

1218
            if is_torch_tpu_available():
1219
                parallel_loader = pl.ParallelLoader(train_dataloader, [args.device]).per_device_loader(args.device)
1220
                epoch_iterator = parallel_loader
1221
            else:
1222
                epoch_iterator = train_dataloader
1223

1224
            # Reset the past mems state at the beginning of each epoch if necessary.
1225
            if args.past_index >= 0:
1226
1227
                self._past = None

1228
            steps_in_epoch = (
1229
                len(epoch_iterator) if train_dataset_is_sized else args.max_steps * args.gradient_accumulation_steps
1230
            )
1231
            self.control = self.callback_handler.on_epoch_begin(args, self.state, self.control)
Sylvain Gugger's avatar
Sylvain Gugger committed
1232

Julien Chaumond's avatar
Julien Chaumond committed
1233
1234
1235
1236
1237
            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
1238
1239
                    if steps_trained_progress_bar is not None:
                        steps_trained_progress_bar.update(1)
1240
1241
                    if steps_trained_in_current_epoch == 0:
                        self._load_rng_state(resume_from_checkpoint)
Julien Chaumond's avatar
Julien Chaumond committed
1242
                    continue
1243
1244
1245
                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
1246

1247
1248
                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
1249

1250
                if (
1251
1252
1253
                    ((step + 1) % args.gradient_accumulation_steps != 0)
                    and args.local_rank != -1
                    and args._no_sync_in_gradient_accumulation
1254
                ):
1255
                    # Avoid unnecessary DDP synchronization since there will be no backward pass on this example.
1256
1257
1258
1259
                    with model.no_sync():
                        tr_loss += self.training_step(model, inputs)
                else:
                    tr_loss += self.training_step(model, inputs)
1260
                self.current_flos += float(self.floating_point_ops(inputs))
Julien Chaumond's avatar
Julien Chaumond committed
1261

1262
1263
1264
1265
                # Optimizer step for deepspeed must be called on every step regardless of the value of gradient_accumulation_steps
                if self.deepspeed:
                    self.deepspeed.step()

1266
                if (step + 1) % args.gradient_accumulation_steps == 0 or (
Julien Chaumond's avatar
Julien Chaumond committed
1267
                    # last step in epoch but step is always smaller than gradient_accumulation_steps
1268
                    steps_in_epoch <= args.gradient_accumulation_steps
1269
                    and (step + 1) == steps_in_epoch
Julien Chaumond's avatar
Julien Chaumond committed
1270
                ):
1271
                    # Gradient clipping
1272
                    if args.max_grad_norm is not None and args.max_grad_norm > 0 and not self.deepspeed:
1273
1274
                        # deepspeed does its own clipping

1275
1276
1277
1278
1279
1280
                        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
1281
                            self.optimizer.clip_grad_norm(args.max_grad_norm)
1282
1283
                        elif hasattr(model, "clip_grad_norm_"):
                            # Some models (like FullyShardedDDP) have a specific way to do gradient clipping
1284
                            model.clip_grad_norm_(args.max_grad_norm)
1285
1286
1287
1288
                        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(),
1289
                                args.max_grad_norm,
1290
1291
1292
                            )

                    # Optimizer step
1293
                    optimizer_was_run = True
Stas Bekman's avatar
Stas Bekman committed
1294
                    if self.deepspeed:
1295
                        pass  # called outside the loop
Stas Bekman's avatar
Stas Bekman committed
1296
                    elif is_torch_tpu_available():
1297
                        xm.optimizer_step(self.optimizer)
1298
                    elif self.use_amp:
1299
                        scale_before = self.scaler.get_scale()
1300
                        self.scaler.step(self.optimizer)
1301
                        self.scaler.update()
1302
1303
                        scale_after = self.scaler.get_scale()
                        optimizer_was_run = scale_before <= scale_after
Lysandre Debut's avatar
Lysandre Debut committed
1304
                    else:
1305
                        self.optimizer.step()
Lysandre Debut's avatar
Lysandre Debut committed
1306

1307
                    if optimizer_was_run and not self.deepspeed:
1308
1309
                        self.lr_scheduler.step()

Julien Chaumond's avatar
Julien Chaumond committed
1310
                    model.zero_grad()
1311
                    self.state.global_step += 1
1312
                    self.state.epoch = epoch + (step + 1) / steps_in_epoch
1313
                    self.control = self.callback_handler.on_step_end(args, self.state, self.control)
Sylvain Gugger's avatar
Sylvain Gugger committed
1314

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

Sylvain Gugger's avatar
Sylvain Gugger committed
1317
                if self.control.should_epoch_stop or self.control.should_training_stop:
Julien Chaumond's avatar
Julien Chaumond committed
1318
                    break
1319

1320
            self.control = self.callback_handler.on_epoch_end(args, self.state, self.control)
1321
            self._maybe_log_save_evaluate(tr_loss, model, trial, epoch)
1322

1323
            if DebugOption.TPU_METRICS_DEBUG in self.args.debug:
1324
1325
1326
1327
1328
1329
1330
1331
                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
1332
            if self.control.should_training_stop:
1333
                break
Julien Chaumond's avatar
Julien Chaumond committed
1334

1335
        if args.past_index and hasattr(self, "_past"):
1336
1337
            # Clean the state at the end of training
            delattr(self, "_past")
Julien Chaumond's avatar
Julien Chaumond committed
1338
1339

        logger.info("\n\nTraining completed. Do not forget to share your model on huggingface.co/models =)\n\n")
1340
        if args.load_best_model_at_end and self.state.best_model_checkpoint is not None:
1341
1342
1343
            # 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")
1344
            elif args.local_rank != -1:
1345
1346
                dist.barrier()

1347
1348
1349
            logger.info(
                f"Loading best model from {self.state.best_model_checkpoint} (score: {self.state.best_metric})."
            )
1350
1351
1352
            # 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!
1353
            self._load_state_dict_in_model(state_dict)
1354

1355
1356
1357
1358
1359
            if self.deepspeed:
                self.deepspeed.load_checkpoint(
                    self.state.best_model_checkpoint, load_optimizer_states=False, load_lr_scheduler_states=False
                )

1360
        metrics = speed_metrics("train", start_time, num_samples=num_train_samples, num_steps=self.state.max_steps)
1361
1362
        self.store_flos()
        metrics["total_flos"] = self.state.total_flos
1363
        self.log(metrics)
1364

1365
        self.control = self.callback_handler.on_train_end(args, self.state, self.control)
1366
1367
        # add remaining tr_loss
        self._total_loss_scalar += tr_loss.item()
Sylvain Gugger's avatar
Sylvain Gugger committed
1368

1369
        self.is_in_train = False
1370

1371
1372
        self._memory_tracker.stop_and_update_metrics(metrics)

1373
        return TrainOutput(self.state.global_step, self._total_loss_scalar / self.state.global_step, metrics)
1374

1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
    def _load_state_dict_in_model(self, state_dict):
        load_result = self.model.load_state_dict(state_dict, strict=False)

        if len(load_result.missing_keys) != 0:
            if set(load_result.missing_keys) == set(self.model._keys_to_ignore_on_save):
                self.model.tie_weights()
            else:
                logger.warn(f"There were missing keys in the checkpoint model loaded: {load_result.missing_keys}.")
        if len(load_result.unexpected_keys) != 0:
            logger.warn(f"There were unexpected keys in the checkpoint model loaded: {load_result.unexpected_keys}.")

1386
    def _maybe_log_save_evaluate(self, tr_loss, model, trial, epoch):
Sylvain Gugger's avatar
Sylvain Gugger committed
1387
1388
1389
        if self.control.should_log:
            logs: Dict[str, float] = {}
            tr_loss_scalar = tr_loss.item()
1390
1391
1392
            # reset tr_loss to zero
            tr_loss -= tr_loss

1393
            logs["loss"] = round(tr_loss_scalar / (self.state.global_step - self._globalstep_last_logged), 4)
1394
            logs["learning_rate"] = self._get_learning_rate()
1395

1396
            self._total_loss_scalar += tr_loss_scalar
1397
            self._globalstep_last_logged = self.state.global_step
Teven's avatar
Teven committed
1398
            self.store_flos()
Sylvain Gugger's avatar
Sylvain Gugger committed
1399
1400
1401
1402
1403
1404
1405

            self.log(logs)

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

Sylvain Gugger's avatar
Sylvain Gugger committed
1407
1408
1409
1410
        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)

1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
    def _load_rng_state(self, checkpoint):
        # Load RNG states from `checkpoint`
        if checkpoint is None:
            return

        local_rank = xm.get_local_ordinal() if is_torch_tpu_available() else self.args.local_rank
        if local_rank != -1:
            rng_file = os.path.join(checkpoint, f"rng_state_{local_rank}.pth")
            if not os.path.isfile(os.path.join(checkpoint, rng_file)):
                logger.info(
                    f"Didn't find an RNG file for process {local_rank}, if you are resuming a training that "
                    "wasn't launched in a distributed fashion, reproducibility is not guaranteed."
                )
                return
        else:
            rng_file = os.path.join(checkpoint, "rng_state.pth")
            if not os.path.isfile(os.path.join(checkpoint, rng_file)):
                logger.info(
                    "Didn't find an RNG file, if you are resuming a training that was launched in a distributed "
                    "fashion, reproducibility is not guaranteed."
                )
                return

        checkpoint_rng_state = torch.load(rng_file)
        random.setstate(checkpoint_rng_state["python"])
        np.random.set_state(checkpoint_rng_state["numpy"])
        torch.random.set_rng_state(checkpoint_rng_state["cpu"])
        if torch.cuda.is_available():
            if self.args.local_rank != -1:
                torch.cuda.random.set_rng_state(checkpoint_rng_state["cuda"])
            else:
                torch.cuda.random.set_rng_state_all(checkpoint_rng_state["cuda"])
        if is_torch_tpu_available():
            xm.set_rng_state(checkpoint_rng_state["xla"])

Sylvain Gugger's avatar
Sylvain Gugger committed
1446
    def _save_checkpoint(self, model, trial, metrics=None):
1447
        # In all cases, including ddp/dp/deepspeed, self.model is always a reference to the model we
1448
        # want to save except FullyShardedDDP.
1449
        # assert unwrap_model(model) is self.model, "internal model should be a reference to self.model"
1450

1451
        # Save model checkpoint
1452
        checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}"
1453

1454
        if self.hp_search_backend is not None and trial is not None:
1455
1456
1457
1458
1459
1460
            if self.hp_search_backend == HPSearchBackend.OPTUNA:
                run_id = trial.number
            else:
                from ray import tune

                run_id = tune.get_trial_id()
1461
            run_name = self.hp_name(trial) if self.hp_name is not None else f"run-{run_id}"
1462
            run_dir = os.path.join(self.args.output_dir, run_name)
1463
        else:
1464
            run_dir = self.args.output_dir
1465
            self.store_flos()
1466

1467
        output_dir = os.path.join(run_dir, checkpoint_folder)
1468
        self.save_model(output_dir)
1469
        if self.deepspeed:
1470
1471
            # 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
1472
            self.deepspeed.save_checkpoint(output_dir)
1473
1474

        # Save optimizer and scheduler
1475
        if self.sharded_ddp == ShardedDDPOption.SIMPLE:
1476
            self.optimizer.consolidate_state_dict()
1477

1478
1479
1480
1481
1482
1483
        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
1484
        elif is_sagemaker_mp_enabled():
1485
1486
1487
1488
1489
1490
1491
1492
1493
            if smp.dp_rank() == 0:
                # 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)
1494
1495
                    if self.use_amp:
                        torch.save(self.scaler.state_dict(), os.path.join(output_dir, "scaler.pt"))
1496
1497
        elif self.is_world_process_zero() and not self.deepspeed:
            # deepspeed.save_checkpoint above saves model/optim/sched
1498
1499
1500
1501
            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)
1502
1503
            if self.use_amp:
                torch.save(self.scaler.state_dict(), os.path.join(output_dir, "scaler.pt"))
1504
1505

        # Determine the new best metric / best model checkpoint
Sylvain Gugger's avatar
Sylvain Gugger committed
1506
        if metrics is not None and self.args.metric_for_best_model is not None:
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
            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"))

1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
        # Save RNG state in non-distributed training
        rng_states = {
            "python": random.getstate(),
            "numpy": np.random.get_state(),
            "cpu": torch.random.get_rng_state(),
        }
        if torch.cuda.is_available():
            if self.args.local_rank == -1:
                # In non distributed, we save the global CUDA RNG state (will take care of DataParallel)
                rng_states["cuda"] = torch.cuda.random.get_rng_state_all()
            else:
                rng_states["cuda"] = torch.cuda.random.get_rng_state()

        if is_torch_tpu_available():
            rng_states["xla"] = xm.get_rng_state()

1541
1542
1543
        # A process can arrive here before the process 0 has a chance to save the model, in which case output_dir may
        # not yet exist.
        os.makedirs(output_dir, exist_ok=True)
1544
1545
1546
1547
1548
1549
        local_rank = xm.get_local_ordinal() if is_torch_tpu_available() else self.args.local_rank
        if local_rank == -1:
            torch.save(rng_states, os.path.join(output_dir, "rng_state.pth"))
        else:
            torch.save(rng_states, os.path.join(output_dir, f"rng_state_{local_rank}.pth"))

1550
1551
1552
1553
        # Maybe delete some older checkpoints.
        if self.is_world_process_zero():
            self._rotate_checkpoints(use_mtime=True, output_dir=run_dir)

1554
    def _load_optimizer_and_scheduler(self, checkpoint):
Sylvain Gugger's avatar
Sylvain Gugger committed
1555
        """If optimizer and scheduler states exist, load them."""
1556
        if checkpoint is None:
1557
1558
            return

1559
        if self.deepspeed:
1560
            # deepspeed loads optimizer/lr_scheduler together with the model in deepspeed_init
1561
1562
            return

1563
1564
        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
1565
1566
1567
1568
        ):
            # 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.
1569
                optimizer_state = torch.load(os.path.join(checkpoint, "optimizer.pt"), map_location="cpu")
Sylvain Gugger's avatar
Sylvain Gugger committed
1570
                with warnings.catch_warnings(record=True) as caught_warnings:
1571
                    lr_scheduler_state = torch.load(os.path.join(checkpoint, "scheduler.pt"), map_location="cpu")
Sylvain Gugger's avatar
Sylvain Gugger committed
1572
1573
1574
1575
1576
1577
1578
1579
                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
1580
                map_location = "cpu" if is_sagemaker_mp_enabled() else self.args.device
Sylvain Gugger's avatar
Sylvain Gugger committed
1581
                self.optimizer.load_state_dict(
Sylvain Gugger's avatar
Sylvain Gugger committed
1582
                    torch.load(os.path.join(checkpoint, "optimizer.pt"), map_location=map_location)
Sylvain Gugger's avatar
Sylvain Gugger committed
1583
1584
                )
                with warnings.catch_warnings(record=True) as caught_warnings:
1585
                    self.lr_scheduler.load_state_dict(torch.load(os.path.join(checkpoint, "scheduler.pt")))
Sylvain Gugger's avatar
Sylvain Gugger committed
1586
                reissue_pt_warnings(caught_warnings)
1587
1588
                if self.use_amp and os.path.isfile(os.path.join(checkpoint, "scaler.pt")):
                    self.scaler.load_state_dict(torch.load(os.path.join(checkpoint, "scaler.pt")))
Sylvain Gugger's avatar
Sylvain Gugger committed
1589

1590
1591
1592
1593
1594
1595
1596
    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,
1597
        hp_name: Optional[Callable[["optuna.Trial"], str]] = None,
1598
        **kwargs,
1599
1600
    ) -> BestRun:
        """
1601
        Launch an hyperparameter search using ``optuna`` or ``Ray Tune``. The optimized quantity is determined by
Sylvain Gugger's avatar
Sylvain Gugger committed
1602
1603
        :obj:`compute_objective`, which defaults to a function returning the evaluation loss when no metric is
        provided, the sum of all metrics otherwise.
1604

Sylvain Gugger's avatar
Sylvain Gugger committed
1605
1606
1607
1608
1609
1610
1611
        .. 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.

1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
        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
1633
                - the documentation of `optuna.create_study
1634
                  <https://optuna.readthedocs.io/en/stable/reference/generated/optuna.study.create_study.html>`__
Sylvain Gugger's avatar
Sylvain Gugger committed
1635
1636
                - the documentation of `tune.run
                  <https://docs.ray.io/en/latest/tune/api_docs/execution.html#tune-run>`__
1637
1638

        Returns:
Tiger's avatar
Tiger committed
1639
            :class:`transformers.trainer_utils.BestRun`: All the information about the best run.
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
        """
        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
1651
            raise RuntimeError("You picked the optuna backend, but it is not installed. Use `pip install optuna`.")
1652
        if backend == HPSearchBackend.RAY and not is_ray_tune_available():
1653
            raise RuntimeError(
Sylvain Gugger's avatar
Sylvain Gugger committed
1654
                "You picked the Ray Tune backend, but it is not installed. Use `pip install 'ray[tune]'`."
1655
1656
            )
        self.hp_search_backend = backend
Sylvain Gugger's avatar
Sylvain Gugger committed
1657
1658
1659
1660
1661
        if self.model_init is None:
            raise RuntimeError(
                "To use hyperparameter search, you need to pass your model through a model_init function."
            )

1662
        self.hp_space = default_hp_space[backend] if hp_space is None else hp_space
1663
        self.hp_name = hp_name
1664
1665
        self.compute_objective = default_compute_objective if compute_objective is None else compute_objective

1666
1667
        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)
1668
1669
1670
1671

        self.hp_search_backend = None
        return best_run

Sylvain Gugger's avatar
Sylvain Gugger committed
1672
    def log(self, logs: Dict[str, float]) -> None:
1673
1674
1675
1676
1677
1678
1679
1680
1681
        """
        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.
        """
1682
        if self.state.epoch is not None:
1683
            logs["epoch"] = round(self.state.epoch, 2)
1684

1685
1686
        output = {**logs, **{"step": self.state.global_step}}
        self.state.log_history.append(output)
1687
        self.control = self.callback_handler.on_log(self.args, self.state, self.control, logs)
Julien Chaumond's avatar
Julien Chaumond committed
1688

sgugger's avatar
Fix CI  
sgugger committed
1689
    def _prepare_inputs(self, inputs: Dict[str, Union[torch.Tensor, Any]]) -> Dict[str, Union[torch.Tensor, Any]]:
1690
1691
1692
1693
        """
        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
1694
        for k, v in inputs.items():
1695
1696
            if isinstance(v, torch.Tensor):
                inputs[k] = v.to(self.args.device)
Julien Chaumond's avatar
Julien Chaumond committed
1697

1698
1699
        if self.args.past_index >= 0 and self._past is not None:
            inputs["mems"] = self._past
1700

1701
1702
        return inputs

1703
    def training_step(self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]]) -> torch.Tensor:
1704
        """
1705
        Perform a training step on a batch of inputs.
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718

        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:
1719
            :obj:`torch.Tensor`: The tensor with training loss on this batch.
1720
1721
        """
        model.train()
1722
        inputs = self._prepare_inputs(inputs)
1723

Sylvain Gugger's avatar
Sylvain Gugger committed
1724
        if is_sagemaker_mp_enabled():
1725
1726
            scaler = self.scaler if self.use_amp else None
            loss_mb = smp_forward_backward(model, inputs, self.args.gradient_accumulation_steps, scaler=scaler)
Sylvain Gugger's avatar
Sylvain Gugger committed
1727
1728
            return loss_mb.reduce_mean().detach().to(self.args.device)

1729
        if self.use_amp:
1730
            with autocast():
Sylvain Gugger's avatar
Sylvain Gugger committed
1731
                loss = self.compute_loss(model, inputs)
1732
        else:
Sylvain Gugger's avatar
Sylvain Gugger committed
1733
            loss = self.compute_loss(model, inputs)
1734

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

1738
1739
        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
1740
1741
            loss = loss / self.args.gradient_accumulation_steps

1742
        if self.use_amp:
1743
            self.scaler.scale(loss).backward()
1744
        elif self.use_apex:
1745
            with amp.scale_loss(loss, self.optimizer) as scaled_loss:
Julien Chaumond's avatar
Julien Chaumond committed
1746
                scaled_loss.backward()
1747
        elif self.deepspeed:
1748
1749
            # loss gets scaled under gradient_accumulation_steps in deepspeed
            loss = self.deepspeed.backward(loss)
Julien Chaumond's avatar
Julien Chaumond committed
1750
1751
1752
        else:
            loss.backward()

1753
        return loss.detach()
Julien Chaumond's avatar
Julien Chaumond committed
1754

1755
    def compute_loss(self, model, inputs, return_outputs=False):
Sylvain Gugger's avatar
Sylvain Gugger committed
1756
1757
1758
1759
1760
        """
        How the loss is computed by Trainer. By default, all models return the loss in the first element.

        Subclass and override for custom behavior.
        """
1761
1762
1763
1764
        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
1765
1766
        outputs = model(**inputs)
        # Save past state if it exists
1767
        # TODO: this needs to be fixed and made cleaner later.
Sylvain Gugger's avatar
Sylvain Gugger committed
1768
1769
        if self.args.past_index >= 0:
            self._past = outputs[self.args.past_index]
Sylvain Gugger's avatar
Sylvain Gugger committed
1770

1771
        if labels is not None:
1772
            loss = self.label_smoother(outputs, labels)
Sylvain Gugger's avatar
Sylvain Gugger committed
1773
1774
        else:
            # We don't use .loss here since the model may return tuples instead of ModelOutput.
1775
1776
1777
            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
1778

1779
1780
    def is_local_process_zero(self) -> bool:
        """
Sylvain Gugger's avatar
Sylvain Gugger committed
1781
1782
        Whether or not this process is the local (e.g., on one machine if training in a distributed fashion on several
        machines) main process.
1783
        """
1784
        if is_torch_tpu_available():
Lysandre Debut's avatar
Lysandre Debut committed
1785
            return xm.is_master_ordinal(local=True)
Sylvain Gugger's avatar
Sylvain Gugger committed
1786
1787
        elif is_sagemaker_mp_enabled():
            return smp.local_rank() == 0
Lysandre Debut's avatar
Lysandre Debut committed
1788
1789
1790
        else:
            return self.args.local_rank in [-1, 0]

1791
1792
    def is_world_process_zero(self) -> bool:
        """
Sylvain Gugger's avatar
Sylvain Gugger committed
1793
1794
        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
1795
        """
1796
        if is_torch_tpu_available():
Lysandre Debut's avatar
Lysandre Debut committed
1797
            return xm.is_master_ordinal(local=False)
Sylvain Gugger's avatar
Sylvain Gugger committed
1798
1799
        elif is_sagemaker_mp_enabled():
            return smp.rank() == 0
Lysandre Debut's avatar
Lysandre Debut committed
1800
        else:
Sylvain Gugger's avatar
Sylvain Gugger committed
1801
            return self.args.process_index == 0
Julien Chaumond's avatar
Julien Chaumond committed
1802
1803
1804

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

1807
        Will only save from the main process.
Julien Chaumond's avatar
Julien Chaumond committed
1808
        """
1809
1810
1811
1812

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

1813
        if is_torch_tpu_available():
1814
            self._save_tpu(output_dir)
Sylvain Gugger's avatar
Sylvain Gugger committed
1815
1816
1817
1818
1819
        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)
1820
1821
1822
1823
        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()
1824

1825
            if self.is_world_process_zero():
1826
                self._save(output_dir, state_dict=state_dict)
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
        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)

1849
1850
        elif self.is_world_process_zero():
            self._save(output_dir)
Julien Chaumond's avatar
Julien Chaumond committed
1851

1852
1853
    def _save_tpu(self, output_dir: Optional[str] = None):
        output_dir = output_dir if output_dir is not None else self.args.output_dir
1854
        logger.info(f"Saving model checkpoint to {output_dir}")
1855
1856
1857
1858
1859
1860
1861
1862

        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")
1863
        if not isinstance(self.model, PreTrainedModel):
1864
1865
1866
1867
1868
1869
1870
            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,
                )
1871
1872
            else:
                logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.")
1873
1874
                state_dict = self.model.state_dict()
                xm.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME))
1875
        else:
1876
            self.model.save_pretrained(output_dir, save_config=self.is_world_process_zero(), save_function=xm.save)
Sylvain Gugger's avatar
Sylvain Gugger committed
1877
        if self.tokenizer is not None and self.is_world_process_zero():
1878
            self.tokenizer.save_pretrained(output_dir)
1879

1880
    def _save(self, output_dir: Optional[str] = None, state_dict=None):
1881
        # 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
1882
1883
        output_dir = output_dir if output_dir is not None else self.args.output_dir
        os.makedirs(output_dir, exist_ok=True)
1884
        logger.info(f"Saving model checkpoint to {output_dir}")
Julien Chaumond's avatar
Julien Chaumond committed
1885
1886
1887
        # Save a trained model and configuration using `save_pretrained()`.
        # They can then be reloaded using `from_pretrained()`
        if not isinstance(self.model, PreTrainedModel):
1888
            if isinstance(unwrap_model(self.model), PreTrainedModel):
1889
1890
1891
                if state_dict is None:
                    state_dict = self.model.state_dict()
                unwrap_model(self.model).save_pretrained(output_dir, state_dict=state_dict)
1892
1893
            else:
                logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.")
1894
1895
                if state_dict is None:
                    state_dict = self.model.state_dict()
1896
                torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME))
1897
        else:
1898
            self.model.save_pretrained(output_dir, state_dict=state_dict)
1899
        if self.tokenizer is not None:
1900
            self.tokenizer.save_pretrained(output_dir)
Julien Chaumond's avatar
Julien Chaumond committed
1901
1902
1903

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

1905
    def store_flos(self):
1906
        # Storing the number of floating-point operations that went into the model
1907
1908
1909
1910
        if self.args.local_rank != -1:
            self.state.total_flos += distributed_broadcast_scalars([self.current_flos]).sum().item()
            self.current_flos = 0
        else:
Teven's avatar
Teven committed
1911
            self.state.total_flos += self.current_flos
1912
            self.current_flos = 0
Julien Chaumond's avatar
Julien Chaumond committed
1913

1914
1915
1916
    def _sorted_checkpoints(
        self, output_dir=None, checkpoint_prefix=PREFIX_CHECKPOINT_DIR, use_mtime=False
    ) -> List[str]:
Julien Chaumond's avatar
Julien Chaumond committed
1917
1918
        ordering_and_checkpoint_path = []

1919
        glob_checkpoints = [str(x) for x in Path(output_dir).glob(f"{checkpoint_prefix}-*")]
Julien Chaumond's avatar
Julien Chaumond committed
1920
1921
1922
1923
1924

        for path in glob_checkpoints:
            if use_mtime:
                ordering_and_checkpoint_path.append((os.path.getmtime(path), path))
            else:
1925
                regex_match = re.match(f".*{checkpoint_prefix}-([0-9]+)", path)
1926
                if regex_match is not None and regex_match.groups() is not None:
Julien Chaumond's avatar
Julien Chaumond committed
1927
1928
1929
1930
                    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]
1931
1932
        # Make sure we don't delete the best model.
        if self.state.best_model_checkpoint is not None:
1933
            best_model_index = checkpoints_sorted.index(str(Path(self.state.best_model_checkpoint)))
1934
1935
            for i in range(best_model_index, len(checkpoints_sorted) - 2):
                checkpoints_sorted[i], checkpoints_sorted[i + 1] = checkpoints_sorted[i + 1], checkpoints_sorted[i]
Julien Chaumond's avatar
Julien Chaumond committed
1936
1937
        return checkpoints_sorted

1938
    def _rotate_checkpoints(self, use_mtime=False, output_dir=None) -> None:
1939
        if self.args.save_total_limit is None or self.args.save_total_limit <= 0:
Julien Chaumond's avatar
Julien Chaumond committed
1940
1941
1942
            return

        # Check if we should delete older checkpoint(s)
1943
        checkpoints_sorted = self._sorted_checkpoints(use_mtime=use_mtime, output_dir=output_dir)
Julien Chaumond's avatar
Julien Chaumond committed
1944
1945
1946
        if len(checkpoints_sorted) <= self.args.save_total_limit:
            return

1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
        # If save_total_limit=1 with load_best_mode_at_end=True, we could end up deleting the last checkpoint, which
        # we don't do to allow resuming.
        save_total_limit = self.args.save_total_limit
        if (
            self.state.best_model_checkpoint is not None
            and self.args.save_total_limit == 1
            and checkpoints_sorted[-1] != self.state.best_model_checkpoint
        ):
            save_total_limit = 2

        number_of_checkpoints_to_delete = max(0, len(checkpoints_sorted) - save_total_limit)
Julien Chaumond's avatar
Julien Chaumond committed
1958
1959
        checkpoints_to_be_deleted = checkpoints_sorted[:number_of_checkpoints_to_delete]
        for checkpoint in checkpoints_to_be_deleted:
1960
            logger.info(f"Deleting older checkpoint [{checkpoint}] due to args.save_total_limit")
Julien Chaumond's avatar
Julien Chaumond committed
1961
1962
            shutil.rmtree(checkpoint)

1963
    def evaluate(
1964
1965
1966
1967
        self,
        eval_dataset: Optional[Dataset] = None,
        ignore_keys: Optional[List[str]] = None,
        metric_key_prefix: str = "eval",
1968
    ) -> Dict[str, float]:
Julien Chaumond's avatar
Julien Chaumond committed
1969
        """
1970
        Run evaluation and returns metrics.
Julien Chaumond's avatar
Julien Chaumond committed
1971

Sylvain Gugger's avatar
Sylvain Gugger committed
1972
1973
        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
1974

1975
1976
        You can also subclass and override this method to inject custom behavior.

Julien Chaumond's avatar
Julien Chaumond committed
1977
        Args:
1978
            eval_dataset (:obj:`Dataset`, `optional`):
1979
                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
1980
1981
                columns not accepted by the ``model.forward()`` method are automatically removed. It must implement the
                :obj:`__len__` method.
1982
1983
1984
            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.
1985
1986
1987
            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)
1988

Julien Chaumond's avatar
Julien Chaumond committed
1989
        Returns:
1990
1991
            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
1992
        """
1993
1994
1995
        # memory metrics - must set up as early as possible
        self._memory_tracker.start()

Julien Chaumond's avatar
Julien Chaumond committed
1996
        eval_dataloader = self.get_eval_dataloader(eval_dataset)
1997
        start_time = time.time()
Julien Chaumond's avatar
Julien Chaumond committed
1998

1999
2000
        eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
        output = eval_loop(
2001
2002
2003
2004
2005
            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,
2006
            ignore_keys=ignore_keys,
2007
            metric_key_prefix=metric_key_prefix,
2008
        )
Lysandre Debut's avatar
Lysandre Debut committed
2009

2010
2011
2012
2013
2014
2015
2016
2017
2018
        total_batch_size = self.args.eval_batch_size * self.args.world_size
        output.metrics.update(
            speed_metrics(
                metric_key_prefix,
                start_time,
                num_samples=output.num_samples,
                num_steps=math.ceil(output.num_samples / total_batch_size),
            )
        )
2019

2020
        self.log(output.metrics)
2021

2022
        if DebugOption.TPU_METRICS_DEBUG in self.args.debug:
Lysandre Debut's avatar
Lysandre Debut committed
2023
2024
2025
            # 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
2026
        self.control = self.callback_handler.on_evaluate(self.args, self.state, self.control, output.metrics)
2027
2028
2029

        self._memory_tracker.stop_and_update_metrics(output.metrics)

Julien Chaumond's avatar
Julien Chaumond committed
2030
2031
        return output.metrics

2032
    def predict(
Bhadresh Savani's avatar
Bhadresh Savani committed
2033
        self, test_dataset: Dataset, ignore_keys: Optional[List[str]] = None, metric_key_prefix: str = "test"
2034
    ) -> PredictionOutput:
Julien Chaumond's avatar
Julien Chaumond committed
2035
        """
2036
        Run prediction and returns predictions and potential metrics.
Julien Chaumond's avatar
Julien Chaumond committed
2037

Sylvain Gugger's avatar
Sylvain Gugger committed
2038
2039
        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()`.
2040
2041
2042

        Args:
            test_dataset (:obj:`Dataset`):
2043
                Dataset to run the predictions on. If it is an :obj:`datasets.Dataset`, columns not accepted by the
2044
                ``model.forward()`` method are automatically removed. Has to implement the method :obj:`__len__`
2045
2046
2047
            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
2048
            metric_key_prefix (:obj:`str`, `optional`, defaults to :obj:`"test"`):
2049
                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
2050
                "test_bleu" if the prefix is "test" (default)
2051

2052
2053
2054
2055
2056
2057
        .. 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
2058
2059
2060
2061
2062
2063
        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
2064
        """
2065
2066
2067
        # memory metrics - must set up as early as possible
        self._memory_tracker.start()

Julien Chaumond's avatar
Julien Chaumond committed
2068
        test_dataloader = self.get_test_dataloader(test_dataset)
2069
        start_time = time.time()
2070

2071
2072
        eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
        output = eval_loop(
2073
2074
            test_dataloader, description="Prediction", ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix
        )
2075
2076
2077
2078
2079
2080
2081
2082
2083
        total_batch_size = self.args.eval_batch_size * self.args.world_size
        output.metrics.update(
            speed_metrics(
                metric_key_prefix,
                start_time,
                num_samples=output.num_samples,
                num_steps=math.ceil(output.num_samples / total_batch_size),
            )
        )
2084
2085
2086

        self._memory_tracker.stop_and_update_metrics(output.metrics)

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

2089
    def evaluation_loop(
2090
2091
2092
2093
2094
        self,
        dataloader: DataLoader,
        description: str,
        prediction_loss_only: Optional[bool] = None,
        ignore_keys: Optional[List[str]] = None,
2095
        metric_key_prefix: str = "eval",
2096
    ) -> EvalLoopOutput:
Julien Chaumond's avatar
Julien Chaumond committed
2097
        """
2098
        Prediction/evaluation loop, shared by :obj:`Trainer.evaluate()` and :obj:`Trainer.predict()`.
Julien Chaumond's avatar
Julien Chaumond committed
2099
2100
2101

        Works both with or without labels.
        """
2102
2103
2104
        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
2105

2106
2107
2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
        # 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
2120

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

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

2128
        batch_size = dataloader.batch_size
2129

2130
        logger.info(f"***** Running {description} *****")
2131
2132
2133
2134
        if isinstance(dataloader.dataset, collections.abc.Sized):
            logger.info(f"  Num examples = {self.num_examples(dataloader)}")
        else:
            logger.info("  Num examples: Unknown")
2135
        logger.info(f"  Batch size = {batch_size}")
2136

Julien Chaumond's avatar
Julien Chaumond committed
2137
2138
        model.eval()

2139
2140
2141
2142
        self.callback_handler.eval_dataloader = dataloader
        # Do this before wrapping.
        eval_dataset = dataloader.dataset

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

2146
        if self.args.past_index >= 0:
2147
            self._past = None
2148

2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
        # 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
2162
        for step, inputs in enumerate(dataloader):
2163
2164
2165
2166
2167
2168
            # 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
2169
            loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)
2170
2171

            # Update containers on host
2172
            if loss is not None:
2173
                losses = self._nested_gather(loss.repeat(batch_size))
2174
                losses_host = losses if losses_host is None else torch.cat((losses_host, losses), dim=0)
2175
            if logits is not None:
2176
2177
                logits = self._pad_across_processes(logits)
                logits = self._nested_gather(logits)
2178
                preds_host = logits if preds_host is None else nested_concat(preds_host, logits, padding_index=-100)
2179
            if labels is not None:
2180
2181
                labels = self._pad_across_processes(labels)
                labels = self._nested_gather(labels)
2182
                labels_host = labels if labels_host is None else nested_concat(labels_host, labels, padding_index=-100)
Sylvain Gugger's avatar
Sylvain Gugger committed
2183
            self.control = self.callback_handler.on_prediction_step(self.args, self.state, self.control)
Julien Chaumond's avatar
Julien Chaumond committed
2184

2185
2186
            # 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:
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
                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)
                    )
2198
2199
2200
2201

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

2202
2203
2204
        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
2205

2206
        # Gather all remaining tensors and put them back on the CPU
2207
2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
2222
2223
2224
2225
2226
2227
2228
2229
2230
2231
2232
2233
2234
2235
2236
        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
2237
2238
        else:
            metrics = {}
2239

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

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

2246
        # Prefix all keys with metric_key_prefix + '_'
2247
        for key in list(metrics.keys()):
2248
2249
            if not key.startswith(f"{metric_key_prefix}_"):
                metrics[f"{metric_key_prefix}_{key}"] = metrics.pop(key)
Julien Chaumond's avatar
Julien Chaumond committed
2250

2251
        return EvalLoopOutput(predictions=all_preds, label_ids=all_labels, metrics=metrics, num_samples=num_samples)
2252

2253
    def _nested_gather(self, tensors, name=None):
2254
2255
2256
2257
2258
2259
2260
        """
        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():
2261
2262
            if name is None:
                name = "nested_gather"
2263
            tensors = nested_xla_mesh_reduce(tensors, name)
Sylvain Gugger's avatar
Sylvain Gugger committed
2264
2265
        elif is_sagemaker_mp_enabled():
            tensors = smp_gather(tensors)
2266
2267
        elif self.args.local_rank != -1:
            tensors = distributed_concat(tensors)
2268
        return tensors
2269

2270
2271
2272
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
    # 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
2302

2303
    def prediction_step(
2304
2305
2306
2307
2308
        self,
        model: nn.Module,
        inputs: Dict[str, Union[torch.Tensor, Any]],
        prediction_loss_only: bool,
        ignore_keys: Optional[List[str]] = None,
2309
    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
        """
        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.
2325
2326
2327
            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.
2328
2329

        Return:
2330
2331
            Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: A tuple with the loss,
            logits and labels (each being optional).
2332
        """
2333
        has_labels = all(inputs.get(k) is not None for k in self.label_names)
2334
        inputs = self._prepare_inputs(inputs)
2335
2336
2337
2338
2339
        if ignore_keys is None:
            if hasattr(self.model, "config"):
                ignore_keys = getattr(self.model.config, "keys_to_ignore_at_inference", [])
            else:
                ignore_keys = []
2340

2341
2342
2343
2344
2345
2346
2347
2348
        # 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

2349
        with torch.no_grad():
Sylvain Gugger's avatar
Sylvain Gugger committed
2350
2351
2352
2353
2354
2355
2356
2357
2358
2359
2360
2361
            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)
2362
                else:
Sylvain Gugger's avatar
Sylvain Gugger committed
2363
2364
2365
2366
2367
2368
                    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)
2369
            else:
Sylvain Gugger's avatar
Sylvain Gugger committed
2370
2371
2372
2373
2374
2375
2376
                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:]
2377
                else:
Sylvain Gugger's avatar
Sylvain Gugger committed
2378
2379
2380
2381
2382
2383
2384
2385
2386
2387
2388
2389
2390
                    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]
2391
2392
2393
2394

        if prediction_loss_only:
            return (loss, None, None)

2395
        logits = nested_detach(logits)
Sylvain Gugger's avatar
Sylvain Gugger committed
2396
2397
2398
2399
        if len(logits) == 1:
            logits = logits[0]

        return (loss, logits, labels)
2400
2401
2402

    def floating_point_ops(self, inputs: Dict[str, Union[torch.Tensor, Any]]):
        """
Sylvain Gugger's avatar
Sylvain Gugger committed
2403
2404
2405
        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.
2406
2407
2408
2409
2410
2411
2412
2413

        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.
        """
2414
2415
        if hasattr(self.model, "floating_point_ops"):
            return self.model.floating_point_ops(inputs)
2416
2417
        else:
            return 0
2418

Sylvain Gugger's avatar
Sylvain Gugger committed
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
    def create_model_card(
        self,
        language: Optional[str] = None,
        license: Optional[str] = None,
        tags: Optional[str] = None,
        model_name: Optional[str] = None,
        finetuned_from: Optional[str] = None,
        dataset_tags: Optional[Union[str, List[str]]] = None,
        dataset: Optional[Union[str, List[str]]] = None,
        dataset_args: Optional[Union[str, List[str]]] = None,
    ):
        training_summary = TrainingSummary.from_trainer(
            self,
            language=language,
            license=license,
            tags=tags,
            model_name=model_name,
            finetuned_from=finetuned_from,
            dataset_tags=dataset_tags,
            dataset=dataset,
            dataset_args=dataset_args,
        )
        model_card = training_summary.to_model_card()
        with open(os.path.join(self.args.output_dir, "README.md"), "w") as f:
            f.write(model_card)

Sylvain Gugger's avatar
Sylvain Gugger committed
2445
2446
2447
2448
2449
2450
2451
2452
    def push_to_hub(
        self,
        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,
Sylvain Gugger's avatar
Sylvain Gugger committed
2453
        **kwargs,
Sylvain Gugger's avatar
Sylvain Gugger committed
2454
2455
2456
2457
2458
2459
    ):
        """
        Upload `self.model` to the 馃 model hub.

        Parameters:
            repo_name (:obj:`str`, `optional`):
Sylvain Gugger's avatar
Sylvain Gugger committed
2460
2461
                Repository name for your model or tokenizer in the hub. If not specified and :obj:`repo_url` is not
                specified either, will default to the stem of :obj:`self.args.output_dir`.
Sylvain Gugger's avatar
Sylvain Gugger committed
2462
2463
2464
2465
2466
2467
2468
2469
2470
2471
2472
2473
2474
2475
2476
            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.
Sylvain Gugger's avatar
Sylvain Gugger committed
2477
2478
            kwargs:
                Additional keyword arguments passed along to :meth:`~transformers.Trainer.create_model_card`.
Sylvain Gugger's avatar
Sylvain Gugger committed
2479
2480
2481
2482
2483
2484
2485
2486
2487
2488
2489
2490

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

Sylvain Gugger's avatar
Sylvain Gugger committed
2491
2492
2493
2494
2495
2496
2497
2498
2499
2500
2501
        if repo_url is None and repo_name is None:
            repo_name = Path(self.args.output_dir).name

        if repo_name is not None:
            model_name = repo_name
        elif repo_url is not None:
            model_name = repo_url.split("/")[-1]
        else:
            model_name = None
        self.create_model_card(model_name=model_name, **kwargs)

Sylvain Gugger's avatar
Sylvain Gugger committed
2502
        with tempfile.TemporaryDirectory() as tmp_dir:
Sylvain Gugger's avatar
Sylvain Gugger committed
2503
2504
2505
2506
            shutil.copy(os.path.join(self.args.output_dir, "README.md"), os.path.join(tmp_dir, "README.md"))
            unwrap_model(self.model).save_pretrained(tmp_dir)
            if self.tokenizer is not None:
                self.tokenizer.save_pretrained(tmp_dir)
Sylvain Gugger's avatar
Sylvain Gugger committed
2507
2508
2509
2510
2511
2512
2513
2514
2515
2516
2517

            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,
            )

2518
2519
2520
2521
2522
2523
2524
2525
2526
2527
2528
2529
2530
2531
2532
2533
2534
2535
2536
2537
2538
2539
2540
2541
2542
2543
2544
2545
2546
2547
2548
2549
2550
2551
2552
2553
2554
2555
2556
2557
2558
    #
    # 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
2559
        # ``train`` is running, halve it first and then put on device
2560
2561
2562
2563
2564
2565
2566
2567
2568
2569
2570
2571
2572
2573
2574
2575
2576
2577
2578
2579
2580
2581
2582
2583
2584
2585
2586
2587
2588
2589
2590
2591
2592
2593
2594
2595
2596
2597
2598
2599
2600
2601
2602
2603
2604
2605
2606
2607
2608
2609
2610
2611
2612
2613
2614
2615
2616
2617
2618
2619
2620
2621
2622
2623
2624
2625
2626
2627
2628
2629
2630
2631
2632
2633
2634
2635
2636
2637
2638
2639
2640
2641
2642
2643
2644
2645
2646
2647
2648
2649
2650
2651
2652
2653
2654
2655
2656
2657
2658
2659
2660
2661
        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)