modeling_utils.py 139 KB
Newer Older
1
# coding=utf-8
thomwolf's avatar
thomwolf committed
2
# Copyright 2018 The Google AI Language Team Authors, Facebook AI Research authors and The HuggingFace Inc. team.
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
#
# 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.

17
import gc
Sylvain Gugger's avatar
Sylvain Gugger committed
18
import json
19
import os
20
import re
Sylvain Gugger's avatar
Sylvain Gugger committed
21
22
import shutil
import tempfile
23
import warnings
24
from contextlib import contextmanager
25
from dataclasses import dataclass
26
from functools import partial
Sylvain Gugger's avatar
Sylvain Gugger committed
27
from pathlib import Path
28
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
29
30

import torch
31
from torch import Tensor, device, nn
32
from torch.nn import CrossEntropyLoss
33

34
from requests import HTTPError
Arthur's avatar
Arthur committed
35
from transformers.utils.hub import convert_file_size_to_int, get_checkpoint_shard_files
36

37
from .activations import get_activation
38
from .configuration_utils import PretrainedConfig
39
from .deepspeed import deepspeed_config, is_deepspeed_zero3_enabled
40
from .dynamic_module_utils import custom_object_save
41
from .generation_utils import GenerationMixin
42
43
44
45
46
47
48
49
from .pytorch_utils import (  # noqa: F401
    Conv1D,
    apply_chunking_to_forward,
    find_pruneable_heads_and_indices,
    prune_conv1d_layer,
    prune_layer,
    prune_linear_layer,
)
50
from .utils import (
Aymeric Augustin's avatar
Aymeric Augustin committed
51
    DUMMY_INPUTS,
52
    FLAX_WEIGHTS_NAME,
53
    HUGGINGFACE_CO_RESOLVE_ENDPOINT,
54
55
    TF2_WEIGHTS_NAME,
    TF_WEIGHTS_NAME,
Sylvain Gugger's avatar
Sylvain Gugger committed
56
    WEIGHTS_INDEX_NAME,
57
    WEIGHTS_NAME,
58
    ContextManagers,
59
    EntryNotFoundError,
60
    ModelOutput,
Sylvain Gugger's avatar
Sylvain Gugger committed
61
    PushToHubMixin,
62
63
    RepositoryNotFoundError,
    RevisionNotFoundError,
64
    cached_path,
65
    has_file,
66
    hf_bucket_url,
67
    is_accelerate_available,
68
    is_offline_mode,
69
    is_remote_url,
70
    logging,
Sylvain Gugger's avatar
Sylvain Gugger committed
71
    replace_return_docstrings,
72
)
73
from .utils.versions import require_version_core
74

Aymeric Augustin's avatar
Aymeric Augustin committed
75

76
77
78
79
80
81
82
83
84
if is_accelerate_available():
    from accelerate import dispatch_model, infer_auto_device_map, init_empty_weights
    from accelerate.utils import (
        load_offloaded_weights,
        offload_weight,
        save_offload_index,
        set_module_tensor_to_device,
    )

Lysandre Debut's avatar
Lysandre Debut committed
85
logger = logging.get_logger(__name__)
86

87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106

_init_weights = True


@contextmanager
def no_init_weights(_enable=True):
    """
    Context manager to globally disable weight initialization to speed up loading large models.

    TODO(Patrick): Delete safety argument `_enable=True` at next major version. .
    """
    global _init_weights
    if _enable:
        _init_weights = False
    try:
        yield
    finally:
        _init_weights = True


thomwolf's avatar
thomwolf committed
107
108
109
110
111
try:
    from torch.nn import Identity
except ImportError:
    # Older PyTorch compatibility
    class Identity(nn.Module):
Lysandre's avatar
Lysandre committed
112
        r"""A placeholder identity operator that is argument-insensitive."""
113

thomwolf's avatar
thomwolf committed
114
        def __init__(self, *args, **kwargs):
Julien Chaumond's avatar
Julien Chaumond committed
115
            super().__init__()
thomwolf's avatar
thomwolf committed
116
117
118
119

        def forward(self, input):
            return input

120

Lysandre Debut's avatar
Lysandre Debut committed
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
def get_parameter_device(parameter: Union[nn.Module, GenerationMixin, "ModuleUtilsMixin"]):
    try:
        return next(parameter.parameters()).device
    except StopIteration:
        # For nn.DataParallel compatibility in PyTorch 1.5

        def find_tensor_attributes(module: nn.Module) -> List[Tuple[str, Tensor]]:
            tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)]
            return tuples

        gen = parameter._named_members(get_members_fn=find_tensor_attributes)
        first_tuple = next(gen)
        return first_tuple[1].device


136
137
138
139
def get_first_parameter_dtype(parameter: Union[nn.Module, GenerationMixin, "ModuleUtilsMixin"]):
    """
    Returns the first parameter dtype (can be non-floating) or asserts if none were found.
    """
Lysandre Debut's avatar
Lysandre Debut committed
140
141
142
    try:
        return next(parameter.parameters()).dtype
    except StopIteration:
Sylvain Gugger's avatar
Sylvain Gugger committed
143
        # For nn.DataParallel compatibility in PyTorch > 1.5
Lysandre Debut's avatar
Lysandre Debut committed
144
145
146
147
148
149
150
151
152
153

        def find_tensor_attributes(module: nn.Module) -> List[Tuple[str, Tensor]]:
            tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)]
            return tuples

        gen = parameter._named_members(get_members_fn=find_tensor_attributes)
        first_tuple = next(gen)
        return first_tuple[1].dtype


154
155
def get_parameter_dtype(parameter: Union[nn.Module, GenerationMixin, "ModuleUtilsMixin"]):
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
156
    Returns the first found floating dtype in parameters if there is one, otherwise returns the last dtype it found.
157
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
158
159
160
161
162
    last_dtype = None
    for t in parameter.parameters():
        last_dtype = t.dtype
        if t.is_floating_point():
            return t.dtype
163

Sylvain Gugger's avatar
Sylvain Gugger committed
164
165
166
    if last_dtype is not None:
        # if no floating dtype was found return whatever the first dtype is
        return last_dtype
167

Sylvain Gugger's avatar
Sylvain Gugger committed
168
169
    else:
        # For nn.DataParallel compatibility in PyTorch > 1.5
170
171
172
173
174
        def find_tensor_attributes(module: nn.Module) -> List[Tuple[str, Tensor]]:
            tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)]
            return tuples

        gen = parameter._named_members(get_members_fn=find_tensor_attributes)
Sylvain Gugger's avatar
Sylvain Gugger committed
175
        last_tuple = None
176
        for tuple in gen:
Sylvain Gugger's avatar
Sylvain Gugger committed
177
            last_tuple = tuple
178
179
            if tuple[1].is_floating_point():
                return tuple[1].dtype
Sylvain Gugger's avatar
Sylvain Gugger committed
180
181
182

        # fallback to the last dtype
        return last_tuple[1].dtype
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197


def get_state_dict_float_dtype(state_dict):
    """
    Returns the first found floating dtype in `state_dict` or asserts if none were found.
    """
    for t in state_dict.values():
        if t.is_floating_point():
            return t.dtype

    raise ValueError("couldn't find any floating point dtypes in state_dict")


def get_state_dict_dtype(state_dict):
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
198
    Returns the first found floating dtype in `state_dict` if there is one, otherwise returns the first dtype.
199
200
201
202
203
204
205
    """
    for t in state_dict.values():
        if t.is_floating_point():
            return t.dtype

    # if no floating dtype was found return whatever the first dtype is
    else:
Sylvain Gugger's avatar
Sylvain Gugger committed
206
        return next(state_dict.values()).dtype
207
208


Sylvain Gugger's avatar
Sylvain Gugger committed
209
210
211
212
213
214
215
216
217
218
219
220
221
def dtype_byte_size(dtype):
    """
    Returns the size (in bytes) occupied by one parameter of type `dtype`.

    Example:

    ```py
    >>> dtype_byte_size(torch.float32)
    4
    ```
    """
    if dtype == torch.bool:
        return 1 / 8
222
    bit_search = re.search(r"[^\d](\d+)$", str(dtype))
Sylvain Gugger's avatar
Sylvain Gugger committed
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
    if bit_search is None:
        raise ValueError(f"`dtype` is not a valid dtype: {dtype}.")
    bit_size = int(bit_search.groups()[0])
    return bit_size // 8


def shard_checkpoint(state_dict: Dict[str, torch.Tensor], max_shard_size: Union[int, str] = "10GB"):
    """
    Splits a model state dictionary in sub-checkpoints so that the final size of each sub-checkpoint does not exceed a
    given size.

    The sub-checkpoints are determined by iterating through the `state_dict` in the order of its keys, so there is no
    optimization made to make each sub-checkpoint as close as possible to the maximum size passed. For example, if the
    limit is 10GB and we have weights of sizes [6GB, 6GB, 2GB, 6GB, 2GB, 2GB] they will get sharded as [6GB], [6+2GB],
    [6+2+2GB] and not [6+2+2GB], [6+2GB], [6GB].

    <Tip warning={true}>

    If one of the model's weight is bigger that `max_sahrd_size`, it will end up in its own sub-checkpoint which will
    have a size greater than `max_shard_size`.

    </Tip>

    Args:
        state_dict (`Dict[str, torch.Tensor]`): The state dictionary of a model to save.
        max_shard_size (`int` or `str`, *optional*, defaults to `"10GB"`):
            The maximum size of each sub-checkpoint. If expressed as a string, needs to be digits followed by a unit
            (like `"5MB"`).
    """
    max_shard_size = convert_file_size_to_int(max_shard_size)

    sharded_state_dicts = []
    current_block = {}
    current_block_size = 0
    total_size = 0

    for key, weight in state_dict.items():
        weight_size = weight.numel() * dtype_byte_size(weight.dtype)

        # If this weight is going to tip up over the maximal size, we split.
        if current_block_size + weight_size > max_shard_size:
            sharded_state_dicts.append(current_block)
            current_block = {}
            current_block_size = 0

        current_block[key] = weight
        current_block_size += weight_size
        total_size += weight_size

    # Add the last block
    sharded_state_dicts.append(current_block)

    # If we only have one shard, we return it
    if len(sharded_state_dicts) == 1:
        return {WEIGHTS_NAME: sharded_state_dicts[0]}, None

    # Otherwise, let's build the index
    weight_map = {}
    shards = {}
    for idx, shard in enumerate(sharded_state_dicts):
        shard_file = WEIGHTS_NAME.replace(".bin", f"-{idx+1:05d}-of-{len(sharded_state_dicts):05d}.bin")
        shards[shard_file] = shard
        for key in shard.keys():
            weight_map[key] = shard_file

    # Add the metadata
    metadata = {"total_size": total_size}
    index = {"metadata": metadata, "weight_map": weight_map}
    return shards, index


294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
def load_sharded_checkpoint(model, folder, strict=True):
    """
    This is the same as
    [`torch.nn.Module.load_state_dict`](https://pytorch.org/docs/stable/generated/torch.nn.Module.html?highlight=load_state_dict#torch.nn.Module.load_state_dict)
    but for a sharded checkpoint.

    This load is performed efficiently: each checkpoint shard is loaded one by one in RAM and deleted after being
    loaded in the model.

    Args:
        model (`torch.nn.Module`): The model in which to load the checkpoint.
        folder (`str` or `os.PathLike`): A path to a folder containing the sharded checkpoint.
        strict (`bool`, *optional`, defaults to `True`):
            Whether to strictly enforce that the keys in the model state dict match the keys in the sharded checkpoint.

    Returns:
        `NamedTuple`: A named tuple with `missing_keys` and `unexpected_keys` fields
            - `missing_keys` is a list of str containing the missing keys
            - `unexpected_keys` is a list of str containing the unexpected keys
    """
    # Load the index
    index_file = os.path.join(folder, WEIGHTS_INDEX_NAME)
    if not os.path.isfile(index_file):
        raise ValueError(f"Can't find a checkpoint index ({WEIGHTS_INDEX_NAME}) in {folder}.")

    with open(index_file, "r", encoding="utf-8") as f:
        index = json.load(f)

    shard_files = list(set(index["weight_map"].values()))

    # If strict=True, error before loading any of the state dicts.
    loaded_keys = index["weight_map"].keys()
    model_keys = model.state_dict().keys()
    missing_keys = [key for key in model_keys if key not in loaded_keys]
    unexpected_keys = [key for key in loaded_keys if key not in model_keys]
    if strict and (len(missing_keys) > 0 or len(unexpected_keys) > 0):
        error_message = f"Error(s) in loading state_dict for {model.__class__.__name__}"
        if len(missing_keys) > 0:
            str_missing_keys = ",".join([f'"{k}"' for k in missing_keys])
            error_message += f"\nMissing key(s): {str_missing_keys}."
        if len(unexpected_keys) > 0:
            str_unexpected_keys = ",".join([f'"{k}"' for k in unexpected_keys])
            error_message += f"\nMissing key(s): {str_unexpected_keys}."
        raise RuntimeError(error_message)

    for shard_file in shard_files:
        state_dict = torch.load(os.path.join(folder, shard_file))
        model.load_state_dict(state_dict, strict=False)

        # Make sure memory is fred before we load the next state dict.
        del state_dict
        gc.collect()

    # Return the same thing as PyTorch load_state_dict function.
    return torch.nn.modules.module._IncompatibleKeys(missing_keys, unexpected_keys)


Sylvain Gugger's avatar
Sylvain Gugger committed
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
def load_state_dict(checkpoint_file: Union[str, os.PathLike]):
    """
    Reads a PyTorch checkpoint file, returning properly formatted errors if they arise.
    """
    try:
        return torch.load(checkpoint_file, map_location="cpu")
    except Exception as e:
        try:
            with open(checkpoint_file) as f:
                if f.read().startswith("version"):
                    raise OSError(
                        "You seem to have cloned a repository without having git-lfs installed. Please install "
                        "git-lfs and run `git lfs install` followed by `git lfs pull` in the folder "
                        "you cloned."
                    )
                else:
                    raise ValueError(
                        f"Unable to locate the file {checkpoint_file} which is necessary to load this pretrained "
                        "model. Make sure you have saved the model properly."
                    ) from e
        except (UnicodeDecodeError, ValueError):
            raise OSError(
                f"Unable to load weights from pytorch checkpoint file for '{checkpoint_file}' "
                f"at '{checkpoint_file}'. "
                "If you tried to load a PyTorch model from a TF 2.0 checkpoint, please set from_tf=True."
            )


def _load_state_dict_into_model(model_to_load, state_dict, start_prefix):
    # Convert old format to new format if needed from a PyTorch state_dict
    old_keys = []
    new_keys = []
    for key in state_dict.keys():
        new_key = None
        if "gamma" in key:
            new_key = key.replace("gamma", "weight")
        if "beta" in key:
            new_key = key.replace("beta", "bias")
        if new_key:
            old_keys.append(key)
            new_keys.append(new_key)
    for old_key, new_key in zip(old_keys, new_keys):
        state_dict[new_key] = state_dict.pop(old_key)

    # copy state_dict so _load_from_state_dict can modify it
    metadata = getattr(state_dict, "_metadata", None)
    state_dict = state_dict.copy()
    if metadata is not None:
        state_dict._metadata = metadata

    error_msgs = []

    # PyTorch's `_load_from_state_dict` does not copy parameters in a module's descendants
    # so we need to apply the function recursively.
    def load(module: nn.Module, prefix=""):
        local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
        args = (state_dict, prefix, local_metadata, True, [], [], error_msgs)
        if is_deepspeed_zero3_enabled():
            import deepspeed

            # because zero3 puts placeholders in model params, this context
            # manager gathers (unpartitions) the params of the current layer, then loads from
            # the state dict and then re-partitions them again
            with deepspeed.zero.GatheredParameters(list(module.parameters(recurse=False)), modifier_rank=0):
                if torch.distributed.get_rank() == 0:
                    module._load_from_state_dict(*args)
        else:
            module._load_from_state_dict(*args)

        for name, child in module._modules.items():
            if child is not None:
                load(child, prefix + name + ".")

    load(model_to_load, prefix=start_prefix)

    return error_msgs


429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
def find_submodule_and_param_name(model, long_key, start_prefix):
    """
    A helper util to find the last sub-module and the param/buffer name. If `start_prefix` is supplied it'll be removed
    from the start of the key
    """

    if len(start_prefix) > 0 and long_key.startswith(start_prefix):
        long_key = ".".join(long_key.split(".")[1:])

    split_key = long_key.split(".")
    submodule = model
    while len(split_key) > 1:
        if hasattr(submodule, split_key[0]):
            submodule = getattr(submodule, split_key[0])
            del split_key[0]
        else:
            submodule = None
            break
    if submodule == model:
        submodule = None
    return submodule, split_key[0]


def _move_model_to_meta(model, loaded_state_dict_keys, start_prefix):
    """
    Moves `loaded_state_dict_keys` in model to meta device which frees up the memory taken by those params.

    `start_prefix` is used for models which insert their name into model keys, e.g. `bert` in
    `bert.pooler.dense.weight`

    """

    # meta device was added in pt=1.9
    require_version_core("torch>=1.9")

    # dematerialize param storage for keys that are going to be replaced by state_dict, by
    # putting those on the meta device
    for k in loaded_state_dict_keys:
        submodule, param_name = find_submodule_and_param_name(model, k, start_prefix)
        if submodule is not None:
            # selectively switch to the meta device only those params/buffers that will
            # be next replaced from state_dict. This a complex way to do p.to_("meta")
            # since we have no in-place to_ for tensors.
            new_val = getattr(submodule, param_name)
            if isinstance(new_val, torch.nn.Parameter):
                # isinstance returns False for Params on meta device, so switch after the check
                new_val = torch.nn.Parameter(new_val.to("meta"))
            else:
                new_val = new_val.to("meta")
            setattr(submodule, param_name, new_val)


481
482
483
484
485
486
487
488
489
490
491
492
493
def _load_state_dict_into_meta_model(
    model,
    state_dict,
    loaded_state_dict_keys,  # left for now but could be removed, see below
    start_prefix,
    expected_keys,
    device_map=None,
    offload_folder=None,
    offload_index=None,
    state_dict_folder=None,
    state_dict_index=None,
    dtype=None,
):
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
    """
    This is somewhat similar to `_load_state_dict_into_model`, but deals with a model that has some or all of its
    params on a `meta` device. It replaces the model params with the data from the `state_dict`, while moving the
    params back to the normal device, but only for `loaded_state_dict_keys`.

    `start_prefix` is used for models which insert their name into model keys, e.g. `bert` in
    `bert.pooler.dense.weight`

    """

    # XXX: remaining features to implement to be fully compatible with _load_state_dict_into_model
    # - deepspeed zero 3 support
    # - need to copy metadata if any - see _load_state_dict_into_model
    # - handling error_msgs - mimicking the error handling in module._load_from_state_dict()
    # - Is there a situation where some keys aren't in `loaded_state_dict_keys` and in which case
    #   they won't get loaded.

    error_msgs = []

513
514
515
516
517
518
519
520
521
522
523
524
525
    old_keys = []
    new_keys = []
    for key in state_dict.keys():
        new_key = None
        if "gamma" in key:
            new_key = key.replace("gamma", "weight")
        if "beta" in key:
            new_key = key.replace("beta", "bias")
        if new_key:
            old_keys.append(key)
            new_keys.append(new_key)
    for old_key, new_key in zip(old_keys, new_keys):
        state_dict[new_key] = state_dict.pop(old_key)
526

527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
    for param_name, param in state_dict.items():
        # First part of the test is always true as load_state_dict_keys always contains state_dict keys.
        if param_name not in loaded_state_dict_keys or param_name not in expected_keys:
            continue

        if param_name.startswith(start_prefix):
            param_name = param_name[len(start_prefix) :]

        module_name = param_name
        # We convert floating dtypes to the `dtype` passed.
        if dtype is not None and not str(param.dtype).startswith("torch.int"):
            param = param.to(dtype)

        if device_map is None:
            param_device = "cpu"
        else:
            # find next higher level module that is defined in device_map:
            # bert.lm_head.weight -> bert.lm_head -> bert -> ''
            while len(module_name) > 0 and module_name not in device_map:
                module_name = ".".join(module_name.split(".")[:-1])
            if module_name == "" and "" not in device_map:
                # TODO: group all errors and raise at the end.
                raise ValueError(f"{param_name} doesn't have any device set.")
            param_device = device_map[module_name]

        if param_device == "disk":
            offload_index = offload_weight(param, param_name, offload_folder, offload_index)
        elif param_device == "cpu" and state_dict_index is not None:
            state_dict_index = offload_weight(param, param_name, state_dict_folder, state_dict_index)
Sylvain Gugger's avatar
Sylvain Gugger committed
556
557
        else:
            set_module_tensor_to_device(model, param_name, param_device, value=param)
558
559

    return error_msgs, offload_index, state_dict_index
560
561


562
class ModuleUtilsMixin:
Julien Chaumond's avatar
Julien Chaumond committed
563
    """
564
    A few utilities for `torch.nn.Modules`, to be used as a mixin.
Julien Chaumond's avatar
Julien Chaumond committed
565
566
    """

567
568
569
570
    @staticmethod
    def _hook_rss_memory_pre_forward(module, *args, **kwargs):
        try:
            import psutil
Sylvain Gugger's avatar
Sylvain Gugger committed
571
        except ImportError:
572
573
574
575
576
577
578
579
580
581
582
            raise ImportError("You need to install psutil (pip install psutil) to use memory tracing.")

        process = psutil.Process(os.getpid())
        mem = process.memory_info()
        module.mem_rss_pre_forward = mem.rss
        return None

    @staticmethod
    def _hook_rss_memory_post_forward(module, *args, **kwargs):
        try:
            import psutil
Sylvain Gugger's avatar
Sylvain Gugger committed
583
        except ImportError:
584
585
586
587
588
589
590
591
592
593
            raise ImportError("You need to install psutil (pip install psutil) to use memory tracing.")

        process = psutil.Process(os.getpid())
        mem = process.memory_info()
        module.mem_rss_post_forward = mem.rss
        mem_rss_diff = module.mem_rss_post_forward - module.mem_rss_pre_forward
        module.mem_rss_diff = mem_rss_diff + (module.mem_rss_diff if hasattr(module, "mem_rss_diff") else 0)
        return None

    def add_memory_hooks(self):
Sylvain Gugger's avatar
Sylvain Gugger committed
594
595
596
        """
        Add a memory hook before and after each sub-module forward pass to record increase in memory consumption.

Sylvain Gugger's avatar
Sylvain Gugger committed
597
598
        Increase in memory consumption is stored in a `mem_rss_diff` attribute for each module and can be reset to zero
        with `model.reset_memory_hooks_state()`.
599
600
601
602
603
604
605
        """
        for module in self.modules():
            module.register_forward_pre_hook(self._hook_rss_memory_pre_forward)
            module.register_forward_hook(self._hook_rss_memory_post_forward)
        self.reset_memory_hooks_state()

    def reset_memory_hooks_state(self):
Sylvain Gugger's avatar
Sylvain Gugger committed
606
        """
Sylvain Gugger's avatar
Sylvain Gugger committed
607
        Reset the `mem_rss_diff` attribute of each module (see [`~modeling_utils.ModuleUtilsMixin.add_memory_hooks`]).
Sylvain Gugger's avatar
Sylvain Gugger committed
608
        """
609
610
611
612
613
        for module in self.modules():
            module.mem_rss_diff = 0
            module.mem_rss_post_forward = 0
            module.mem_rss_pre_forward = 0

614
    @property
615
    def device(self) -> device:
616
        """
617
        `torch.device`: The device on which the module is (assuming that all the module parameters are on the same
618
        device).
619
        """
Lysandre Debut's avatar
Lysandre Debut committed
620
        return get_parameter_device(self)
621

622
    @property
623
    def dtype(self) -> torch.dtype:
624
        """
625
        `torch.dtype`: The dtype of the module (assuming that all the module parameters have the same dtype).
626
        """
Lysandre Debut's avatar
Lysandre Debut committed
627
        return get_parameter_dtype(self)
628
629

    def invert_attention_mask(self, encoder_attention_mask: Tensor) -> Tensor:
Sylvain Gugger's avatar
Sylvain Gugger committed
630
631
632
633
        """
        Invert an attention mask (e.g., switches 0. and 1.).

        Args:
634
            encoder_attention_mask (`torch.Tensor`): An attention mask.
Sylvain Gugger's avatar
Sylvain Gugger committed
635
636

        Returns:
637
            `torch.Tensor`: The inverted attention mask.
Sylvain Gugger's avatar
Sylvain Gugger committed
638
        """
639
640
641
642
643
644
645
646
647
648
        if encoder_attention_mask.dim() == 3:
            encoder_extended_attention_mask = encoder_attention_mask[:, None, :, :]
        if encoder_attention_mask.dim() == 2:
            encoder_extended_attention_mask = encoder_attention_mask[:, None, None, :]
        # T5 has a mask that can compare sequence ids, we can simulate this here with this transposition
        # Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow
        # /transformer/transformer_layers.py#L270
        # encoder_extended_attention_mask = (encoder_extended_attention_mask ==
        # encoder_extended_attention_mask.transpose(-1, -2))
        encoder_extended_attention_mask = encoder_extended_attention_mask.to(dtype=self.dtype)  # fp16 compatibility
Yih-Dar's avatar
Yih-Dar committed
649
        encoder_extended_attention_mask = (1.0 - encoder_extended_attention_mask) * torch.finfo(self.dtype).min
650

651
652
        return encoder_extended_attention_mask

653
    @staticmethod
654
655
656
657
658
659
660
    def create_extended_attention_mask_for_decoder(input_shape, attention_mask, device=None):
        if device is not None:
            warnings.warn(
                "The `device` argument is deprecated and will be removed in v5 of Transformers.", FutureWarning
            )
        else:
            device = attention_mask.device
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
        batch_size, seq_length = input_shape
        seq_ids = torch.arange(seq_length, device=device)
        causal_mask = seq_ids[None, None, :].repeat(batch_size, seq_length, 1) <= seq_ids[None, :, None]
        # in case past_key_values are used we need to add a prefix ones mask to the causal mask
        # causal and attention masks must have same type with pytorch version < 1.3
        causal_mask = causal_mask.to(attention_mask.dtype)

        if causal_mask.shape[1] < attention_mask.shape[1]:
            prefix_seq_len = attention_mask.shape[1] - causal_mask.shape[1]
            causal_mask = torch.cat(
                [
                    torch.ones((batch_size, seq_length, prefix_seq_len), device=device, dtype=causal_mask.dtype),
                    causal_mask,
                ],
                axis=-1,
            )

        extended_attention_mask = causal_mask[:, None, :, :] * attention_mask[:, None, None, :]
        return extended_attention_mask

681
    def get_extended_attention_mask(
Yih-Dar's avatar
Yih-Dar committed
682
        self, attention_mask: Tensor, input_shape: Tuple[int], device: device = None, dtype: torch.float = None
683
    ) -> Tensor:
Sylvain Gugger's avatar
Sylvain Gugger committed
684
685
        """
        Makes broadcastable attention and causal masks so that future and masked tokens are ignored.
686
687

        Arguments:
688
            attention_mask (`torch.Tensor`):
Sylvain Gugger's avatar
Sylvain Gugger committed
689
                Mask with ones indicating tokens to attend to, zeros for tokens to ignore.
690
            input_shape (`Tuple[int]`):
Sylvain Gugger's avatar
Sylvain Gugger committed
691
                The shape of the input to the model.
692
693

        Returns:
694
            `torch.Tensor` The extended attention mask, with a the same dtype as `attention_mask.dtype`.
695
        """
Yih-Dar's avatar
Yih-Dar committed
696
697
698
        if dtype is None:
            dtype = self.dtype

699
700
701
702
703
704
        if not (attention_mask.dim() == 2 and self.config.is_decoder):
            # show warning only if it won't be shown in `create_extended_attention_mask_for_decoder`
            if device is not None:
                warnings.warn(
                    "The `device` argument is deprecated and will be removed in v5 of Transformers.", FutureWarning
                )
705
706
707
708
709
710
711
712
713
        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
        # ourselves in which case we just need to make it broadcastable to all heads.
        if attention_mask.dim() == 3:
            extended_attention_mask = attention_mask[:, None, :, :]
        elif attention_mask.dim() == 2:
            # Provided a padding mask of dimensions [batch_size, seq_length]
            # - if the model is a decoder, apply a causal mask in addition to the padding mask
            # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
            if self.config.is_decoder:
714
                extended_attention_mask = ModuleUtilsMixin.create_extended_attention_mask_for_decoder(
715
716
                    input_shape, attention_mask, device
                )
717
718
719
720
            else:
                extended_attention_mask = attention_mask[:, None, None, :]
        else:
            raise ValueError(
721
                f"Wrong shape for input_ids (shape {input_shape}) or attention_mask (shape {attention_mask.shape})"
722
723
724
725
726
727
728
            )

        # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
        # masked positions, this operation will create a tensor which is 0.0 for
        # positions we want to attend and -10000.0 for masked positions.
        # Since we are adding it to the raw scores before the softmax, this is
        # effectively the same as removing these entirely.
Yih-Dar's avatar
Yih-Dar committed
729
730
        extended_attention_mask = extended_attention_mask.to(dtype=dtype)  # fp16 compatibility
        extended_attention_mask = (1.0 - extended_attention_mask) * torch.finfo(dtype).min
731
732
        return extended_attention_mask

Sylvain Gugger's avatar
Sylvain Gugger committed
733
734
735
    def get_head_mask(
        self, head_mask: Optional[Tensor], num_hidden_layers: int, is_attention_chunked: bool = False
    ) -> Tensor:
736
        """
Sylvain Gugger's avatar
Sylvain Gugger committed
737
738
739
        Prepare the head mask if needed.

        Args:
740
            head_mask (`torch.Tensor` with shape `[num_heads]` or `[num_hidden_layers x num_heads]`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
741
                The mask indicating if we should keep the heads or not (1.0 for keep, 0.0 for discard).
742
            num_hidden_layers (`int`):
Sylvain Gugger's avatar
Sylvain Gugger committed
743
                The number of hidden layers in the model.
744
            is_attention_chunked: (`bool`, *optional*, defaults to `False`):
Sylvain Gugger's avatar
Sylvain Gugger committed
745
746
                Whether or not the attentions scores are computed by chunks or not.

747
        Returns:
Sylvain Gugger's avatar
Sylvain Gugger committed
748
749
            `torch.Tensor` with shape `[num_hidden_layers x batch x num_heads x seq_length x seq_length]` or list with
            `[None]` for each layer.
750
751
752
        """
        if head_mask is not None:
            head_mask = self._convert_head_mask_to_5d(head_mask, num_hidden_layers)
Patrick von Platen's avatar
Patrick von Platen committed
753
754
            if is_attention_chunked is True:
                head_mask = head_mask.unsqueeze(-1)
755
756
757
758
759
760
761
762
763
764
765
766
767
        else:
            head_mask = [None] * num_hidden_layers

        return head_mask

    def _convert_head_mask_to_5d(self, head_mask, num_hidden_layers):
        """-> [num_hidden_layers x batch x num_heads x seq_length x seq_length]"""
        if head_mask.dim() == 1:
            head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
            head_mask = head_mask.expand(num_hidden_layers, -1, -1, -1, -1)
        elif head_mask.dim() == 2:
            head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1)  # We can specify head_mask for each layer
        assert head_mask.dim() == 5, f"head_mask.dim != 5, instead {head_mask.dim()}"
768
        head_mask = head_mask.to(dtype=self.dtype)  # switch to float if need + fp16 compatibility
769
770
        return head_mask

771
772
773
774
775
    def num_parameters(self, only_trainable: bool = False, exclude_embeddings: bool = False) -> int:
        """
        Get number of (optionally, trainable or non-embeddings) parameters in the module.

        Args:
776
            only_trainable (`bool`, *optional*, defaults to `False`):
777
778
                Whether or not to return only the number of trainable parameters

779
            exclude_embeddings (`bool`, *optional*, defaults to `False`):
780
781
782
                Whether or not to return only the number of non-embeddings parameters

        Returns:
783
            `int`: The number of parameters.
784
785
        """

786
787
788
789
790
791
792
793
794
795
        if exclude_embeddings:
            embedding_param_names = [
                f"{name}.weight" for name, module_type in self.named_modules() if isinstance(module_type, nn.Embedding)
            ]
            non_embedding_parameters = [
                parameter for name, parameter in self.named_parameters() if name not in embedding_param_names
            ]
            return sum(p.numel() for p in non_embedding_parameters if p.requires_grad or not only_trainable)
        else:
            return sum(p.numel() for p in self.parameters() if p.requires_grad or not only_trainable)
796
797
798
799
800
801

    def estimate_tokens(self, input_dict: Dict[str, Union[torch.Tensor, Any]]) -> int:
        """
        Helper function to estimate the total number of tokens from the model inputs.

        Args:
802
            inputs (`dict`): The model inputs.
803
804

        Returns:
805
            `int`: The total number of tokens.
806
        """
807
808
        if not hasattr(self, "warnings_issued"):
            self.warnings_issued = {}
809
810
        if self.main_input_name in input_dict:
            return input_dict[self.main_input_name].numel()
811
        elif "estimate_tokens" not in self.warnings_issued:
812
            logger.warning(
813
814
                "Could not estimate the number of tokens of the input, floating-point operations will not be computed"
            )
815
816
            self.warnings_issued["estimate_tokens"] = True
        return 0
817
818
819
820
821
822
823

    def floating_point_ops(
        self, input_dict: Dict[str, Union[torch.Tensor, Any]], exclude_embeddings: bool = True
    ) -> int:
        """
        Get number of (optionally, non-embeddings) floating-point operations for the forward and backward passes of a
        batch with this transformer model. Default approximation neglects the quadratic dependency on the number of
Sylvain Gugger's avatar
Sylvain Gugger committed
824
825
        tokens (valid if `12 * d_model << sequence_length`) as laid out in [this
        paper](https://arxiv.org/pdf/2001.08361.pdf) section 2.1. Should be overridden for transformers with parameter
Sylvain Gugger's avatar
Sylvain Gugger committed
826
        re-use e.g. Albert or Universal Transformers, or if doing long-range modeling with very high sequence lengths.
827
828

        Args:
829
            batch_size (`int`):
830
831
                The batch size for the forward pass.

832
            sequence_length (`int`):
833
834
                The number of tokens in each line of the batch.

835
            exclude_embeddings (`bool`, *optional*, defaults to `True`):
836
837
838
                Whether or not to count embedding and softmax operations.

        Returns:
839
            `int`: The number of floating-point operations.
840
841
842
843
        """

        return 6 * self.estimate_tokens(input_dict) * self.num_parameters(exclude_embeddings=exclude_embeddings)

Julien Chaumond's avatar
Julien Chaumond committed
844

Sylvain Gugger's avatar
Sylvain Gugger committed
845
class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMixin):
846
847
    r"""
    Base class for all models.
848

Sylvain Gugger's avatar
Sylvain Gugger committed
849
850
    [`PreTrainedModel`] takes care of storing the configuration of the models and handles methods for loading,
    downloading and saving models as well as a few methods common to all models to:
851

852
853
        - resize the input embeddings,
        - prune heads in the self-attention heads.
854

855
    Class attributes (overridden by derived classes):
Sylvain Gugger's avatar
Sylvain Gugger committed
856

Sylvain Gugger's avatar
Sylvain Gugger committed
857
858
859
860
        - **config_class** ([`PretrainedConfig`]) -- A subclass of [`PretrainedConfig`] to use as configuration class
          for this model architecture.
        - **load_tf_weights** (`Callable`) -- A python *method* for loading a TensorFlow checkpoint in a PyTorch model,
          taking as arguments:
861

Sylvain Gugger's avatar
Sylvain Gugger committed
862
863
            - **model** ([`PreTrainedModel`]) -- An instance of the model on which to load the TensorFlow checkpoint.
            - **config** ([`PreTrainedConfig`]) -- An instance of the configuration associated to the model.
864
            - **path** (`str`) -- A path to the TensorFlow checkpoint.
865

Sylvain Gugger's avatar
Sylvain Gugger committed
866
867
        - **base_model_prefix** (`str`) -- A string indicating the attribute associated to the base model in derived
          classes of the same architecture adding modules on top of the base model.
868
        - **is_parallelizable** (`bool`) -- A flag indicating whether this model supports model parallelization.
Sylvain Gugger's avatar
Sylvain Gugger committed
869
870
        - **main_input_name** (`str`) -- The name of the principal input to the model (often `input_ids` for NLP
          models, `pixel_values` for vision models and `input_values` for speech models).
871
    """
872
    config_class = None
873
    base_model_prefix = ""
874
    main_input_name = "input_ids"
875
    _auto_class = None
876
    _no_split_modules = None
877

878
879
    # a list of `re` patterns of `state_dict` keys that should be removed from the list of missing
    # keys we find (keys inside the model but not in the checkpoint) and avoid unnecessary warnings.
880
    _keys_to_ignore_on_load_missing = None
881
882
883
    # a list of `re` patterns of `state_dict` keys that should be removed from the list of
    # unexpected keys we find (keys inside the checkpoint but not the model) and avoid unnecessary
    # warnings.
884
    _keys_to_ignore_on_load_unexpected = None
885
886
    # a list of `state_dict` keys to ignore when saving the model (useful for keys that aren't
    # trained, but which are either deterministic or tied variables)
887
    _keys_to_ignore_on_save = None
888

889
    is_parallelizable = False
890
    supports_gradient_checkpointing = False
891

892
    @property
893
    def dummy_inputs(self) -> Dict[str, torch.Tensor]:
894
        """
895
        `Dict[str, torch.Tensor]`: Dummy inputs to do a forward pass in the network.
896
        """
897
        return {"input_ids": torch.tensor(DUMMY_INPUTS)}
898

899
900
901
902
903
904
905
    @property
    def framework(self) -> str:
        """
        :str: Identifies that this is a PyTorch model.
        """
        return "pt"

906
    def __init__(self, config: PretrainedConfig, *inputs, **kwargs):
Julien Chaumond's avatar
Julien Chaumond committed
907
        super().__init__()
908
909
        if not isinstance(config, PretrainedConfig):
            raise ValueError(
910
911
912
                f"Parameter config in `{self.__class__.__name__}(config)` should be an instance of class "
                "`PretrainedConfig`. To create a model from a pretrained model use "
                f"`model = {self.__class__.__name__}.from_pretrained(PRETRAINED_MODEL_NAME)`"
913
            )
914
        # Save config and origin of the pretrained weights if given in model
915
        self.config = config
916
        self.name_or_path = config.name_or_path
917
        self.warnings_issued = {}
918
919
920
921
922
923
924
925
926
927
928
929
930
931

    def post_init(self):
        """
        A method executed at the end of each Transformer model initialization, to execute code that needs the model's
        modules properly initialized (such as weight initialization).
        """
        self.init_weights()
        self._backward_compatibility_gradient_checkpointing()

    def _backward_compatibility_gradient_checkpointing(self):
        if self.supports_gradient_checkpointing and getattr(self.config, "gradient_checkpointing", False):
            self.gradient_checkpointing_enable()
            # Remove the attribute now that is has been consumed, so it's no saved in the config.
            delattr(self.config, "gradient_checkpointing")
932

933
934
935
936
937
938
    @classmethod
    def _from_config(cls, config, **kwargs):
        """
        All context managers that the model should be initialized under go here.

        Args:
939
940
            torch_dtype (`torch.dtype`, *optional*):
                Override the default `torch.dtype` and load the model under this dtype.
941
942
943
944
945
946
947
948
949
950
951
952
953
954
        """
        torch_dtype = kwargs.pop("torch_dtype", None)

        # override default dtype if needed
        dtype_orig = None
        if torch_dtype is not None:
            dtype_orig = cls._set_default_torch_dtype(torch_dtype)

        if is_deepspeed_zero3_enabled():
            import deepspeed

            logger.info("Detected DeepSpeed ZeRO-3: activating zero.init() for this model")
            # this immediately partitions the model across all gpus, to avoid the overhead in time
            # and memory copying it on CPU or each GPU first
955
            with deepspeed.zero.Init(config_dict_or_path=deepspeed_config()):
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
                model = cls(config, **kwargs)
        else:
            model = cls(config, **kwargs)

        # restore default dtype if it was modified
        if dtype_orig is not None:
            torch.set_default_dtype(dtype_orig)

        return model

    @classmethod
    def _set_default_torch_dtype(cls, dtype: torch.dtype) -> torch.dtype:
        """
        Change the default dtype and return the previous one. This is needed when wanting to instantiate the model
        under specific dtype.

        Args:
973
            dtype (`torch.dtype`):
974
975
976
                a floating dtype to set to.

        Returns:
Sylvain Gugger's avatar
Sylvain Gugger committed
977
978
            `torch.dtype`: the original `dtype` that can be used to restore `torch.set_default_dtype(dtype)` if it was
            modified. If it wasn't, returns `None`.
979

980
981
        Note `set_default_dtype` currently only works with floating-point types and asserts if for example,
        `torch.int64` is passed. So if a non-float `dtype` is passed this functions will throw an exception.
982
983
984
985
986
987
988
989
990
991
992
        """
        if not dtype.is_floating_point:
            raise ValueError(
                f"Can't instantiate {cls.__name__} model under dtype={dtype} since it is not a floating point dtype"
            )

        logger.info(f"Instantiating {cls.__name__} model under default dtype {dtype}.")
        dtype_orig = torch.get_default_dtype()
        torch.set_default_dtype(dtype)
        return dtype_orig

993
    @property
994
995
    def base_model(self) -> nn.Module:
        """
996
        `torch.nn.Module`: The main body of the model.
997
        """
998
        return getattr(self, self.base_model_prefix, self)
thomwolf's avatar
thomwolf committed
999

1000
    def get_input_embeddings(self) -> nn.Module:
1001
1002
1003
1004
        """
        Returns the model's input embeddings.

        Returns:
1005
            `nn.Module`: A torch module mapping vocabulary to hidden states.
thomwolf's avatar
thomwolf committed
1006
        """
1007
        base_model = getattr(self, self.base_model_prefix, self)
thomwolf's avatar
thomwolf committed
1008
1009
1010
1011
        if base_model is not self:
            return base_model.get_input_embeddings()
        else:
            raise NotImplementedError
thomwolf's avatar
thomwolf committed
1012

1013
    def set_input_embeddings(self, value: nn.Module):
1014
        """
Sylvain Gugger's avatar
Sylvain Gugger committed
1015
        Set model's input embeddings.
1016
1017

        Args:
1018
            value (`nn.Module`): A module mapping vocabulary to hidden states.
thomwolf's avatar
thomwolf committed
1019
1020
1021
1022
1023
1024
        """
        base_model = getattr(self, self.base_model_prefix, self)
        if base_model is not self:
            base_model.set_input_embeddings(value)
        else:
            raise NotImplementedError
thomwolf's avatar
thomwolf committed
1025

1026
    def get_output_embeddings(self) -> nn.Module:
1027
1028
1029
1030
        """
        Returns the model's output embeddings.

        Returns:
1031
            `nn.Module`: A torch module mapping hidden states to vocabulary.
thomwolf's avatar
thomwolf committed
1032
        """
1033
        return None  # Overwrite for models with output embeddings
thomwolf's avatar
thomwolf committed
1034

1035
1036
1037
1038
    def _init_weights(self, module):
        """
        Initialize the weights. This method should be overridden by derived class.
        """
1039
        raise NotImplementedError(f"Make sure `_init_weights` is implemented for {self.__class__}")
1040

1041
    def tie_weights(self):
1042
1043
        """
        Tie the weights between the input embeddings and the output embeddings.
1044

Sylvain Gugger's avatar
Sylvain Gugger committed
1045
1046
        If the `torchscript` flag is set in the configuration, can't handle parameter sharing so we are cloning the
        weights instead.
thomwolf's avatar
thomwolf committed
1047
        """
1048
1049
1050
1051
        if getattr(self.config, "tie_word_embeddings", True):
            output_embeddings = self.get_output_embeddings()
            if output_embeddings is not None:
                self._tie_or_clone_weights(output_embeddings, self.get_input_embeddings())
thomwolf's avatar
thomwolf committed
1052

1053
        if getattr(self.config, "is_encoder_decoder", False) and getattr(self.config, "tie_encoder_decoder", False):
Weizhen's avatar
Weizhen committed
1054
1055
            if hasattr(self, self.base_model_prefix):
                self = getattr(self, self.base_model_prefix)
1056
1057
            self._tie_encoder_decoder_weights(self.encoder, self.decoder, self.base_model_prefix)

Sylvain Gugger's avatar
Sylvain Gugger committed
1058
1059
1060
1061
        for module in self.modules():
            if hasattr(module, "_tie_weights"):
                module._tie_weights()

1062
1063
1064
    @staticmethod
    def _tie_encoder_decoder_weights(encoder: nn.Module, decoder: nn.Module, base_model_prefix: str):
        uninitialized_encoder_weights: List[str] = []
Weizhen's avatar
Weizhen committed
1065
1066
        if decoder.__class__ != encoder.__class__:
            logger.info(
Sylvain Gugger's avatar
Sylvain Gugger committed
1067
1068
                f"{decoder.__class__} and {encoder.__class__} are not equal. In this case make sure that all encoder"
                " weights are correctly initialized."
Weizhen's avatar
Weizhen committed
1069
            )
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079

        def tie_encoder_to_decoder_recursively(
            decoder_pointer: nn.Module,
            encoder_pointer: nn.Module,
            module_name: str,
            uninitialized_encoder_weights: List[str],
            depth=0,
        ):
            assert isinstance(decoder_pointer, nn.Module) and isinstance(
                encoder_pointer, nn.Module
1080
            ), f"{decoder_pointer} and {encoder_pointer} have to be of type nn.Module"
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
            if hasattr(decoder_pointer, "weight"):
                assert hasattr(encoder_pointer, "weight")
                encoder_pointer.weight = decoder_pointer.weight
                if hasattr(decoder_pointer, "bias"):
                    assert hasattr(encoder_pointer, "bias")
                    encoder_pointer.bias = decoder_pointer.bias
                return

            encoder_modules = encoder_pointer._modules
            decoder_modules = decoder_pointer._modules
            if len(decoder_modules) > 0:
                assert (
                    len(encoder_modules) > 0
                ), f"Encoder module {encoder_pointer} does not match decoder module {decoder_pointer}"

                all_encoder_weights = set([module_name + "/" + sub_name for sub_name in encoder_modules.keys()])
                encoder_layer_pos = 0
                for name, module in decoder_modules.items():
                    if name.isdigit():
                        encoder_name = str(int(name) + encoder_layer_pos)
                        decoder_name = name
Weizhen's avatar
Weizhen committed
1102
1103
1104
                        if not isinstance(decoder_modules[decoder_name], type(encoder_modules[encoder_name])) and len(
                            encoder_modules
                        ) != len(decoder_modules):
1105
1106
                            # this can happen if the name corresponds to the position in a list module list of layers
                            # in this case the decoder has added a cross-attention that the encoder does not have
1107
                            # thus skip this step and subtract one layer pos from encoder
1108
1109
1110
1111
1112
1113
                            encoder_layer_pos -= 1
                            continue
                    elif name not in encoder_modules:
                        continue
                    elif depth > 500:
                        raise ValueError(
Sylvain Gugger's avatar
Sylvain Gugger committed
1114
1115
                            "Max depth of recursive function `tie_encoder_to_decoder` reached. It seems that there is"
                            " a circular dependency between two or more `nn.Modules` of your model."
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
                        )
                    else:
                        decoder_name = encoder_name = name
                    tie_encoder_to_decoder_recursively(
                        decoder_modules[decoder_name],
                        encoder_modules[encoder_name],
                        module_name + "/" + name,
                        uninitialized_encoder_weights,
                        depth=depth + 1,
                    )
                    all_encoder_weights.remove(module_name + "/" + encoder_name)

                uninitialized_encoder_weights += list(all_encoder_weights)

        # tie weights recursively
        tie_encoder_to_decoder_recursively(decoder, encoder, base_model_prefix, uninitialized_encoder_weights)
        if len(uninitialized_encoder_weights) > 0:
            logger.warning(
                f"The following encoder weights were not tied to the decoder {uninitialized_encoder_weights}"
            )

1137
    def _tie_or_clone_weights(self, output_embeddings, input_embeddings):
Lysandre's avatar
Lysandre committed
1138
        """Tie or clone module weights depending of whether we are using TorchScript or not"""
thomwolf's avatar
thomwolf committed
1139
        if self.config.torchscript:
1140
            output_embeddings.weight = nn.Parameter(input_embeddings.weight.clone())
thomwolf's avatar
thomwolf committed
1141
        else:
1142
            output_embeddings.weight = input_embeddings.weight
thomwolf's avatar
thomwolf committed
1143

Sam Shleifer's avatar
Sam Shleifer committed
1144
        if getattr(output_embeddings, "bias", None) is not None:
1145
            output_embeddings.bias.data = nn.functional.pad(
1146
                output_embeddings.bias.data,
Lysandre's avatar
Lysandre committed
1147
1148
1149
1150
                (
                    0,
                    output_embeddings.weight.shape[0] - output_embeddings.bias.shape[0],
                ),
1151
1152
                "constant",
                0,
1153
            )
1154
        if hasattr(output_embeddings, "out_features") and hasattr(input_embeddings, "num_embeddings"):
1155
            output_embeddings.out_features = input_embeddings.num_embeddings
1156

1157
    def resize_token_embeddings(self, new_num_tokens: Optional[int] = None) -> nn.Embedding:
1158
        """
1159
        Resizes input token embeddings matrix of the model if `new_num_tokens != config.vocab_size`.
1160

1161
        Takes care of tying weights embeddings afterwards if the model class has a `tie_weights()` method.
thomwolf's avatar
thomwolf committed
1162

1163
        Arguments:
1164
            new_num_tokens (`int`, *optional*):
1165
                The number of new tokens in the embedding matrix. Increasing the size will add newly initialized
Sylvain Gugger's avatar
Sylvain Gugger committed
1166
1167
                vectors at the end. Reducing the size will remove vectors from the end. If not provided or `None`, just
                returns a pointer to the input tokens `torch.nn.Embedding` module of the model without doing anything.
1168
1169

        Return:
1170
            `torch.nn.Embedding`: Pointer to the input tokens Embeddings Module of the model.
thomwolf's avatar
thomwolf committed
1171
        """
1172
        model_embeds = self._resize_token_embeddings(new_num_tokens)
thomwolf's avatar
thomwolf committed
1173
1174
        if new_num_tokens is None:
            return model_embeds
thomwolf's avatar
thomwolf committed
1175
1176
1177

        # Update base model and current model config
        self.config.vocab_size = new_num_tokens
1178
        self.vocab_size = new_num_tokens
thomwolf's avatar
thomwolf committed
1179
1180

        # Tie weights again if needed
1181
        self.tie_weights()
thomwolf's avatar
thomwolf committed
1182

thomwolf's avatar
thomwolf committed
1183
1184
        return model_embeds

1185
    def _resize_token_embeddings(self, new_num_tokens):
thomwolf's avatar
thomwolf committed
1186
1187
1188
        old_embeddings = self.get_input_embeddings()
        new_embeddings = self._get_resized_embeddings(old_embeddings, new_num_tokens)
        self.set_input_embeddings(new_embeddings)
1189
1190
1191
1192
1193
1194
1195

        # if word embeddings are not tied, make sure that lm head is resized as well
        if self.get_output_embeddings() is not None and not self.config.tie_word_embeddings:
            old_lm_head = self.get_output_embeddings()
            new_lm_head = self._get_resized_lm_head(old_lm_head, new_num_tokens)
            self.set_output_embeddings(new_lm_head)

thomwolf's avatar
thomwolf committed
1196
        return self.get_input_embeddings()
1197

1198
    def _get_resized_embeddings(
1199
1200
        self, old_embeddings: nn.Embedding, new_num_tokens: Optional[int] = None
    ) -> nn.Embedding:
1201
1202
1203
        """
        Build a resized Embedding Module from a provided token Embedding Module. Increasing the size will add newly
        initialized vectors at the end. Reducing the size will remove vectors from the end
1204
1205

        Args:
1206
            old_embeddings (`torch.nn.Embedding`):
1207
                Old embeddings to be resized.
1208
            new_num_tokens (`int`, *optional*):
1209
                New number of tokens in the embedding matrix.
1210
1211

                Increasing the size will add newly initialized vectors at the end. Reducing the size will remove
1212
1213
                vectors from the end. If not provided or `None`, just returns a pointer to the input tokens
                ``torch.nn.Embedding``` module of the model without doing anything.
1214
1215

        Return:
1216
1217
            `torch.nn.Embedding`: Pointer to the resized Embedding Module or the old Embedding Module if
            `new_num_tokens` is `None`
1218
1219
1220
1221
        """
        if new_num_tokens is None:
            return old_embeddings

1222
1223
1224
1225
1226
1227
1228
1229
        if is_deepspeed_zero3_enabled():
            import deepspeed

            with deepspeed.zero.GatheredParameters(old_embeddings.weight, modifier_rank=None):
                old_num_tokens, old_embedding_dim = old_embeddings.weight.size()
        else:
            old_num_tokens, old_embedding_dim = old_embeddings.weight.size()

1230
1231
1232
        if old_num_tokens == new_num_tokens:
            return old_embeddings

1233
1234
        if not isinstance(old_embeddings, nn.Embedding):
            raise TypeError(
Sylvain Gugger's avatar
Sylvain Gugger committed
1235
1236
1237
                f"Old embeddings are of type {type(old_embeddings)}, which is not an instance of {nn.Embedding}. You"
                " should either use a different resize function or make sure that `old_embeddings` are an instance of"
                f" {nn.Embedding}."
1238
1239
            )

1240
        # Build new embeddings
1241
        new_embeddings = nn.Embedding(new_num_tokens, old_embedding_dim)
1242
        new_embeddings.to(old_embeddings.weight.device, dtype=old_embeddings.weight.dtype)
1243
1244
1245
1246

        # initialize all new embeddings (in particular added tokens)
        self._init_weights(new_embeddings)

1247
        # Copy token embeddings from the previous weights
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258

        # numbers of tokens to copy
        n = min(old_num_tokens, new_num_tokens)
        if is_deepspeed_zero3_enabled():
            import deepspeed

            with deepspeed.zero.GatheredParameters(old_embeddings.weight, modifier_rank=0):
                if torch.distributed.get_rank() == 0:
                    new_embeddings.weight.data[:n, :] = old_embeddings.weight.data[:n, :]
        else:
            new_embeddings.weight.data[:n, :] = old_embeddings.weight.data[:n, :]
1259
1260
1261

        return new_embeddings

1262
    def _get_resized_lm_head(
1263
1264
        self, old_lm_head: nn.Linear, new_num_tokens: Optional[int] = None, transposed: Optional[bool] = False
    ) -> nn.Linear:
1265
1266
1267
1268
1269
        """
        Build a resized Linear Module from a provided old Linear Module. Increasing the size will add newly initialized
        vectors at the end. Reducing the size will remove vectors from the end

        Args:
1270
            old_lm_head (`torch.nn.Linear`):
1271
                Old lm head liner layer to be resized.
1272
            new_num_tokens (`int`, *optional*):
1273
1274
1275
                New number of tokens in the linear matrix.

                Increasing the size will add newly initialized vectors at the end. Reducing the size will remove
1276
                vectors from the end. If not provided or `None`, just returns a pointer to the input tokens
Sylvain Gugger's avatar
Sylvain Gugger committed
1277
1278
1279
                ``torch.nn.Linear``` module of the model without doing anything. transposed (`bool`, *optional*,
                defaults to `False`): Whether `old_lm_head` is transposed or not. If True `old_lm_head.size()` is
                `lm_head_dim, vocab_size` else `vocab_size, lm_head_dim`.
1280
1281

        Return:
Sylvain Gugger's avatar
Sylvain Gugger committed
1282
1283
            `torch.nn.Linear`: Pointer to the resized Linear Module or the old Linear Module if `new_num_tokens` is
            `None`
1284
1285
1286
1287
        """
        if new_num_tokens is None:
            return old_lm_head

1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
        if is_deepspeed_zero3_enabled():
            import deepspeed

            with deepspeed.zero.GatheredParameters(old_lm_head.weight, modifier_rank=None):
                old_num_tokens, old_lm_head_dim = (
                    old_lm_head.weight.size() if not transposed else old_lm_head.weight.t().size()
                )
        else:
            old_num_tokens, old_lm_head_dim = (
                old_lm_head.weight.size() if not transposed else old_lm_head.weight.t().size()
            )
1299
1300
1301
1302
1303
1304

        if old_num_tokens == new_num_tokens:
            return old_lm_head

        if not isinstance(old_lm_head, nn.Linear):
            raise TypeError(
Sylvain Gugger's avatar
Sylvain Gugger committed
1305
1306
1307
                f"Old language model head is of type {type(old_lm_head)}, which is not an instance of {nn.Linear}. You"
                " should either use a different resize function or make sure that `old_lm_head` are an instance of"
                f" {nn.Linear}."
1308
1309
1310
1311
1312
            )

        # Build new lm head
        new_lm_head_shape = (old_lm_head_dim, new_num_tokens) if not transposed else (new_num_tokens, old_lm_head_dim)
        has_new_lm_head_bias = old_lm_head.bias is not None
1313
        new_lm_head = nn.Linear(*new_lm_head_shape, bias=has_new_lm_head_bias)
1314
        new_lm_head = new_lm_head.to(old_lm_head.weight.device, dtype=old_lm_head.weight.dtype)
1315
1316
1317
1318
1319
1320

        # initialize new lm head (in particular added tokens)
        self._init_weights(new_lm_head)

        num_tokens_to_copy = min(old_num_tokens, new_num_tokens)

1321
1322
1323
1324
        # XXX: put the long block of code in a wrapper
        if is_deepspeed_zero3_enabled():
            import deepspeed

1325
1326
            params = [old_lm_head.weight, old_lm_head.bias, new_lm_head.weight, new_lm_head.bias]
            with deepspeed.zero.GatheredParameters(params, modifier_rank=0):
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
                if torch.distributed.get_rank() == 0:
                    # Copy old lm head weights to new lm head
                    if not transposed:
                        new_lm_head.weight.data[:num_tokens_to_copy, :] = old_lm_head.weight.data[
                            :num_tokens_to_copy, :
                        ]
                    else:
                        new_lm_head.weight.data[:, :num_tokens_to_copy] = old_lm_head.weight.data[
                            :, :num_tokens_to_copy
                        ]

                    # Copy bias weights to new lm head
                    if has_new_lm_head_bias:
                        new_lm_head.bias.data[:num_tokens_to_copy] = old_lm_head.bias.data[:num_tokens_to_copy]
1341
        else:
1342
1343
1344
1345
1346
            # Copy old lm head weights to new lm head
            if not transposed:
                new_lm_head.weight.data[:num_tokens_to_copy, :] = old_lm_head.weight.data[:num_tokens_to_copy, :]
            else:
                new_lm_head.weight.data[:, :num_tokens_to_copy] = old_lm_head.weight.data[:, :num_tokens_to_copy]
1347

1348
1349
1350
            # Copy bias weights to new lm head
            if has_new_lm_head_bias:
                new_lm_head.bias.data[:num_tokens_to_copy] = old_lm_head.bias.data[:num_tokens_to_copy]
1351
1352
1353

        return new_lm_head

1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
    def resize_position_embeddings(self, new_num_position_embeddings: int):
        raise NotImplementedError(
            f"`resize_position_embeddings` is not implemented for {self.__class__}`. To implement it, you should "
            f"overwrite this method in the class {self.__class__} in `modeling_{self.__class__.__module__}.py`"
        )

    def get_position_embeddings(self) -> Union[nn.Embedding, Tuple[nn.Embedding]]:
        raise NotImplementedError(
            f"`get_position_embeddings` is not implemented for {self.__class__}`. To implement it, you should "
            f"overwrite this method in the class {self.__class__} in `modeling_{self.__class__.__module__}.py`"
        )

1366
    def init_weights(self):
1367
        """
1368
        If needed prunes and maybe initializes weights.
1369
        """
1370
1371
1372
1373
        # Prune heads if needed
        if self.config.pruned_heads:
            self.prune_heads(self.config.pruned_heads)

1374
1375
1376
1377
1378
1379
1380
        if _init_weights:
            # Initialize weights
            self.apply(self._init_weights)

            # Tie weights should be skipped when not initializing all weights
            # since from_pretrained(...) calls tie weights anyways
            self.tie_weights()
1381

1382
1383
1384
    def prune_heads(self, heads_to_prune: Dict[int, List[int]]):
        """
        Prunes heads of the base model.
1385

1386
        Arguments:
1387
            heads_to_prune (`Dict[int, List[int]]`):
Sylvain Gugger's avatar
Sylvain Gugger committed
1388
1389
1390
                Dictionary with keys being selected layer indices (`int`) and associated values being the list of heads
                to prune in said layer (list of `int`). For instance {1: [0, 2], 2: [2, 3]} will prune heads 0 and 2 on
                layer 1 and heads 2 and 3 on layer 2.
thomwolf's avatar
thomwolf committed
1391
        """
1392
        # save new sets of pruned heads as union of previously stored pruned heads and newly pruned heads
1393
        for layer, heads in heads_to_prune.items():
1394
1395
1396
            union_heads = set(self.config.pruned_heads.get(layer, [])) | set(heads)
            self.config.pruned_heads[layer] = list(union_heads)  # Unfortunately we have to store it as list for JSON

1397
        self.base_model._prune_heads(heads_to_prune)
thomwolf's avatar
thomwolf committed
1398

1399
    def gradient_checkpointing_enable(self):
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
        """
        Activates gradient checkpointing for the current model.

        Note that in other frameworks this feature can be referred to as "activation checkpointing" or "checkpoint
        activations".
        """
        if not self.supports_gradient_checkpointing:
            raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.")
        self.apply(partial(self._set_gradient_checkpointing, value=True))

1410
    def gradient_checkpointing_disable(self):
1411
1412
1413
1414
1415
1416
1417
1418
1419
        """
        Deactivates gradient checkpointing for the current model.

        Note that in other frameworks this feature can be referred to as "activation checkpointing" or "checkpoint
        activations".
        """
        if self.supports_gradient_checkpointing:
            self.apply(partial(self._set_gradient_checkpointing, value=False))

1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
    @property
    def is_gradient_checkpointing(self) -> bool:
        """
        Whether gradient checkpointing is activated for this model or not.

        Note that in other frameworks this feature can be referred to as "activation checkpointing" or "checkpoint
        activations".
        """
        return any(hasattr(m, "gradient_checkpointing") and m.gradient_checkpointing for m in self.modules())

1430
1431
1432
    def save_pretrained(
        self,
        save_directory: Union[str, os.PathLike],
1433
        is_main_process: bool = True,
1434
1435
        state_dict: Optional[dict] = None,
        save_function: Callable = torch.save,
Sylvain Gugger's avatar
Sylvain Gugger committed
1436
        push_to_hub: bool = False,
Sylvain Gugger's avatar
Sylvain Gugger committed
1437
        max_shard_size: Union[int, str] = "10GB",
Sylvain Gugger's avatar
Sylvain Gugger committed
1438
        **kwargs,
1439
    ):
1440
1441
        """
        Save a model and its configuration file to a directory, so that it can be re-loaded using the
1442
        `[`~PreTrainedModel.from_pretrained`]` class method.
1443

1444
        Arguments:
1445
            save_directory (`str` or `os.PathLike`):
1446
                Directory to which to save. Will be created if it doesn't exist.
1447
1448
1449
1450
            is_main_process (`bool`, *optional*, defaults to `True`):
                Whether the process calling this is the main process or not. Useful when in distributed training like
                TPUs and need to call this function on all processes. In this case, set `is_main_process=True` only on
                the main process to avoid race conditions.
1451
            state_dict (nested dictionary of `torch.Tensor`):
Sylvain Gugger's avatar
Sylvain Gugger committed
1452
1453
1454
                The state dictionary of the model to save. Will default to `self.state_dict()`, but can be used to only
                save parts of the model or if special precautions need to be taken when recovering the state dictionary
                of a model (like when using model parallelism).
1455
            save_function (`Callable`):
1456
                The function to use to save the state dictionary. Useful on distributed training like TPUs when one
1457
1458
                need to replace `torch.save` by another method.
            push_to_hub (`bool`, *optional*, defaults to `False`):
Sylvain Gugger's avatar
Sylvain Gugger committed
1459
                Whether or not to push your model to the Hugging Face model hub after saving it.
1460

1461
                <Tip warning={true}>
1462

Sylvain Gugger's avatar
Sylvain Gugger committed
1463
1464
1465
                Using `push_to_hub=True` will synchronize the repository you are pushing to with `save_directory`,
                which requires `save_directory` to be a local clone of the repo you are pushing to if it's an existing
                folder. Pass along `temp_dir=True` to use a temporary directory instead.
1466
1467

                </Tip>
1468

Sylvain Gugger's avatar
Sylvain Gugger committed
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
            max_shard_size (`int` or `str`, *optional*, defaults to `"10GB"`):
                The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size
                lower than this size. If expressed as a string, needs to be digits followed by a unit (like `"5MB"`).

                <Tip warning={true}>

                If a single weight of the model is bigger than `max_shard_size`, it will be in its own checkpoint shard
                which will be bigger than `max_shard_size`.

                </Tip>

Sylvain Gugger's avatar
Sylvain Gugger committed
1480
            kwargs:
1481
                Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
1482
        """
1483
1484
1485
1486
1487
1488
        if "save_config" in kwargs:
            warnings.warn(
                "`save_config` is deprecated and will be removed in v5 of Transformers. Use `is_main_process` instead."
            )
            is_main_process = kwargs.pop("save_config")

1489
        if os.path.isfile(save_directory):
1490
            logger.error(f"Provided path ({save_directory}) should be a directory, not a file")
1491
            return
1492
1493
1494
1495
1496

        if push_to_hub:
            commit_message = kwargs.pop("commit_message", None)
            repo = self._create_or_get_repo(save_directory, **kwargs)

1497
        os.makedirs(save_directory, exist_ok=True)
1498

Julien Chaumond's avatar
Julien Chaumond committed
1499
        # Only save the model itself if we are using distributed training
1500
        model_to_save = unwrap_model(self)
1501

1502
1503
1504
1505
1506
        # save the string version of dtype to the config, e.g. convert torch.float32 => "float32"
        # we currently don't use this setting automatically, but may start to use with v5
        dtype = get_parameter_dtype(model_to_save)
        model_to_save.config.torch_dtype = str(dtype).split(".")[1]

Julien Chaumond's avatar
Julien Chaumond committed
1507
1508
1509
        # Attach architecture to the config
        model_to_save.config.architectures = [model_to_save.__class__.__name__]

1510
1511
1512
1513
1514
        # If we have a custom model, we copy the file defining it in the folder and set the attributes so it can be
        # loaded from the Hub.
        if self._auto_class is not None:
            custom_object_save(self, save_directory, config=self.config)

1515
        # Save the config
1516
        if is_main_process:
1517
1518
1519
1520
1521
            model_to_save.config.save_pretrained(save_directory)

        # Save the model
        if state_dict is None:
            state_dict = model_to_save.state_dict()
1522
1523

        # Handle the case where some state_dict keys shouldn't be saved
1524
        if self._keys_to_ignore_on_save is not None:
1525
            for ignore_key in self._keys_to_ignore_on_save:
1526
1527
                if ignore_key in state_dict.keys():
                    del state_dict[ignore_key]
1528

Sylvain Gugger's avatar
Sylvain Gugger committed
1529
1530
1531
1532
1533
1534
        # Shard the model if it is too big.
        shards, index = shard_checkpoint(state_dict, max_shard_size=max_shard_size)

        # Clean the folder from a previous save
        for filename in os.listdir(save_directory):
            full_filename = os.path.join(save_directory, filename)
1535
1536
1537
1538
1539
1540
1541
1542
            # If we have a shard file that is not going to be replaced, we delete it, but only from the main process
            # in distributed settings to avoid race conditions.
            if (
                filename.startswith(WEIGHTS_NAME[:-4])
                and os.path.isfile(full_filename)
                and filename not in shards.keys()
                and is_main_process
            ):
Sylvain Gugger's avatar
Sylvain Gugger committed
1543
                os.remove(full_filename)
1544

Sylvain Gugger's avatar
Sylvain Gugger committed
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
        # Save the model
        for shard_file, shard in shards.items():
            save_function(shard, os.path.join(save_directory, shard_file))

        if index is None:
            logger.info(f"Model weights saved in {os.path.join(save_directory, WEIGHTS_NAME)}")
        else:
            save_index_file = os.path.join(save_directory, WEIGHTS_INDEX_NAME)
            # Save the index as well
            with open(save_index_file, "w", encoding="utf-8") as f:
                content = json.dumps(index, indent=2, sort_keys=True) + "\n"
                f.write(content)
            logger.info(
                f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be "
                f"split in {len(shards)} checkpoint shards. You can find where each parameters has been saved in the "
                f"index located at {save_index_file}."
            )
1562

Sylvain Gugger's avatar
Sylvain Gugger committed
1563
        if push_to_hub:
1564
            url = self._push_to_hub(repo, commit_message=commit_message)
Sylvain Gugger's avatar
Sylvain Gugger committed
1565
1566
            logger.info(f"Model pushed to the hub in this commit: {url}")

1567
    @classmethod
1568
    def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *model_args, **kwargs):
1569
1570
        r"""
        Instantiate a pretrained pytorch model from a pre-trained model configuration.
1571

Sylvain Gugger's avatar
Sylvain Gugger committed
1572
1573
        The model is set in evaluation mode by default using `model.eval()` (Dropout modules are deactivated). To train
        the model, you should first set it back in training mode with `model.train()`.
1574

1575
        The warning *Weights from XXX not initialized from pretrained model* means that the weights of XXX do not come
1576
1577
        pretrained with the rest of the model. It is up to you to train those weights with a downstream fine-tuning
        task.
1578

1579
        The warning *Weights from XXX not used in YYY* means that the layer XXX is not used by YYY, therefore those
1580
        weights are discarded.
1581

1582
        Parameters:
1583
            pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*):
1584
1585
                Can be either:

1586
                    - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
Sylvain Gugger's avatar
Sylvain Gugger committed
1587
1588
                      Valid model ids can be located at the root-level, like `bert-base-uncased`, or namespaced under a
                      user or organization name, like `dbmdz/bert-base-german-cased`.
1589
1590
1591
                    - A path to a *directory* containing model weights saved using
                      [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
                    - A path or url to a *tensorflow index checkpoint file* (e.g, `./tf_model/model.ckpt.index`). In
Sylvain Gugger's avatar
Sylvain Gugger committed
1592
1593
1594
                      this case, `from_tf` should be set to `True` and a configuration object should be provided as
                      `config` argument. This loading path is slower than converting the TensorFlow checkpoint in a
                      PyTorch model using the provided conversion scripts and loading the PyTorch model afterwards.
1595
                    - A path or url to a model folder containing a *flax checkpoint file* in *.msgpack* format (e.g,
Sylvain Gugger's avatar
Sylvain Gugger committed
1596
1597
                      `./flax_model/` containing `flax_model.msgpack`). In this case, `from_flax` should be set to
                      `True`.
1598
1599
1600
1601
1602
                    - `None` if you are both providing the configuration and state dictionary (resp. with keyword
                      arguments `config` and `state_dict`).
            model_args (sequence of positional arguments, *optional*):
                All remaining positional arguments will be passed to the underlying model's `__init__` method.
            config (`Union[PretrainedConfig, str, os.PathLike]`, *optional*):
1603
1604
                Can be either:

1605
1606
                    - an instance of a class derived from [`PretrainedConfig`],
                    - a string or path valid as input to [`~PretrainedConfig.from_pretrained`].
1607

1608
                Configuration for the model to use instead of an automatically loaded configuration. Configuration can
1609
1610
                be automatically loaded when:

1611
                    - The model is a model provided by the library (loaded with the *model id* string of a pretrained
1612
                      model).
Sylvain Gugger's avatar
Sylvain Gugger committed
1613
1614
                    - The model was saved using [`~PreTrainedModel.save_pretrained`] and is reloaded by supplying the
                      save directory.
1615
1616
1617
                    - The model is loaded by supplying a local directory as `pretrained_model_name_or_path` and a
                      configuration JSON file named *config.json* is found in the directory.
            state_dict (`Dict[str, torch.Tensor]`, *optional*):
1618
1619
1620
                A state dictionary to use instead of a state dictionary loaded from saved weights file.

                This option can be used if you want to create a model from a pretrained configuration but load your own
Sylvain Gugger's avatar
Sylvain Gugger committed
1621
                weights. In this case though, you should check if using [`~PreTrainedModel.save_pretrained`] and
1622
1623
                [`~PreTrainedModel.from_pretrained`] is not a simpler option.
            cache_dir (`Union[str, os.PathLike]`, *optional*):
1624
1625
                Path to a directory in which a downloaded pretrained model configuration should be cached if the
                standard cache should not be used.
1626
            from_tf (`bool`, *optional*, defaults to `False`):
1627
                Load the model weights from a TensorFlow checkpoint save file (see docstring of
1628
1629
                `pretrained_model_name_or_path` argument).
            from_flax (`bool`, *optional*, defaults to `False`):
1630
                Load the model weights from a Flax checkpoint save file (see docstring of
1631
1632
                `pretrained_model_name_or_path` argument).
            ignore_mismatched_sizes (`bool`, *optional*, defaults to `False`):
1633
1634
1635
                Whether or not to raise an error if some of the weights from the checkpoint do not have the same size
                as the weights of the model (if for instance, you are instantiating a model with 10 labels from a
                checkpoint with 3 labels).
1636
            force_download (`bool`, *optional*, defaults to `False`):
1637
1638
                Whether or not to force the (re-)download of the model weights and configuration files, overriding the
                cached versions if they exist.
1639
            resume_download (`bool`, *optional*, defaults to `False`):
1640
1641
                Whether or not to delete incompletely received files. Will attempt to resume the download if such a
                file exists.
1642
            proxies (`Dict[str, str]`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
1643
1644
                A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
                'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
1645
            output_loading_info(`bool`, *optional*, defaults to `False`):
Sylvain Gugger's avatar
Sylvain Gugger committed
1646
                Whether ot not to also return a dictionary containing missing keys, unexpected keys and error messages.
1647
            local_files_only(`bool`, *optional*, defaults to `False`):
Stas Bekman's avatar
Stas Bekman committed
1648
                Whether or not to only look at local files (i.e., do not try to download the model).
1649
            use_auth_token (`str` or *bool*, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
1650
1651
                The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
                when running `transformers-cli login` (stored in `~/.huggingface`).
1652
            revision (`str`, *optional*, defaults to `"main"`):
Julien Chaumond's avatar
Julien Chaumond committed
1653
                The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
1654
                git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
Julien Chaumond's avatar
Julien Chaumond committed
1655
                identifier allowed by git.
1656
            mirror (`str`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
1657
1658
1659
                Mirror source to accelerate downloads in China. If you are from China and have an accessibility
                problem, you can set this option to resolve it. Note that we do not guarantee the timeliness or safety.
                Please refer to the mirror site for more information.
1660
            _fast_init(`bool`, *optional*, defaults to `True`):
1661
1662
                Whether or not to disable fast initialization.

1663
1664
                <Tip warning={true}>

Sylvain Gugger's avatar
Sylvain Gugger committed
1665
1666
1667
                One should only disable *_fast_init* to ensure backwards compatibility with `transformers.__version__ <
                4.6.0` for seeded model initialization. This argument will be removed at the next major version. See
                [pull request 11471](https://github.com/huggingface/transformers/pull/11471) for more information.
1668

1669
                </Tip>
1670

1671
1672
1673
            > Parameters for big model inference

            low_cpu_mem_usage(`bool`, *optional*):
1674
1675
1676
1677
1678
                Tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model.
                This is an experimental feature and a subject to change at any moment.
            torch_dtype (`str` or `torch.dtype`, *optional*):
                Override the default `torch.dtype` and load the model under this dtype. If `"auto"` is passed the dtype
                will be automatically derived from the model's weights.
1679
1680
1681
1682
1683
1684
            device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*):
                A map that specifies where each submodule should go. It doesn't need to be refined to each
                parameter/buffer name, once a given module name is inside, every submodule of it will be sent to the
                same device.

                To have Accelerate compute the most optimized `device_map` automatically, set `device_map="auto"`.
1685
1686
1687
            max_memory (`Dict`, *optional*):
                A dictionary device identifier to maximum memory. Will default to the maximum memory available for each
                GPU and the available CPU RAM if unset.
1688
1689
            offload_folder (`str` or `os.PathLike`, *optional*):
                If the `device_map` contains any value `"disk"`, the folder where we will offload weights.
1690
            offload_state_dict (`bool`, *optional*):
1691
                If `True`, will temporarily offload the CPU state dict to the hard drive to avoid getting out of CPU
1692
1693
                RAM if the weight of the CPU state dict + the biggest shard of the checkpoint does not fit. Defaults to
                `True` when there is some disk offload.
1694

1695
            kwargs (remaining dictionary of keyword arguments, *optional*):
1696
                Can be used to update the configuration object (after it being loaded) and initiate the model (e.g.,
1697
                `output_attentions=True`). Behaves differently depending on whether a `config` is provided or
1698
1699
                automatically loaded:

1700
1701
                    - If a configuration is provided with `config`, `**kwargs` will be directly passed to the
                      underlying model's `__init__` method (we assume all relevant updates to the configuration have
1702
                      already been done)
1703
                    - If a configuration is not provided, `kwargs` will be first passed to the configuration class
Sylvain Gugger's avatar
Sylvain Gugger committed
1704
1705
1706
1707
                      initialization function ([`~PretrainedConfig.from_pretrained`]). Each key of `kwargs` that
                      corresponds to a configuration attribute will be used to override said attribute with the
                      supplied `kwargs` value. Remaining keys that do not correspond to any configuration attribute
                      will be passed to the underlying model's `__init__` function.
1708
1709
1710
1711
1712
1713
1714
1715
1716

        <Tip>

        Passing `use_auth_token=True`` is required when you want to use a private model.

        </Tip>

        <Tip>

Sylvain Gugger's avatar
Sylvain Gugger committed
1717
1718
        Activate the special ["offline-mode"](https://huggingface.co/transformers/installation.html#offline-mode) to
        use this method in a firewalled environment.
1719
1720
1721
1722
1723
1724
1725

        </Tip>

        Examples:

        ```python
        >>> from transformers import BertConfig, BertModel
Sylvain Gugger's avatar
Sylvain Gugger committed
1726

1727
        >>> # Download model and configuration from huggingface.co and cache.
Sylvain Gugger's avatar
Sylvain Gugger committed
1728
        >>> model = BertModel.from_pretrained("bert-base-uncased")
1729
        >>> # Model was saved using *save_pretrained('./test/saved_model/')* (for example purposes, not runnable).
Sylvain Gugger's avatar
Sylvain Gugger committed
1730
        >>> model = BertModel.from_pretrained("./test/saved_model/")
1731
        >>> # Update configuration during loading.
Sylvain Gugger's avatar
Sylvain Gugger committed
1732
        >>> model = BertModel.from_pretrained("bert-base-uncased", output_attentions=True)
1733
1734
        >>> assert model.config.output_attentions == True
        >>> # Loading from a TF checkpoint file instead of a PyTorch model (slower, for example purposes, not runnable).
Sylvain Gugger's avatar
Sylvain Gugger committed
1735
1736
        >>> config = BertConfig.from_json_file("./tf_model/my_tf_model_config.json")
        >>> model = BertModel.from_pretrained("./tf_model/my_tf_checkpoint.ckpt.index", from_tf=True, config=config)
1737
        >>> # Loading from a Flax checkpoint file instead of a PyTorch model (slower)
Sylvain Gugger's avatar
Sylvain Gugger committed
1738
        >>> model = BertModel.from_pretrained("bert-base-uncased", from_flax=True)
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
        ```

        * `low_cpu_mem_usage` algorithm:

        This is an experimental function that loads the model using ~1x model size CPU memory

        Here is how it works:

        1. save which state_dict keys we have
        2. drop state_dict before the model is created, since the latter takes 1x model size CPU memory
        3. after the model has been instantiated switch to the meta device all params/buffers that
        are going to be replaced from the loaded state_dict
        4. load state_dict 2nd time
        5. replace the params/buffers from the state_dict

        Currently, it can't handle deepspeed ZeRO stage 3 and ignores loading errors

        """
1757
1758
1759
1760
        config = kwargs.pop("config", None)
        state_dict = kwargs.pop("state_dict", None)
        cache_dir = kwargs.pop("cache_dir", None)
        from_tf = kwargs.pop("from_tf", False)
1761
        from_flax = kwargs.pop("from_flax", False)
1762
        ignore_mismatched_sizes = kwargs.pop("ignore_mismatched_sizes", False)
1763
1764
1765
1766
        force_download = kwargs.pop("force_download", False)
        resume_download = kwargs.pop("resume_download", False)
        proxies = kwargs.pop("proxies", None)
        output_loading_info = kwargs.pop("output_loading_info", False)
1767
        local_files_only = kwargs.pop("local_files_only", False)
1768
        use_auth_token = kwargs.pop("use_auth_token", None)
Julien Chaumond's avatar
Julien Chaumond committed
1769
        revision = kwargs.pop("revision", None)
1770
        mirror = kwargs.pop("mirror", None)
1771
1772
        from_pipeline = kwargs.pop("_from_pipeline", None)
        from_auto_class = kwargs.pop("_from_auto", False)
1773
        _fast_init = kwargs.pop("_fast_init", True)
1774
        torch_dtype = kwargs.pop("torch_dtype", None)
1775
1776
        low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", None)
        device_map = kwargs.pop("device_map", None)
1777
        max_memory = kwargs.pop("max_memory", None)
1778
        offload_folder = kwargs.pop("offload_folder", None)
1779
        offload_state_dict = kwargs.pop("offload_state_dict", None)
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
1797
1798

        if device_map is not None:
            if low_cpu_mem_usage is None:
                low_cpu_mem_usage = True
            elif not low_cpu_mem_usage:
                raise ValueError("Passing along a `device_map` requires `low_cpu_mem_usage=True`")

        if low_cpu_mem_usage:
            # low_cpu_mem_usage requires PyTorch >= 1.9 to have the meta device.
            require_version_core("torch>=1.9")

            if is_deepspeed_zero3_enabled():
                raise ValueError(
                    "DeepSpeed Zero-3 is not compatible with `low_cpu_mem_usage=True` or with passing a `device_map`."
                )
            elif not is_accelerate_available():
                raise ImportError(
                    "Using `low_cpu_mem_usage=True` or a `device_map` requires Accelerate: `pip install accelerate`"
                )
1799
1800

        from_pt = not (from_tf | from_flax)
1801
1802
1803
1804

        user_agent = {"file_type": "model", "framework": "pytorch", "from_auto_class": from_auto_class}
        if from_pipeline is not None:
            user_agent["using_pipeline"] = from_pipeline
thomwolf's avatar
thomwolf committed
1805

1806
1807
1808
1809
        if is_offline_mode() and not local_files_only:
            logger.info("Offline mode: forcing local_files_only=True")
            local_files_only = True

1810
1811
1812
        # Load config if we don't provide a configuration
        if not isinstance(config, PretrainedConfig):
            config_path = config if config is not None else pretrained_model_name_or_path
1813
            config, model_kwargs = cls.config_class.from_pretrained(
1814
1815
1816
                config_path,
                cache_dir=cache_dir,
                return_unused_kwargs=True,
1817
                force_download=force_download,
1818
                resume_download=resume_download,
1819
                proxies=proxies,
1820
                local_files_only=local_files_only,
1821
                use_auth_token=use_auth_token,
Julien Chaumond's avatar
Julien Chaumond committed
1822
                revision=revision,
1823
1824
                _from_auto=from_auto_class,
                _from_pipeline=from_pipeline,
1825
                **kwargs,
1826
1827
1828
            )
        else:
            model_kwargs = kwargs
1829

Sylvain Gugger's avatar
Sylvain Gugger committed
1830
1831
1832
1833
        # This variable will flag if we're loading a sharded checkpoint. In this case the archive file is just the
        # index of the files.
        is_sharded = False
        sharded_metadata = None
thomwolf's avatar
thomwolf committed
1834
        # Load model
Yih-Dar's avatar
Yih-Dar committed
1835
1836
        loading_info = None

thomwolf's avatar
thomwolf committed
1837
        if pretrained_model_name_or_path is not None:
1838
            pretrained_model_name_or_path = str(pretrained_model_name_or_path)
1839
            if os.path.isdir(pretrained_model_name_or_path):
thomwolf's avatar
thomwolf committed
1840
                if from_tf and os.path.isfile(os.path.join(pretrained_model_name_or_path, TF_WEIGHTS_NAME + ".index")):
1841
                    # Load from a TF 1.0 checkpoint in priority if from_tf
thomwolf's avatar
thomwolf committed
1842
                    archive_file = os.path.join(pretrained_model_name_or_path, TF_WEIGHTS_NAME + ".index")
thomwolf's avatar
thomwolf committed
1843
                elif from_tf and os.path.isfile(os.path.join(pretrained_model_name_or_path, TF2_WEIGHTS_NAME)):
1844
                    # Load from a TF 2.0 checkpoint in priority if from_tf
thomwolf's avatar
thomwolf committed
1845
                    archive_file = os.path.join(pretrained_model_name_or_path, TF2_WEIGHTS_NAME)
1846
1847
1848
                elif from_flax and os.path.isfile(os.path.join(pretrained_model_name_or_path, FLAX_WEIGHTS_NAME)):
                    # Load from a Flax checkpoint in priority if from_flax
                    archive_file = os.path.join(pretrained_model_name_or_path, FLAX_WEIGHTS_NAME)
thomwolf's avatar
thomwolf committed
1849
1850
                elif os.path.isfile(os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)):
                    # Load from a PyTorch checkpoint
thomwolf's avatar
thomwolf committed
1851
                    archive_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)
Sylvain Gugger's avatar
Sylvain Gugger committed
1852
1853
1854
1855
                elif os.path.isfile(os.path.join(pretrained_model_name_or_path, WEIGHTS_INDEX_NAME)):
                    # Load from a sharded PyTorch checkpoint
                    archive_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_INDEX_NAME)
                    is_sharded = True
1856
1857
1858
1859
1860
1861
1862
1863
1864
                # At this stage we don't have a weight file so we will raise an error.
                elif os.path.isfile(
                    os.path.join(pretrained_model_name_or_path, TF_WEIGHTS_NAME + ".index")
                ) or os.path.isfile(os.path.join(pretrained_model_name_or_path, TF2_WEIGHTS_NAME)):
                    raise EnvironmentError(
                        f"Error no file named {WEIGHTS_NAME} found in directory {pretrained_model_name_or_path} but "
                        "there is a file for TensorFlow weights. Use `from_tf=True` to load this model from those "
                        "weights."
                    )
Nathan Dahlberg's avatar
Nathan Dahlberg committed
1865
                elif os.path.isfile(os.path.join(pretrained_model_name_or_path, FLAX_WEIGHTS_NAME)):
1866
1867
1868
1869
1870
                    raise EnvironmentError(
                        f"Error no file named {WEIGHTS_NAME} found in directory {pretrained_model_name_or_path} but "
                        "there is a file for Flax weights. Use `from_flax=True` to load this model from those "
                        "weights."
                    )
thomwolf's avatar
thomwolf committed
1871
                else:
1872
                    raise EnvironmentError(
1873
1874
                        f"Error no file named {WEIGHTS_NAME}, {TF2_WEIGHTS_NAME}, {TF_WEIGHTS_NAME + '.index'} or "
                        f"{FLAX_WEIGHTS_NAME} found in directory {pretrained_model_name_or_path}."
1875
                    )
1876
            elif os.path.isfile(pretrained_model_name_or_path) or is_remote_url(pretrained_model_name_or_path):
1877
                archive_file = pretrained_model_name_or_path
1878
            elif os.path.isfile(pretrained_model_name_or_path + ".index"):
1879
1880
1881
1882
1883
                if not from_tf:
                    raise ValueError(
                        f"We found a TensorFlow checkpoint at {pretrained_model_name_or_path + '.index'}, please set "
                        "from_tf to True to load from this checkpoint."
                    )
1884
                archive_file = pretrained_model_name_or_path + ".index"
1885
            else:
1886
1887
1888
1889
1890
1891
1892
1893
                # set correct filename
                if from_tf:
                    filename = TF2_WEIGHTS_NAME
                elif from_flax:
                    filename = FLAX_WEIGHTS_NAME
                else:
                    filename = WEIGHTS_NAME

thomwolf's avatar
thomwolf committed
1894
                archive_file = hf_bucket_url(
1895
                    pretrained_model_name_or_path, filename=filename, revision=revision, mirror=mirror
thomwolf's avatar
thomwolf committed
1896
                )
1897

thomwolf's avatar
thomwolf committed
1898
            try:
1899
                # Load from URL or cache if already cached
1900
1901
1902
1903
1904
1905
                resolved_archive_file = cached_path(
                    archive_file,
                    cache_dir=cache_dir,
                    force_download=force_download,
                    proxies=proxies,
                    resume_download=resume_download,
1906
                    local_files_only=local_files_only,
1907
                    use_auth_token=use_auth_token,
1908
                    user_agent=user_agent,
1909
                )
1910

1911
            except RepositoryNotFoundError:
1912
1913
1914
1915
1916
1917
                raise EnvironmentError(
                    f"{pretrained_model_name_or_path} is not a local folder and is not a valid model identifier "
                    "listed on 'https://huggingface.co/models'\nIf this is a private repository, make sure to pass a "
                    "token having permission to this repo with `use_auth_token` or log in with `huggingface-cli "
                    "login` and pass `use_auth_token=True`."
                )
1918
            except RevisionNotFoundError:
1919
1920
1921
1922
1923
                raise EnvironmentError(
                    f"{revision} is not a valid git identifier (branch name, tag name or commit id) that exists for "
                    "this model name. Check the model page at "
                    f"'https://huggingface.co/{pretrained_model_name_or_path}' for available revisions."
                )
1924
            except EntryNotFoundError:
1925
                if filename == WEIGHTS_NAME:
Sylvain Gugger's avatar
Sylvain Gugger committed
1926
1927
1928
1929
1930
1931
1932
                    try:
                        # Maybe the checkpoint is sharded, we try to grab the index name in this case.
                        archive_file = hf_bucket_url(
                            pretrained_model_name_or_path,
                            filename=WEIGHTS_INDEX_NAME,
                            revision=revision,
                            mirror=mirror,
1933
                        )
Sylvain Gugger's avatar
Sylvain Gugger committed
1934
1935
1936
1937
1938
1939
1940
1941
1942
                        resolved_archive_file = cached_path(
                            archive_file,
                            cache_dir=cache_dir,
                            force_download=force_download,
                            proxies=proxies,
                            resume_download=resume_download,
                            local_files_only=local_files_only,
                            use_auth_token=use_auth_token,
                            user_agent=user_agent,
1943
                        )
Sylvain Gugger's avatar
Sylvain Gugger committed
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
                        is_sharded = True
                    except EntryNotFoundError:
                        # Otherwise, maybe there is a TF or Flax model file.  We try those to give a helpful error
                        # message.
                        has_file_kwargs = {
                            "revision": revision,
                            "mirror": mirror,
                            "proxies": proxies,
                            "use_auth_token": use_auth_token,
                        }
                        if has_file(pretrained_model_name_or_path, TF2_WEIGHTS_NAME, **has_file_kwargs):
                            raise EnvironmentError(
Sylvain Gugger's avatar
Sylvain Gugger committed
1956
1957
1958
                                f"{pretrained_model_name_or_path} does not appear to have a file named"
                                f" {WEIGHTS_NAME} but there is a file for TensorFlow weights. Use `from_tf=True` to"
                                " load this model from those weights."
Sylvain Gugger's avatar
Sylvain Gugger committed
1959
1960
1961
                            )
                        elif has_file(pretrained_model_name_or_path, FLAX_WEIGHTS_NAME, **has_file_kwargs):
                            raise EnvironmentError(
Sylvain Gugger's avatar
Sylvain Gugger committed
1962
1963
1964
                                f"{pretrained_model_name_or_path} does not appear to have a file named"
                                f" {WEIGHTS_NAME} but there is a file for Flax weights. Use `from_flax=True` to load"
                                " this model from those weights."
Sylvain Gugger's avatar
Sylvain Gugger committed
1965
1966
1967
                            )
                        else:
                            raise EnvironmentError(
Sylvain Gugger's avatar
Sylvain Gugger committed
1968
1969
                                f"{pretrained_model_name_or_path} does not appear to have a file named {WEIGHTS_NAME},"
                                f" {TF2_WEIGHTS_NAME}, {TF_WEIGHTS_NAME} or {FLAX_WEIGHTS_NAME}."
Sylvain Gugger's avatar
Sylvain Gugger committed
1970
                            )
1971
1972
1973
1974
                else:
                    raise EnvironmentError(
                        f"{pretrained_model_name_or_path} does not appear to have a file named {filename}."
                    )
1975
            except HTTPError as err:
1976
                raise EnvironmentError(
1977
1978
1979
1980
1981
                    f"There was a specific connection error when trying to load {pretrained_model_name_or_path}:\n"
                    f"{err}"
                )
            except ValueError:
                raise EnvironmentError(
Sylvain Gugger's avatar
Sylvain Gugger committed
1982
1983
1984
1985
1986
                    f"We couldn't connect to '{HUGGINGFACE_CO_RESOLVE_ENDPOINT}' to load this model, couldn't find it"
                    f" in the cached files and it looks like {pretrained_model_name_or_path} is not the path to a"
                    f" directory containing a file named {WEIGHTS_NAME}, {TF2_WEIGHTS_NAME}, {TF_WEIGHTS_NAME} or"
                    f" {FLAX_WEIGHTS_NAME}.\nCheckout your internet connection or see how to run the library in"
                    " offline mode at 'https://huggingface.co/docs/transformers/installation#offline-mode'."
1987
                )
1988
            except EnvironmentError:
1989
1990
1991
1992
1993
1994
                raise EnvironmentError(
                    f"Can't load the model for '{pretrained_model_name_or_path}'. If you were trying to load it from "
                    "'https://huggingface.co/models', make sure you don't have a local directory with the same name. "
                    f"Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory "
                    f"containing a file named {WEIGHTS_NAME}, {TF2_WEIGHTS_NAME}, {TF_WEIGHTS_NAME} or "
                    f"{FLAX_WEIGHTS_NAME}."
1995
                )
1996

thomwolf's avatar
thomwolf committed
1997
            if resolved_archive_file == archive_file:
1998
                logger.info(f"loading weights file {archive_file}")
1999
            else:
2000
                logger.info(f"loading weights file {archive_file} from cache at {resolved_archive_file}")
2001
        else:
thomwolf's avatar
thomwolf committed
2002
            resolved_archive_file = None
2003

Sylvain Gugger's avatar
Sylvain Gugger committed
2004
2005
2006
2007
2008
2009
2010
2011
2012
2013
2014
2015
2016
2017
2018
2019
2020
        # We'll need to download and cache each checkpoint shard if the checkpoint is sharded.
        if is_sharded:
            # resolved_archive_file becomes a list of files that point to the different checkpoint shards in this case.
            resolved_archive_file, sharded_metadata = get_checkpoint_shard_files(
                pretrained_model_name_or_path,
                resolved_archive_file,
                cache_dir=cache_dir,
                force_download=force_download,
                proxies=proxies,
                resume_download=resume_download,
                local_files_only=local_files_only,
                use_auth_token=use_auth_token,
                user_agent=user_agent,
                revision=revision,
                mirror=mirror,
            )

2021
2022
        # load pt weights early so that we know which dtype to init the model under
        if from_pt:
2023
            if not is_sharded and state_dict is None:
Sylvain Gugger's avatar
Sylvain Gugger committed
2024
2025
                # Time to load the checkpoint
                state_dict = load_state_dict(resolved_archive_file)
2026

2027
2028
2029
            # set dtype to instantiate the model under:
            # 1. If torch_dtype is not None, we use that dtype
            # 2. If torch_dtype is "auto", we auto-detect dtype from the loaded state_dict, by checking its first
2030
            #    weights entry that is of a floating type - we assume all floating dtype weights are of the same dtype
2031
2032
2033
2034
2035
            # we also may have config.torch_dtype available, but we won't rely on it till v5
            dtype_orig = None
            if torch_dtype is not None:
                if isinstance(torch_dtype, str):
                    if torch_dtype == "auto":
Sylvain Gugger's avatar
Sylvain Gugger committed
2036
2037
2038
                        if is_sharded and "dtype" in sharded_metadata:
                            torch_dtype = sharded_metadata["dtype"]
                        elif not is_sharded:
2039
                            torch_dtype = get_state_dict_dtype(state_dict)
Sylvain Gugger's avatar
Sylvain Gugger committed
2040
2041
                        else:
                            one_state_dict = load_state_dict(resolved_archive_file)
2042
                            torch_dtype = get_state_dict_dtype(one_state_dict)
Sylvain Gugger's avatar
Sylvain Gugger committed
2043
                            del one_state_dict  # free CPU memory
2044
2045
2046
2047
2048
2049
                    else:
                        raise ValueError(
                            f"`torch_dtype` can be either a `torch.dtype` or `auto`, but received {torch_dtype}"
                        )
                dtype_orig = cls._set_default_torch_dtype(torch_dtype)

2050
2051
2052
2053
            if is_sharded:
                loaded_state_dict_keys = sharded_metadata["all_checkpoint_keys"]
            else:
                loaded_state_dict_keys = [k for k in state_dict.keys()]
2054
            if low_cpu_mem_usage:
2055
                state_dict = None
2056

2057
2058
        config.name_or_path = pretrained_model_name_or_path

2059
        # Instantiate model.
2060
2061
        init_contexts = [no_init_weights(_enable=_fast_init)]

2062
2063
2064
2065
        if is_deepspeed_zero3_enabled():
            import deepspeed

            logger.info("Detected DeepSpeed ZeRO-3: activating zero.init() for this model")
2066
2067
2068
2069
2070
2071
2072
2073
2074
2075
2076
            init_contexts = [deepspeed.zero.Init(config_dict_or_path=deepspeed_config())] + init_contexts
        elif low_cpu_mem_usage:
            init_contexts.append(init_empty_weights())

        with ContextManagers(init_contexts):
            model = cls(config, *model_args, **model_kwargs)

        if device_map == "auto":
            if model._no_split_modules is None:
                raise ValueError(f"{model.__class__.__name__} does not support `device_map='auto'` yet.")
            no_split_modules = model._no_split_modules
2077
2078
            # Make sure tied weights are tied before creating the device map.
            model.tie_weights()
2079
2080
2081
            device_map = infer_auto_device_map(
                model, no_split_module_classes=no_split_modules, dtype=torch_dtype, max_memory=max_memory
            )
2082
2083

        if from_tf:
2084
            if resolved_archive_file.endswith(".index"):
2085
2086
2087
2088
2089
                # Load from a TensorFlow 1.X checkpoint - provided by original authors
                model = cls.load_tf_weights(model, config, resolved_archive_file[:-6])  # Remove the '.index'
            else:
                # Load from our TensorFlow 2.0 checkpoints
                try:
2090
                    from .modeling_tf_pytorch_utils import load_tf2_checkpoint_in_pytorch_model
2091

Yih-Dar's avatar
Yih-Dar committed
2092
2093
2094
                    model, loading_info = load_tf2_checkpoint_in_pytorch_model(
                        model, resolved_archive_file, allow_missing_keys=True, output_loading_info=True
                    )
2095
                except ImportError:
2096
                    logger.error(
Sylvain Gugger's avatar
Sylvain Gugger committed
2097
2098
2099
                        "Loading a TensorFlow model in PyTorch, requires both PyTorch and TensorFlow to be installed."
                        " Please see https://pytorch.org/ and https://www.tensorflow.org/install/ for installation"
                        " instructions."
2100
                    )
2101
                    raise
2102
2103
2104
2105
2106
2107
2108
        elif from_flax:
            try:
                from .modeling_flax_pytorch_utils import load_flax_checkpoint_in_pytorch_model

                model = load_flax_checkpoint_in_pytorch_model(model, resolved_archive_file)
            except ImportError:
                logger.error(
Sylvain Gugger's avatar
Sylvain Gugger committed
2109
2110
2111
                    "Loading a Flax model in PyTorch, requires both PyTorch and Flax to be installed. Please see"
                    " https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for"
                    " installation instructions."
2112
2113
                )
                raise
2114
        elif from_pt:
2115

2116
2117
2118
2119
2120
2121
2122
2123
2124
2125
2126
2127
2128
2129
            # restore default dtype
            if dtype_orig is not None:
                torch.set_default_dtype(dtype_orig)

            model, missing_keys, unexpected_keys, mismatched_keys, error_msgs = cls._load_pretrained_model(
                model,
                state_dict,
                loaded_state_dict_keys,  # XXX: rename?
                resolved_archive_file,
                pretrained_model_name_or_path,
                ignore_mismatched_sizes=ignore_mismatched_sizes,
                sharded_metadata=sharded_metadata,
                _fast_init=_fast_init,
                low_cpu_mem_usage=low_cpu_mem_usage,
2130
2131
2132
2133
                device_map=device_map,
                offload_folder=offload_folder,
                offload_state_dict=offload_state_dict,
                dtype=torch_dtype,
2134
            )
2135

2136
2137
        # make sure token embedding weights are still tied if needed
        model.tie_weights()
2138

2139
        # Set model in evaluation mode to deactivate DropOut modules by default
2140
2141
        model.eval()

2142
2143
2144
2145
        # Dispatch model with hooks on all devices if necessary
        if device_map is not None:
            dispatch_model(model, device_map=device_map, offload_dir=offload_folder)

thomwolf's avatar
thomwolf committed
2146
        if output_loading_info:
Yih-Dar's avatar
Yih-Dar committed
2147
2148
2149
2150
2151
2152
2153
            if loading_info is None:
                loading_info = {
                    "missing_keys": missing_keys,
                    "unexpected_keys": unexpected_keys,
                    "mismatched_keys": mismatched_keys,
                    "error_msgs": error_msgs,
                }
thomwolf's avatar
thomwolf committed
2154
2155
            return model, loading_info

2156
2157
        return model

2158
    @classmethod
Sylvain Gugger's avatar
Sylvain Gugger committed
2159
2160
2161
2162
    def _load_pretrained_model(
        cls,
        model,
        state_dict,
2163
        loaded_keys,
Sylvain Gugger's avatar
Sylvain Gugger committed
2164
2165
2166
2167
2168
        resolved_archive_file,
        pretrained_model_name_or_path,
        ignore_mismatched_sizes=False,
        sharded_metadata=None,
        _fast_init=True,
2169
        low_cpu_mem_usage=False,
2170
2171
        device_map=None,
        offload_folder=None,
2172
        offload_state_dict=None,
2173
        dtype=None,
2174
    ):
Sylvain Gugger's avatar
Sylvain Gugger committed
2175
2176
2177
2178
2179
2180
2181
        if device_map is not None and "disk" in device_map.values():
            if offload_folder is None:
                raise ValueError(
                    "The current `device_map` had weights offloaded to the disk. Please provide an `offload_folder`"
                    " for them."
                )
            os.makedirs(offload_folder, exist_ok=True)
2182
2183
2184
            if offload_state_dict is None:
                offload_state_dict = True

2185
        # Retrieve missing & unexpected_keys
2186
2187
        model_state_dict = model.state_dict()
        expected_keys = list(model_state_dict.keys())
2188
2189
        prefix = model.base_model_prefix

Sylvain Gugger's avatar
Sylvain Gugger committed
2190
2191
2192
2193
2194
2195
2196
        def _fix_key(key):
            if "beta" in key:
                return key.replace("beta", "bias")
            if "gamma" in key:
                return key.replace("gamma", "weight")
            return key

2197
        original_loaded_keys = loaded_keys
Sylvain Gugger's avatar
Sylvain Gugger committed
2198
2199
        loaded_keys = [_fix_key(key) for key in loaded_keys]

2200
2201
2202
2203
2204
2205
        if len(prefix) > 0:
            has_prefix_module = any(s.startswith(prefix) for s in loaded_keys)
            expects_prefix_module = any(s.startswith(prefix) for s in expected_keys)
        else:
            has_prefix_module = False
            expects_prefix_module = False
Patrick von Platen's avatar
Patrick von Platen committed
2206
2207
2208

        # key re-naming operations are never done on the keys
        # that are loaded, but always on the keys of the newly initialized model
2209
2210
        remove_prefix_from_model = not has_prefix_module and expects_prefix_module
        add_prefix_to_model = has_prefix_module and not expects_prefix_module
2211

2212
        if remove_prefix_from_model:
2213
            expected_keys_not_prefixed = [s for s in expected_keys if not s.startswith(prefix)]
2214
            expected_keys = [".".join(s.split(".")[1:]) if s.startswith(prefix) else s for s in expected_keys]
2215
        elif add_prefix_to_model:
2216
2217
2218
2219
2220
2221
2222
2223
2224
2225
2226
2227
2228
2229
2230
            expected_keys = [".".join([prefix, s]) for s in expected_keys]

        missing_keys = list(set(expected_keys) - set(loaded_keys))
        unexpected_keys = list(set(loaded_keys) - set(expected_keys))

        # Some models may have keys that are not in the state by design, removing them before needlessly warning
        # the user.
        if cls._keys_to_ignore_on_load_missing is not None:
            for pat in cls._keys_to_ignore_on_load_missing:
                missing_keys = [k for k in missing_keys if re.search(pat, k) is None]

        if cls._keys_to_ignore_on_load_unexpected is not None:
            for pat in cls._keys_to_ignore_on_load_unexpected:
                unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None]

2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
2241
        # retrieve weights on meta device and put them back on CPU.
        # This is not ideal in terms of memory, but if we don't do that not, we can't initialize them in the next step
        if low_cpu_mem_usage:
            for key in missing_keys:
                if key.startswith(prefix):
                    key = ".".join(key.split(".")[1:])
                param = model_state_dict[key]
                if param.device == torch.device("meta"):
                    set_module_tensor_to_device(model, key, "cpu", torch.empty(*param.size()))

        # retrieve unintialized modules and initialize before maybe overriding that with the pretrained weights.
2242
        if _fast_init:
2243
            uninitialized_modules = model.retrieve_modules_from_names(
2244
                missing_keys, add_prefix=add_prefix_to_model, remove_prefix=remove_prefix_from_model
2245
            )
2246
            for module in uninitialized_modules:
2247
2248
                model._init_weights(module)

2249
2250
2251
        # Make sure we are able to load base models as well as derived models (with heads)
        start_prefix = ""
        model_to_load = model
2252
        if len(cls.base_model_prefix) > 0 and not hasattr(model, cls.base_model_prefix) and has_prefix_module:
2253
            start_prefix = cls.base_model_prefix + "."
2254
        if len(cls.base_model_prefix) > 0 and hasattr(model, cls.base_model_prefix) and not has_prefix_module:
2255
            model_to_load = getattr(model, cls.base_model_prefix)
2256
2257
            if any(key in expected_keys_not_prefixed for key in loaded_keys):
                raise ValueError(
2258
                    "The state dictionary of the model you are trying to load is corrupted. Are you sure it was "
2259
2260
                    "properly saved?"
                )
2261
2262
            if device_map is not None:
                device_map = {k.replace(f"{cls.base_model_prefix}.", ""): v for k, v in device_map.items()}
2263

2264
2265
2266
2267
2268
2269
2270
2271
        def _find_mismatched_keys(
            state_dict,
            model_state_dict,
            loaded_keys,
            add_prefix_to_model,
            remove_prefix_from_model,
            ignore_mismatched_sizes,
        ):
Sylvain Gugger's avatar
Sylvain Gugger committed
2272
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
            mismatched_keys = []
            if ignore_mismatched_sizes:
                for checkpoint_key in loaded_keys:
                    model_key = checkpoint_key
                    if remove_prefix_from_model:
                        # The model key starts with `prefix` but `checkpoint_key` doesn't so we add it.
                        model_key = f"{prefix}.{checkpoint_key}"
                    elif add_prefix_to_model:
                        # The model key doesn't start with `prefix` but `checkpoint_key` does so we remove it.
                        model_key = ".".join(checkpoint_key.split(".")[1:])

                    if (
                        model_key in model_state_dict
                        and state_dict[checkpoint_key].shape != model_state_dict[model_key].shape
                    ):
                        mismatched_keys.append(
                            (checkpoint_key, state_dict[checkpoint_key].shape, model_state_dict[model_key].shape)
                        )
                        del state_dict[checkpoint_key]
2291
2292
2293
2294
2295
2296
2297
            return mismatched_keys

        if state_dict is not None:
            # Whole checkpoint
            mismatched_keys = _find_mismatched_keys(
                state_dict,
                model_state_dict,
2298
                original_loaded_keys,
2299
2300
2301
2302
                add_prefix_to_model,
                remove_prefix_from_model,
                ignore_mismatched_sizes,
            )
Sylvain Gugger's avatar
Sylvain Gugger committed
2303
2304
            error_msgs = _load_state_dict_into_model(model_to_load, state_dict, start_prefix)
        else:
2305
2306
            # Sharded checkpoint or whole but low_cpu_mem_usage==True

Sylvain Gugger's avatar
Sylvain Gugger committed
2307
2308
2309
2310
2311
            # This should always be a list but, just to be sure.
            if not isinstance(resolved_archive_file, list):
                resolved_archive_file = [resolved_archive_file]

            error_msgs = []
2312
            mismatched_keys = []
2313
2314
2315
2316
2317
2318
2319
2320
            offload_index = {} if device_map is not None and "disk" in device_map.values() else None
            if offload_state_dict:
                state_dict_folder = tempfile.mkdtemp()
                state_dict_index = {}
            else:
                state_dict_folder = None
                state_dict_index = None

Sylvain Gugger's avatar
Sylvain Gugger committed
2321
2322
            for shard_file in resolved_archive_file:
                state_dict = load_state_dict(shard_file)
2323

Sylvain Gugger's avatar
Sylvain Gugger committed
2324
2325
                # Mistmatched keys contains tuples key/shape1/shape2 of weights in the checkpoint that have a shape not
                # matching the weights in the model.
2326
2327
2328
                mismatched_keys += _find_mismatched_keys(
                    state_dict,
                    model_state_dict,
2329
                    original_loaded_keys,
2330
2331
2332
2333
                    add_prefix_to_model,
                    remove_prefix_from_model,
                    ignore_mismatched_sizes,
                )
2334
2335

                if low_cpu_mem_usage:
2336
2337
2338
2339
2340
2341
2342
2343
2344
2345
2346
2347
                    new_error_msgs, offload_index, state_dict_index = _load_state_dict_into_meta_model(
                        model_to_load,
                        state_dict,
                        loaded_keys,
                        start_prefix,
                        expected_keys,
                        device_map=device_map,
                        offload_folder=offload_folder,
                        offload_index=offload_index,
                        state_dict_folder=state_dict_folder,
                        state_dict_index=state_dict_index,
                        dtype=dtype,
2348
                    )
2349
                    error_msgs += new_error_msgs
2350
2351
                else:
                    error_msgs += _load_state_dict_into_model(model_to_load, state_dict, start_prefix)
2352

2353
2354
2355
2356
                # force memory release
                del state_dict
                gc.collect()

2357
            if offload_index is not None and len(offload_index) > 0:
Sylvain Gugger's avatar
Sylvain Gugger committed
2358
2359
2360
2361
2362
2363
2364
2365
2366
                if model != model_to_load:
                    # We need to add the prefix of the base model
                    prefix = cls.base_model_prefix
                    for weight_name in offload_index:
                        shutil.move(
                            os.path.join(offload_folder, f"{weight_name}.dat"),
                            os.path.join(offload_folder, f"{prefix}.{weight_name}.dat"),
                        )
                    offload_index = {f"{prefix}.{key}": value for key, value in offload_index.items()}
2367
                save_offload_index(offload_index, offload_folder)
2368
2369
2370

            if offload_state_dict:
                # Load back temporarily offloaded state dict
2371
                load_offloaded_weights(model_to_load, state_dict_index, state_dict_folder)
2372
2373
                shutil.rmtree(state_dict_folder)

2374
2375
        if len(error_msgs) > 0:
            error_msg = "\n\t".join(error_msgs)
2376
2377
2378
2379
            if "size mismatch" in error_msg:
                error_msg += (
                    "\n\tYou may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method."
                )
2380
2381
            raise RuntimeError(f"Error(s) in loading state_dict for {model.__class__.__name__}:\n\t{error_msg}")

2382
2383
        if len(unexpected_keys) > 0:
            logger.warning(
Sylvain Gugger's avatar
Sylvain Gugger committed
2384
2385
2386
2387
2388
2389
2390
                f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when"
                f" initializing {model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are"
                f" initializing {model.__class__.__name__} from the checkpoint of a model trained on another task or"
                " with another architecture (e.g. initializing a BertForSequenceClassification model from a"
                " BertForPreTraining model).\n- This IS NOT expected if you are initializing"
                f" {model.__class__.__name__} from the checkpoint of a model that you expect to be exactly identical"
                " (initializing a BertForSequenceClassification model from a BertForSequenceClassification model)."
2391
2392
2393
2394
2395
            )
        else:
            logger.info(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n")
        if len(missing_keys) > 0:
            logger.warning(
Sylvain Gugger's avatar
Sylvain Gugger committed
2396
2397
2398
                f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at"
                f" {pretrained_model_name_or_path} and are newly initialized: {missing_keys}\nYou should probably"
                " TRAIN this model on a down-stream task to be able to use it for predictions and inference."
2399
            )
2400
        elif len(mismatched_keys) == 0:
2401
            logger.info(
Sylvain Gugger's avatar
Sylvain Gugger committed
2402
2403
2404
2405
                f"All the weights of {model.__class__.__name__} were initialized from the model checkpoint at"
                f" {pretrained_model_name_or_path}.\nIf your task is similar to the task the model of the checkpoint"
                f" was trained on, you can already use {model.__class__.__name__} for predictions without further"
                " training."
2406
            )
2407
2408
2409
2410
2411
2412
2413
2414
        if len(mismatched_keys) > 0:
            mismatched_warning = "\n".join(
                [
                    f"- {key}: found shape {shape1} in the checkpoint and {shape2} in the model instantiated"
                    for key, shape1, shape2 in mismatched_keys
                ]
            )
            logger.warning(
Sylvain Gugger's avatar
Sylvain Gugger committed
2415
2416
2417
2418
                f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at"
                f" {pretrained_model_name_or_path} and are newly initialized because the shapes did not"
                f" match:\n{mismatched_warning}\nYou should probably TRAIN this model on a down-stream task to be able"
                " to use it for predictions and inference."
2419
            )
2420

2421
        return model, missing_keys, unexpected_keys, mismatched_keys, error_msgs
2422
2423
2424
2425

    def retrieve_modules_from_names(self, names, add_prefix=False, remove_prefix=False):
        module_keys = set([".".join(key.split(".")[:-1]) for key in names])

Patrick von Platen's avatar
Patrick von Platen committed
2426
2427
2428
2429
        # torch.nn.ParameterList is a special case where two parameter keywords
        # are appended to the module name, *e.g.* bert.special_embeddings.0
        module_keys = module_keys.union(set([".".join(key.split(".")[:-2]) for key in names if key[-1].isdigit()]))

2430
2431
2432
2433
2434
2435
        retrieved_modules = []
        # retrieve all modules that has at least one missing weight name
        for name, module in self.named_modules():
            if remove_prefix:
                name = ".".join(name.split(".")[1:]) if name.startswith(self.base_model_prefix) else name
            elif add_prefix:
Patrick von Platen's avatar
Patrick von Platen committed
2436
                name = ".".join([self.base_model_prefix, name]) if len(name) > 0 else self.base_model_prefix
2437
2438
2439
2440
2441
2442

            if name in module_keys:
                retrieved_modules.append(module)

        return retrieved_modules

2443
    @staticmethod
2444
    def _load_pretrained_model_low_mem(model, loaded_state_dict_keys, resolved_archive_file, start_prefix=""):
2445
2446
2447
        """
        This is an experimental function that loads the model using ~1.x model size CPU memory

2448
        Before you call it do:
2449

2450
        1. save which state_dict keys are available
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
2461
        2. drop state_dict before model is created, since the latter takes 1x model size memory

        Here then we continue:

        3. switch to the meta device all params/buffers that are going to be replaced from the loaded state_dict
        4. load state_dict 2nd time
        5. replace the params/buffers from the state_dict

        Currently, it doesn't handle missing_keys, unexpected_keys, mismatched_keys. It can't handle deepspeed.
        """

2462
2463
2464
2465
        _move_model_to_meta(model, loaded_state_dict_keys, start_prefix)
        state_dict = load_state_dict(resolved_archive_file)
        error_msgs = _load_state_dict_into_meta_model(model, state_dict, loaded_state_dict_keys, start_prefix)
        return error_msgs
2466

2467
2468
2469
2470
2471
2472
    @classmethod
    def register_for_auto_class(cls, auto_class="AutoModel"):
        """
        Register this class with a given auto class. This should only be used for custom models as the ones in the
        library are already mapped with an auto class.

2473
2474
2475
2476
2477
2478
        <Tip warning={true}>

        This API is experimental and may have some slight breaking changes in the next releases.

        </Tip>

2479
2480
2481
2482
2483
2484
2485
2486
2487
2488
2489
2490
2491
2492
        Args:
            auto_class (`str` or `type`, *optional*, defaults to `"AutoModel"`):
                The auto class to register this new model with.
        """
        if not isinstance(auto_class, str):
            auto_class = auto_class.__name__

        import transformers.models.auto as auto_module

        if not hasattr(auto_module, auto_class):
            raise ValueError(f"{auto_class} is not a valid auto class.")

        cls._auto_class = auto_class

Sylvain Gugger's avatar
Sylvain Gugger committed
2493
2494
2495
2496
2497
2498
2499
2500
2501
2502
2503
2504
2505
2506
    def push_to_hub(
        self,
        repo_path_or_name: Optional[str] = None,
        repo_url: Optional[str] = None,
        use_temp_dir: bool = False,
        commit_message: str = "add model",
        organization: Optional[str] = None,
        private: Optional[bool] = None,
        use_auth_token: Optional[Union[bool, str]] = None,
        max_shard_size: Union[int, str] = "10GB",
        **model_card_kwargs
    ) -> str:
        """
        Upload the model files to the 🤗 Model Hub while synchronizing a local clone of the repo in `repo_path_or_name`.
thomwolf's avatar
thomwolf committed
2507

Sylvain Gugger's avatar
Sylvain Gugger committed
2508
2509
2510
2511
2512
2513
2514
2515
2516
2517
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
2559
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
        Parameters:
            repo_path_or_name (`str`, *optional*):
                Can either be a repository name for your model in the Hub or a path to a local folder (in which case
                the repository will have the name of that local folder). If not specified, will default to the name
                given by `repo_url` and a local directory with that name will be created.
            repo_url (`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 `organization`) with `repo_name`.
            use_temp_dir (`bool`, *optional*, defaults to `False`):
                Whether or not to clone the distant repo in a temporary directory or in `repo_path_or_name` inside the
                current working directory. This will slow things down if you are making changes in an existing repo
                since you will need to clone the repo before every push.
            commit_message (`str`, *optional*, defaults to `"add model"`):
                Message to commit while pushing.
            organization (`str`, *optional*):
                Organization in which you want to push your {object} (you must be a member of this organization).
            private (`bool`, *optional*):
                Whether or not the repository created should be private (requires a paying subscription).
            use_auth_token (`bool` or `str`, *optional*):
                The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
                when running `transformers-cli login` (stored in `~/.huggingface`). Will default to `True` if
                `repo_url` is not specified.
            max_shard_size (`int` or `str`, *optional*, defaults to `"10GB"`):
                The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size
                lower than this size. If expressed as a string, needs to be digits followed by a unit (like `"5MB"`).

                <Tip warning={true}>

                If a single weight of the model is bigger than `max_shard_size`, it will be in its own checkpoint shard
                which will be bigger than `max_shard_size`.

                </Tip>

        Returns:
            `str`: The url of the commit of your {object} in the given repository.

        Examples:

        ```python
        from transformers import AutoModel

        model = AutoModel.from_pretrained("bert-base-cased")

        # Push the model to your namespace with the name "my-finetuned-bert" and have a local clone in the
        # *my-finetuned-bert* folder.
        model.push_to_hub("my-finetuned-bert")

        # Push the model to your namespace with the name "my-finetuned-bert" with no local clone.
        model.push_to_hub("my-finetuned-bert", use_temp_dir=True)

        # Push the model to an organization with the name "my-finetuned-bert" and have a local clone in the
        # *my-finetuned-bert* folder.
        model.push_to_hub("my-finetuned-bert", organization="huggingface")

        # Make a change to an existing repo that has been cloned locally in *my-finetuned-bert*.
        model.push_to_hub("my-finetuned-bert", repo_url="https://huggingface.co/sgugger/my-finetuned-bert")
        ```
        """
        if use_temp_dir:
            # Make sure we use the right `repo_name` for the `repo_url` before replacing it.
            if repo_url is None:
                if use_auth_token is None:
                    use_auth_token = True
                repo_name = Path(repo_path_or_name).name
                repo_url = self._get_repo_url_from_name(
                    repo_name, organization=organization, private=private, use_auth_token=use_auth_token
                )
            repo_path_or_name = tempfile.mkdtemp()

        # Create or clone the repo. If the repo is already cloned, this just retrieves the path to the repo.
        repo = self._create_or_get_repo(
            repo_path_or_name=repo_path_or_name,
            repo_url=repo_url,
            organization=organization,
            private=private,
            use_auth_token=use_auth_token,
        )
        # Save the files in the cloned repo
        self.save_pretrained(repo_path_or_name, max_shard_size=max_shard_size)

        # Commit and push!
        url = self._push_to_hub(repo, commit_message=commit_message)

        # Clean up! Clean up! Everybody everywhere!
        if use_temp_dir:
            shutil.rmtree(repo_path_or_name)

        return url
2596
2597


thomwolf's avatar
thomwolf committed
2598
class PoolerStartLogits(nn.Module):
Sylvain Gugger's avatar
Sylvain Gugger committed
2599
2600
    """
    Compute SQuAD start logits from sequence hidden states.
2601

Sylvain Gugger's avatar
Sylvain Gugger committed
2602
    Args:
2603
2604
        config ([`PretrainedConfig`]):
            The config used by the model, will be used to grab the `hidden_size` of the model.
Sylvain Gugger's avatar
Sylvain Gugger committed
2605
2606
2607
    """

    def __init__(self, config: PretrainedConfig):
Julien Chaumond's avatar
Julien Chaumond committed
2608
        super().__init__()
thomwolf's avatar
thomwolf committed
2609
2610
        self.dense = nn.Linear(config.hidden_size, 1)

Sylvain Gugger's avatar
Sylvain Gugger committed
2611
2612
2613
2614
2615
    def forward(
        self, hidden_states: torch.FloatTensor, p_mask: Optional[torch.FloatTensor] = None
    ) -> torch.FloatTensor:
        """
        Args:
2616
            hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`):
Sylvain Gugger's avatar
Sylvain Gugger committed
2617
                The final hidden states of the model.
2618
            p_mask (`torch.FloatTensor` of shape `(batch_size, seq_len)`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
2619
2620
                Mask for tokens at invalid position, such as query and special symbols (PAD, SEP, CLS). 1.0 means token
                should be masked.
Sylvain Gugger's avatar
Sylvain Gugger committed
2621
2622

        Returns:
2623
            `torch.FloatTensor`: The start logits for SQuAD.
thomwolf's avatar
thomwolf committed
2624
        """
thomwolf's avatar
thomwolf committed
2625
2626
2627
        x = self.dense(hidden_states).squeeze(-1)

        if p_mask is not None:
Lysandre Debut's avatar
Lysandre Debut committed
2628
            if get_parameter_dtype(self) == torch.float16:
2629
2630
2631
                x = x * (1 - p_mask) - 65500 * p_mask
            else:
                x = x * (1 - p_mask) - 1e30 * p_mask
thomwolf's avatar
thomwolf committed
2632
2633
2634
2635
2636
2637

        return x


class PoolerEndLogits(nn.Module):
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
2638
    Compute SQuAD end logits from sequence hidden states.
2639

Sylvain Gugger's avatar
Sylvain Gugger committed
2640
    Args:
2641
        config ([`PretrainedConfig`]):
Sylvain Gugger's avatar
Sylvain Gugger committed
2642
2643
            The config used by the model, will be used to grab the `hidden_size` of the model and the `layer_norm_eps`
            to use.
Sylvain Gugger's avatar
Sylvain Gugger committed
2644
2645
2646
    """

    def __init__(self, config: PretrainedConfig):
Julien Chaumond's avatar
Julien Chaumond committed
2647
        super().__init__()
thomwolf's avatar
thomwolf committed
2648
2649
2650
2651
2652
        self.dense_0 = nn.Linear(config.hidden_size * 2, config.hidden_size)
        self.activation = nn.Tanh()
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dense_1 = nn.Linear(config.hidden_size, 1)

Sylvain Gugger's avatar
Sylvain Gugger committed
2653
2654
2655
2656
2657
2658
2659
2660
2661
    def forward(
        self,
        hidden_states: torch.FloatTensor,
        start_states: Optional[torch.FloatTensor] = None,
        start_positions: Optional[torch.LongTensor] = None,
        p_mask: Optional[torch.FloatTensor] = None,
    ) -> torch.FloatTensor:
        """
        Args:
2662
            hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`):
Sylvain Gugger's avatar
Sylvain Gugger committed
2663
                The final hidden states of the model.
2664
            start_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
2665
                The hidden states of the first tokens for the labeled span.
2666
            start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
2667
                The position of the first token for the labeled span.
2668
            p_mask (`torch.FloatTensor` of shape `(batch_size, seq_len)`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
2669
2670
                Mask for tokens at invalid position, such as query and special symbols (PAD, SEP, CLS). 1.0 means token
                should be masked.
Sylvain Gugger's avatar
Sylvain Gugger committed
2671

2672
        <Tip>
Sylvain Gugger's avatar
Sylvain Gugger committed
2673

Stas Bekman's avatar
Stas Bekman committed
2674
2675
        One of `start_states` or `start_positions` should be not `None`. If both are set, `start_positions` overrides
        `start_states`.
2676
2677

        </Tip>
Sylvain Gugger's avatar
Sylvain Gugger committed
2678
2679

        Returns:
2680
            `torch.FloatTensor`: The end logits for SQuAD.
thomwolf's avatar
thomwolf committed
2681
        """
2682
2683
2684
        assert (
            start_states is not None or start_positions is not None
        ), "One of start_states, start_positions should be not None"
thomwolf's avatar
thomwolf committed
2685
        if start_positions is not None:
2686
            slen, hsz = hidden_states.shape[-2:]
2687
2688
2689
            start_positions = start_positions[:, None, None].expand(-1, -1, hsz)  # shape (bsz, 1, hsz)
            start_states = hidden_states.gather(-2, start_positions)  # shape (bsz, 1, hsz)
            start_states = start_states.expand(-1, slen, -1)  # shape (bsz, slen, hsz)
thomwolf's avatar
thomwolf committed
2690
2691
2692
2693
2694
2695
2696

        x = self.dense_0(torch.cat([hidden_states, start_states], dim=-1))
        x = self.activation(x)
        x = self.LayerNorm(x)
        x = self.dense_1(x).squeeze(-1)

        if p_mask is not None:
Lysandre Debut's avatar
Lysandre Debut committed
2697
            if get_parameter_dtype(self) == torch.float16:
2698
2699
2700
                x = x * (1 - p_mask) - 65500 * p_mask
            else:
                x = x * (1 - p_mask) - 1e30 * p_mask
thomwolf's avatar
thomwolf committed
2701
2702
2703
2704
2705

        return x


class PoolerAnswerClass(nn.Module):
Sylvain Gugger's avatar
Sylvain Gugger committed
2706
2707
2708
2709
    """
    Compute SQuAD 2.0 answer class from classification and start tokens hidden states.

    Args:
2710
2711
        config ([`PretrainedConfig`]):
            The config used by the model, will be used to grab the `hidden_size` of the model.
Sylvain Gugger's avatar
Sylvain Gugger committed
2712
    """
2713

thomwolf's avatar
thomwolf committed
2714
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
2715
        super().__init__()
thomwolf's avatar
thomwolf committed
2716
2717
2718
2719
        self.dense_0 = nn.Linear(config.hidden_size * 2, config.hidden_size)
        self.activation = nn.Tanh()
        self.dense_1 = nn.Linear(config.hidden_size, 1, bias=False)

Sylvain Gugger's avatar
Sylvain Gugger committed
2720
2721
2722
2723
2724
2725
2726
    def forward(
        self,
        hidden_states: torch.FloatTensor,
        start_states: Optional[torch.FloatTensor] = None,
        start_positions: Optional[torch.LongTensor] = None,
        cls_index: Optional[torch.LongTensor] = None,
    ) -> torch.FloatTensor:
2727
2728
        """
        Args:
2729
            hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`):
Sylvain Gugger's avatar
Sylvain Gugger committed
2730
                The final hidden states of the model.
2731
            start_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
2732
                The hidden states of the first tokens for the labeled span.
2733
            start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
2734
                The position of the first token for the labeled span.
2735
2736
2737
2738
            cls_index (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
                Position of the CLS token for each sentence in the batch. If `None`, takes the last token.

        <Tip>
Sylvain Gugger's avatar
Sylvain Gugger committed
2739

Stas Bekman's avatar
Stas Bekman committed
2740
2741
        One of `start_states` or `start_positions` should be not `None`. If both are set, `start_positions` overrides
        `start_states`.
Sylvain Gugger's avatar
Sylvain Gugger committed
2742

2743
        </Tip>
Sylvain Gugger's avatar
Sylvain Gugger committed
2744
2745

        Returns:
2746
            `torch.FloatTensor`: The SQuAD 2.0 answer class.
thomwolf's avatar
thomwolf committed
2747
        """
Sylvain Gugger's avatar
Sylvain Gugger committed
2748
        # No dependency on end_feature so that we can obtain one single `cls_logits` for each sample.
2749
        hsz = hidden_states.shape[-1]
2750
2751
2752
        assert (
            start_states is not None or start_positions is not None
        ), "One of start_states, start_positions should be not None"
thomwolf's avatar
thomwolf committed
2753
        if start_positions is not None:
2754
2755
            start_positions = start_positions[:, None, None].expand(-1, -1, hsz)  # shape (bsz, 1, hsz)
            start_states = hidden_states.gather(-2, start_positions).squeeze(-2)  # shape (bsz, hsz)
thomwolf's avatar
thomwolf committed
2756
2757

        if cls_index is not None:
2758
2759
            cls_index = cls_index[:, None, None].expand(-1, -1, hsz)  # shape (bsz, 1, hsz)
            cls_token_state = hidden_states.gather(-2, cls_index).squeeze(-2)  # shape (bsz, hsz)
thomwolf's avatar
thomwolf committed
2760
        else:
2761
            cls_token_state = hidden_states[:, -1, :]  # shape (bsz, hsz)
thomwolf's avatar
thomwolf committed
2762
2763
2764
2765
2766
2767
2768
2769

        x = self.dense_0(torch.cat([start_states, cls_token_state], dim=-1))
        x = self.activation(x)
        x = self.dense_1(x).squeeze(-1)

        return x


2770
2771
2772
@dataclass
class SquadHeadOutput(ModelOutput):
    """
2773
    Base class for outputs of question answering models using a [`~modeling_utils.SQuADHead`].
2774
2775

    Args:
2776
        loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned if both `start_positions` and `end_positions` are provided):
Sylvain Gugger's avatar
Sylvain Gugger committed
2777
2778
            Classification loss as the sum of start token, end token (and is_impossible if provided) classification
            losses.
2779
        start_top_log_probs (`torch.FloatTensor` of shape `(batch_size, config.start_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided):
2780
            Log probabilities for the top config.start_n_top start token possibilities (beam-search).
2781
        start_top_index (`torch.LongTensor` of shape `(batch_size, config.start_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided):
2782
            Indices for the top config.start_n_top start token possibilities (beam-search).
2783
2784
        end_top_log_probs (`torch.FloatTensor` of shape `(batch_size, config.start_n_top * config.end_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided):
            Log probabilities for the top `config.start_n_top * config.end_n_top` end token possibilities
Sylvain Gugger's avatar
Sylvain Gugger committed
2785
            (beam-search).
2786
2787
2788
2789
        end_top_index (`torch.LongTensor` of shape `(batch_size, config.start_n_top * config.end_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided):
            Indices for the top `config.start_n_top * config.end_n_top` end token possibilities (beam-search).
        cls_logits (`torch.FloatTensor` of shape `(batch_size,)`, *optional*, returned if `start_positions` or `end_positions` is not provided):
            Log probabilities for the `is_impossible` label of the answers.
2790
2791
2792
2793
2794
2795
2796
2797
2798
2799
2800

    """

    loss: Optional[torch.FloatTensor] = None
    start_top_log_probs: Optional[torch.FloatTensor] = None
    start_top_index: Optional[torch.LongTensor] = None
    end_top_log_probs: Optional[torch.FloatTensor] = None
    end_top_index: Optional[torch.LongTensor] = None
    cls_logits: Optional[torch.FloatTensor] = None


thomwolf's avatar
thomwolf committed
2801
class SQuADHead(nn.Module):
Sylvain Gugger's avatar
Sylvain Gugger committed
2802
2803
    r"""
    A SQuAD head inspired by XLNet.
2804

Sylvain Gugger's avatar
Sylvain Gugger committed
2805
    Args:
2806
        config ([`PretrainedConfig`]):
Sylvain Gugger's avatar
Sylvain Gugger committed
2807
2808
            The config used by the model, will be used to grab the `hidden_size` of the model and the `layer_norm_eps`
            to use.
thomwolf's avatar
thomwolf committed
2809
    """
2810

thomwolf's avatar
thomwolf committed
2811
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
2812
        super().__init__()
thomwolf's avatar
thomwolf committed
2813
2814
2815
2816
2817
2818
2819
        self.start_n_top = config.start_n_top
        self.end_n_top = config.end_n_top

        self.start_logits = PoolerStartLogits(config)
        self.end_logits = PoolerEndLogits(config)
        self.answer_class = PoolerAnswerClass(config)

Sylvain Gugger's avatar
Sylvain Gugger committed
2820
    @replace_return_docstrings(output_type=SquadHeadOutput, config_class=PretrainedConfig)
2821
    def forward(
2822
        self,
Sylvain Gugger's avatar
Sylvain Gugger committed
2823
2824
2825
2826
2827
2828
        hidden_states: torch.FloatTensor,
        start_positions: Optional[torch.LongTensor] = None,
        end_positions: Optional[torch.LongTensor] = None,
        cls_index: Optional[torch.LongTensor] = None,
        is_impossible: Optional[torch.LongTensor] = None,
        p_mask: Optional[torch.FloatTensor] = None,
2829
        return_dict: bool = False,
Sylvain Gugger's avatar
Sylvain Gugger committed
2830
2831
    ) -> Union[SquadHeadOutput, Tuple[torch.FloatTensor]]:
        """
Lysandre's avatar
Lysandre committed
2832
        Args:
2833
            hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`):
Lysandre's avatar
Lysandre committed
2834
                Final hidden states of the model on the sequence tokens.
2835
            start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
Lysandre's avatar
Lysandre committed
2836
                Positions of the first token for the labeled span.
2837
            end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
Lysandre's avatar
Lysandre committed
2838
                Positions of the last token for the labeled span.
2839
2840
2841
            cls_index (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
                Position of the CLS token for each sentence in the batch. If `None`, takes the last token.
            is_impossible (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
Lysandre's avatar
Lysandre committed
2842
                Whether the question has a possible answer in the paragraph or not.
2843
            p_mask (`torch.FloatTensor` of shape `(batch_size, seq_len)`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
2844
2845
                Mask for tokens at invalid position, such as query and special symbols (PAD, SEP, CLS). 1.0 means token
                should be masked.
2846
            return_dict (`bool`, *optional*, defaults to `False`):
2847
                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
Sylvain Gugger's avatar
Sylvain Gugger committed
2848

Lysandre's avatar
Lysandre committed
2849
        Returns:
Sylvain Gugger's avatar
Sylvain Gugger committed
2850
        """
thomwolf's avatar
thomwolf committed
2851
        start_logits = self.start_logits(hidden_states, p_mask=p_mask)
thomwolf's avatar
thomwolf committed
2852
2853
2854
2855
2856
2857
2858
2859
2860
2861
2862
2863
2864
2865
2866
2867
2868
2869
2870
2871
2872
2873
2874

        if start_positions is not None and end_positions is not None:
            # If we are on multi-GPU, let's remove the dimension added by batch splitting
            for x in (start_positions, end_positions, cls_index, is_impossible):
                if x is not None and x.dim() > 1:
                    x.squeeze_(-1)

            # during training, compute the end logits based on the ground truth of the start position
            end_logits = self.end_logits(hidden_states, start_positions=start_positions, p_mask=p_mask)

            loss_fct = CrossEntropyLoss()
            start_loss = loss_fct(start_logits, start_positions)
            end_loss = loss_fct(end_logits, end_positions)
            total_loss = (start_loss + end_loss) / 2

            if cls_index is not None and is_impossible is not None:
                # Predict answerability from the representation of CLS and START
                cls_logits = self.answer_class(hidden_states, start_positions=start_positions, cls_index=cls_index)
                loss_fct_cls = nn.BCEWithLogitsLoss()
                cls_loss = loss_fct_cls(cls_logits, is_impossible)

                # note(zhiliny): by default multiply the loss by 0.5 so that the scale is comparable to start_loss and end_loss
                total_loss += cls_loss * 0.5
2875

2876
            return SquadHeadOutput(loss=total_loss) if return_dict else (total_loss,)
thomwolf's avatar
thomwolf committed
2877
2878
2879
2880

        else:
            # during inference, compute the end logits based on beam search
            bsz, slen, hsz = hidden_states.size()
2881
            start_log_probs = nn.functional.softmax(start_logits, dim=-1)  # shape (bsz, slen)
2882
2883
2884
2885
2886
2887
2888
2889
2890
2891
2892

            start_top_log_probs, start_top_index = torch.topk(
                start_log_probs, self.start_n_top, dim=-1
            )  # shape (bsz, start_n_top)
            start_top_index_exp = start_top_index.unsqueeze(-1).expand(-1, -1, hsz)  # shape (bsz, start_n_top, hsz)
            start_states = torch.gather(hidden_states, -2, start_top_index_exp)  # shape (bsz, start_n_top, hsz)
            start_states = start_states.unsqueeze(1).expand(-1, slen, -1, -1)  # shape (bsz, slen, start_n_top, hsz)

            hidden_states_expanded = hidden_states.unsqueeze(2).expand_as(
                start_states
            )  # shape (bsz, slen, start_n_top, hsz)
thomwolf's avatar
thomwolf committed
2893
2894
            p_mask = p_mask.unsqueeze(-1) if p_mask is not None else None
            end_logits = self.end_logits(hidden_states_expanded, start_states=start_states, p_mask=p_mask)
2895
            end_log_probs = nn.functional.softmax(end_logits, dim=1)  # shape (bsz, slen, start_n_top)
thomwolf's avatar
thomwolf committed
2896

2897
2898
2899
            end_top_log_probs, end_top_index = torch.topk(
                end_log_probs, self.end_n_top, dim=1
            )  # shape (bsz, end_n_top, start_n_top)
thomwolf's avatar
thomwolf committed
2900
2901
2902
2903
2904
2905
            end_top_log_probs = end_top_log_probs.view(-1, self.start_n_top * self.end_n_top)
            end_top_index = end_top_index.view(-1, self.start_n_top * self.end_n_top)

            start_states = torch.einsum("blh,bl->bh", hidden_states, start_log_probs)
            cls_logits = self.answer_class(hidden_states, start_states=start_states, cls_index=cls_index)

2906
            if not return_dict:
2907
2908
2909
2910
2911
2912
2913
2914
2915
                return (start_top_log_probs, start_top_index, end_top_log_probs, end_top_index, cls_logits)
            else:
                return SquadHeadOutput(
                    start_top_log_probs=start_top_log_probs,
                    start_top_index=start_top_index,
                    end_top_log_probs=end_top_log_probs,
                    end_top_index=end_top_index,
                    cls_logits=cls_logits,
                )
thomwolf's avatar
thomwolf committed
2916
2917
2918


class SequenceSummary(nn.Module):
Sylvain Gugger's avatar
Sylvain Gugger committed
2919
2920
2921
2922
    r"""
    Compute a single vector summary of a sequence hidden states.

    Args:
2923
        config ([`PretrainedConfig`]):
Sylvain Gugger's avatar
Sylvain Gugger committed
2924
2925
            The config used by the model. Relevant arguments in the config class of the model are (refer to the actual
            config class of your model for the default values it uses):
Sylvain Gugger's avatar
Sylvain Gugger committed
2926

2927
            - **summary_type** (`str`) -- The method to use to make this summary. Accepted values are:
Sylvain Gugger's avatar
Sylvain Gugger committed
2928

2929
2930
2931
2932
2933
                - `"last"` -- Take the last token hidden state (like XLNet)
                - `"first"` -- Take the first token hidden state (like Bert)
                - `"mean"` -- Take the mean of all tokens hidden states
                - `"cls_index"` -- Supply a Tensor of classification token position (GPT/GPT-2)
                - `"attn"` -- Not implemented now, use multi-head attention
Sylvain Gugger's avatar
Sylvain Gugger committed
2934

2935
            - **summary_use_proj** (`bool`) -- Add a projection after the vector extraction.
Sylvain Gugger's avatar
Sylvain Gugger committed
2936
2937
2938
2939
2940
2941
            - **summary_proj_to_labels** (`bool`) -- If `True`, the projection outputs to `config.num_labels` classes
              (otherwise to `config.hidden_size`).
            - **summary_activation** (`Optional[str]`) -- Set to `"tanh"` to add a tanh activation to the output,
              another string or `None` will add no activation.
            - **summary_first_dropout** (`float`) -- Optional dropout probability before the projection and activation.
            - **summary_last_dropout** (`float`)-- Optional dropout probability after the projection and activation.
thomwolf's avatar
thomwolf committed
2942
    """
2943

2944
    def __init__(self, config: PretrainedConfig):
Julien Chaumond's avatar
Julien Chaumond committed
2945
        super().__init__()
thomwolf's avatar
thomwolf committed
2946

2947
        self.summary_type = getattr(config, "summary_type", "last")
2948
        if self.summary_type == "attn":
thomwolf's avatar
thomwolf committed
2949
2950
2951
2952
2953
            # We should use a standard multi-head attention module with absolute positional embedding for that.
            # Cf. https://github.com/zihangdai/xlnet/blob/master/modeling.py#L253-L276
            # We can probably just use the multi-head attention module of PyTorch >=1.1.0
            raise NotImplementedError

thomwolf's avatar
thomwolf committed
2954
        self.summary = Identity()
2955
2956
        if hasattr(config, "summary_use_proj") and config.summary_use_proj:
            if hasattr(config, "summary_proj_to_labels") and config.summary_proj_to_labels and config.num_labels > 0:
2957
                num_classes = config.num_labels
thomwolf's avatar
thomwolf committed
2958
2959
2960
2961
            else:
                num_classes = config.hidden_size
            self.summary = nn.Linear(config.hidden_size, num_classes)

2962
        activation_string = getattr(config, "summary_activation", None)
Lysandre's avatar
Lysandre committed
2963
        self.activation: Callable = get_activation(activation_string) if activation_string else Identity()
thomwolf's avatar
thomwolf committed
2964

thomwolf's avatar
thomwolf committed
2965
        self.first_dropout = Identity()
2966
        if hasattr(config, "summary_first_dropout") and config.summary_first_dropout > 0:
2967
2968
            self.first_dropout = nn.Dropout(config.summary_first_dropout)

thomwolf's avatar
thomwolf committed
2969
        self.last_dropout = Identity()
2970
        if hasattr(config, "summary_last_dropout") and config.summary_last_dropout > 0:
2971
            self.last_dropout = nn.Dropout(config.summary_last_dropout)
thomwolf's avatar
thomwolf committed
2972

Sylvain Gugger's avatar
Sylvain Gugger committed
2973
2974
2975
2976
2977
2978
2979
    def forward(
        self, hidden_states: torch.FloatTensor, cls_index: Optional[torch.LongTensor] = None
    ) -> torch.FloatTensor:
        """
        Compute a single vector summary of a sequence hidden states.

        Args:
2980
            hidden_states (`torch.FloatTensor` of shape `[batch_size, seq_len, hidden_size]`):
Sylvain Gugger's avatar
Sylvain Gugger committed
2981
                The hidden states of the last layer.
2982
            cls_index (`torch.LongTensor` of shape `[batch_size]` or `[batch_size, ...]` where ... are optional leading dimensions of `hidden_states`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
2983
                Used if `summary_type == "cls_index"` and takes the last token of the sequence as classification token.
Sylvain Gugger's avatar
Sylvain Gugger committed
2984
2985

        Returns:
2986
            `torch.FloatTensor`: The summary of the sequence hidden states.
thomwolf's avatar
thomwolf committed
2987
        """
2988
        if self.summary_type == "last":
thomwolf's avatar
thomwolf committed
2989
            output = hidden_states[:, -1]
2990
        elif self.summary_type == "first":
thomwolf's avatar
thomwolf committed
2991
            output = hidden_states[:, 0]
2992
        elif self.summary_type == "mean":
thomwolf's avatar
thomwolf committed
2993
            output = hidden_states.mean(dim=1)
2994
        elif self.summary_type == "cls_index":
thomwolf's avatar
thomwolf committed
2995
            if cls_index is None:
Lysandre's avatar
Lysandre committed
2996
2997
2998
2999
3000
                cls_index = torch.full_like(
                    hidden_states[..., :1, :],
                    hidden_states.shape[-2] - 1,
                    dtype=torch.long,
                )
thomwolf's avatar
thomwolf committed
3001
            else:
thomwolf's avatar
thomwolf committed
3002
                cls_index = cls_index.unsqueeze(-1).unsqueeze(-1)
3003
                cls_index = cls_index.expand((-1,) * (cls_index.dim() - 1) + (hidden_states.size(-1),))
thomwolf's avatar
thomwolf committed
3004
            # shape of cls_index: (bsz, XX, 1, hidden_size) where XX are optional leading dim of hidden_states
3005
3006
            output = hidden_states.gather(-2, cls_index).squeeze(-2)  # shape (bsz, XX, hidden_size)
        elif self.summary_type == "attn":
thomwolf's avatar
thomwolf committed
3007
3008
            raise NotImplementedError

3009
        output = self.first_dropout(output)
thomwolf's avatar
thomwolf committed
3010
3011
        output = self.summary(output)
        output = self.activation(output)
3012
        output = self.last_dropout(output)
thomwolf's avatar
thomwolf committed
3013
3014
3015
3016

        return output


3017
def unwrap_model(model: nn.Module) -> nn.Module:
3018
3019
3020
3021
    """
    Recursively unwraps a model from potential containers (as used in distributed training).

    Args:
3022
        model (`torch.nn.Module`): The model to unwrap.
3023
3024
3025
3026
3027
3028
    """
    # since there could be multiple levels of wrapping, unwrap recursively
    if hasattr(model, "module"):
        return unwrap_model(model.module)
    else:
        return model