bert.py 26.5 KB
Newer Older
Tri Dao's avatar
Tri Dao committed
1
2
3
4
5
6
7
8
# Copyright (c) 2022, Tri Dao.
# This BERT implementation is based on our MLPerf 2.0 and MLPerf 2.1 BERT implementation.
# https://github.com/mlcommons/training_results_v2.0/blob/main/HazyResearch/benchmarks/bert/implementations/pytorch/modeling.py
# https://github.com/mlcommons/training_results_v2.1/blob/main/Azure-HazyResearch/benchmarks/bert/implementations/ND96amsr_A100_v4/modeling.py

# Inspired by https://github.com/huggingface/transformers/blob/main/src/transformers/models/bert/modeling_bert.py

import logging
Tri Dao's avatar
Tri Dao committed
9
import re
Tri Dao's avatar
Tri Dao committed
10
from collections import OrderedDict
Tri Dao's avatar
Tri Dao committed
11
12
from collections.abc import Sequence
from functools import partial
Tri Dao's avatar
Tri Dao committed
13
14
15
16
17

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
Tri Dao's avatar
Tri Dao committed
18
19
20
21
22
23
24
25
26
27
28
29
from transformers import BertConfig
from transformers.models.bert.modeling_bert import (
    BaseModelOutputWithPoolingAndCrossAttentions,
    BertForPreTrainingOutput,
)

from flash_attn.bert_padding import (
    index_first_axis,
    index_first_axis_residual,
    pad_input,
    unpad_input,
)
Tri Dao's avatar
Tri Dao committed
30
31
from flash_attn.modules.block import Block
from flash_attn.modules.embedding import BertEmbeddings
Tri Dao's avatar
Tri Dao committed
32
33
from flash_attn.modules.mha import MHA
from flash_attn.modules.mlp import FusedMLP, Mlp
34
from flash_attn.utils.pretrained import state_dict_from_pretrained
Tri Dao's avatar
Tri Dao committed
35
36

try:
Tri Dao's avatar
Tri Dao committed
37
    from flash_attn.ops.fused_dense import FusedDense
Tri Dao's avatar
Tri Dao committed
38
except ImportError:
Tri Dao's avatar
Tri Dao committed
39
    FusedDense = None
Tri Dao's avatar
Tri Dao committed
40
41
42
43
44
45
46

try:
    from flash_attn.ops.layer_norm import dropout_add_layer_norm, layer_norm
except ImportError:
    dropout_add_layer_norm, layer_norm = None, None

try:
47
    from flash_attn.losses.cross_entropy import CrossEntropyLoss
Tri Dao's avatar
Tri Dao committed
48
except ImportError:
49
    CrossEntropyLoss = None
Tri Dao's avatar
Tri Dao committed
50
51
52
53
54


logger = logging.getLogger(__name__)


55
def create_mixer_cls(config, cross_attn=False, return_residual=False):
Tri Dao's avatar
Tri Dao committed
56
57
    use_flash_attn = getattr(config, "use_flash_attn", False)
    fused_bias_fc = getattr(config, "fused_bias_fc", False)
58
59
60
61
62
63
    rotary_kwargs = {}
    if config.position_embedding_type == "rotary":
        rotary_kwargs["rotary_emb_dim"] = getattr(config, "rotary_emb_dim", config.hidden_size)
        rotary_kwargs["rotary_emb_base"] = getattr(config, "rotary_emb_base", 10000.0)
        rotary_kwargs["rotary_emb_scale_base"] = getattr(config, "rotary_emb_scale_base", None)
        rotary_kwargs["rotary_emb_interleaved"] = getattr(config, "rotary_emb_interleaved", False)
Tri Dao's avatar
Tri Dao committed
64
65
66
67
68
69
70
71
72
73
74
    mixer_cls = partial(
        MHA,
        num_heads=config.num_attention_heads,
        cross_attn=cross_attn,
        dropout=config.attention_probs_dropout_prob,
        causal=False,
        fused_bias_fc=fused_bias_fc,
        use_flash_attn=use_flash_attn,
        return_residual=return_residual,
        **rotary_kwargs,
    )
Tri Dao's avatar
Tri Dao committed
75
76
77
    return mixer_cls


78
def create_mlp_cls(config, layer_idx=None, return_residual=False):
Tri Dao's avatar
Tri Dao committed
79
    inner_dim = config.intermediate_size
Tri Dao's avatar
Tri Dao committed
80
    fused_mlp = getattr(config, "fused_mlp", False)
81
    if fused_mlp:
Tri Dao's avatar
Tri Dao committed
82
83
84
        assert config.hidden_act in ["gelu_new", "gelu_fast"], (
            "fused_mlp only " "supports approximate gelu"
        )
85
    if not fused_mlp:
Tri Dao's avatar
Tri Dao committed
86
87
88
89
90
91
92
        approximate = "tanh" if config.hidden_act in ["gelu_new", "gelu_fast"] else "none"
        mlp_cls = partial(
            Mlp,
            hidden_features=inner_dim,
            activation=partial(F.gelu, approximate=approximate),
            return_residual=return_residual,
        )
Tri Dao's avatar
Tri Dao committed
93
    else:
94
        if FusedMLP is None:
Tri Dao's avatar
Tri Dao committed
95
96
            raise ImportError("fused_dense is not installed")
        mlp_checkpoint_lvl = getattr(config, "mlp_checkpoint_lvl", 0)
Tri Dao's avatar
Tri Dao committed
97
98
99
100
        # mlp_checkpoint_lvl could be a list, which contains the checkpoint_lvl for each layer
        if isinstance(mlp_checkpoint_lvl, Sequence):
            assert layer_idx is not None
            mlp_checkpoint_lvl = mlp_checkpoint_lvl[layer_idx]
Tri Dao's avatar
Tri Dao committed
101
102
103
104
105
106
        mlp_cls = partial(
            FusedMLP,
            hidden_features=inner_dim,
            checkpoint_lvl=mlp_checkpoint_lvl,
            return_residual=return_residual,
        )
Tri Dao's avatar
Tri Dao committed
107
108
109
110
    return mlp_cls


def create_block(config, layer_idx=None):
Tri Dao's avatar
Tri Dao committed
111
112
    last_layer_subset = getattr(config, "last_layer_subset", False)
    cross_attn = last_layer_subset and layer_idx == config.num_hidden_layers - 1
113
114
115
116
117
118
    # TD [2022-12-19]: For cross attention (last layer), we actually want to return the
    # residual x_kv, not residual x. But it's annoying to change the API (and it only affects
    # one layer) so we just choose not to return residual in this case.
    return_residual = not cross_attn
    mixer_cls = create_mixer_cls(config, cross_attn, return_residual=return_residual)
    mlp_cls = create_mlp_cls(config, layer_idx, return_residual=return_residual)
Tri Dao's avatar
Tri Dao committed
119
    norm_cls = partial(nn.LayerNorm, eps=config.layer_norm_eps)
Tri Dao's avatar
Tri Dao committed
120
121
122
123
124
125
126
127
128
129
130
    block = Block(
        config.hidden_size,
        mixer_cls,
        mlp_cls,
        norm_cls=norm_cls,
        prenorm=False,
        resid_dropout1=config.hidden_dropout_prob,
        resid_dropout2=config.hidden_dropout_prob,
        fused_dropout_add_ln=getattr(config, "fused_dropout_add_ln", False),
        return_residual=return_residual,
    )
Tri Dao's avatar
Tri Dao committed
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
    return block


# https://github.com/huggingface/transformers/blob/7032e0203262ebb2ebf55da8d2e01f873973e835/src/transformers/models/bert/modeling_bert.py#L748
def _init_weights(module, initializer_range=0.02):
    if isinstance(module, nn.Linear):
        nn.init.normal_(module.weight, std=initializer_range)
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Embedding):
        nn.init.normal_(module.weight, std=initializer_range)
        if module.padding_idx is not None:
            nn.init.zeros_(module.weight[module.padding_idx])


class BertEncoder(nn.Module):
    def __init__(self, config: BertConfig):
        super().__init__()
Tri Dao's avatar
Tri Dao committed
149
150
151
152
        self.use_flash_attn = getattr(config, "use_flash_attn", False)
        self.layers = nn.ModuleList(
            [create_block(config, layer_idx=i) for i in range(config.num_hidden_layers)]
        )
Tri Dao's avatar
Tri Dao committed
153

154
155
156
157
158
    def forward(self, hidden_states, key_padding_mask=None, subset_mask=None):
        """If subset_mask is not None, we only want output for the subset of the sequence.
        This means that we only compute the last layer output for these tokens.
        subset_mask: (batch, seqlen), dtype=torch.bool
        """
Tri Dao's avatar
Tri Dao committed
159
        if key_padding_mask is None or not self.use_flash_attn:
Tri Dao's avatar
Tri Dao committed
160
161
162
            mixer_kwargs = (
                {"key_padding_mask": key_padding_mask} if key_padding_mask is not None else None
            )
Tri Dao's avatar
Tri Dao committed
163
164
            for layer in self.layers:
                hidden_states = layer(hidden_states, mixer_kwargs=mixer_kwargs)
165
166
            if subset_mask is not None:
                hidden_states = hidden_states[subset_mask]
Tri Dao's avatar
Tri Dao committed
167
168
169
170
171
        else:
            batch, seqlen = hidden_states.shape[:2]
            hidden_states, indices, cu_seqlens, max_seqlen_in_batch = unpad_input(
                hidden_states, key_padding_mask
            )
Tri Dao's avatar
Tri Dao committed
172
            mixer_kwargs = {"cu_seqlens": cu_seqlens, "max_seqlen": max_seqlen_in_batch}
173
174
175
176
177
178
179
180
            if subset_mask is None:
                for layer in self.layers:
                    hidden_states = layer(hidden_states, mixer_kwargs=mixer_kwargs)
                hidden_states = pad_input(hidden_states, indices, batch, seqlen)
            else:
                for layer in self.layers[:-1]:
                    hidden_states = layer(hidden_states, mixer_kwargs=mixer_kwargs)
                if key_padding_mask is not None:
Tri Dao's avatar
Tri Dao committed
181
182
183
                    subset_idx = torch.nonzero(
                        subset_mask[key_padding_mask], as_tuple=False
                    ).flatten()
184
                    subset_seqlens = (subset_mask & key_padding_mask).sum(dim=-1, dtype=torch.int32)
Tri Dao's avatar
Tri Dao committed
185
186
187
                    subset_cu_seqlens = F.pad(
                        torch.cumsum(subset_seqlens, dim=0, dtype=torch.torch.int32), (1, 0)
                    )
188
189
190
                else:
                    subset_idx = torch.nonzero(subset_mask, as_tuple=False).flatten()
                    subset_seqlens = subset_mask.sum(dim=-1, dtype=torch.int32)
Tri Dao's avatar
Tri Dao committed
191
192
193
                    subset_cu_seqlens = F.pad(
                        torch.cumsum(subset_seqlens, dim=0, dtype=torch.torch.int32), (1, 0)
                    )
194
195
196
197
                hidden_states_subset, hidden_states = index_first_axis_residual(
                    hidden_states, subset_idx
                )
                # It's ok to set max_seqlen_q to be much larger
Tri Dao's avatar
Tri Dao committed
198
199
200
201
202
203
204
                mixer_kwargs = {
                    "x_kv": hidden_states,
                    "cu_seqlens": subset_cu_seqlens,
                    "max_seqlen": max_seqlen_in_batch,
                    "cu_seqlens_k": cu_seqlens,
                    "max_seqlen_k": max_seqlen_in_batch,
                }
205
                hidden_states = self.layers[-1](hidden_states_subset, mixer_kwargs=mixer_kwargs)
Tri Dao's avatar
Tri Dao committed
206
207
208
209
210
211
        return hidden_states


class BertPooler(nn.Module):
    def __init__(self, config):
        super().__init__()
Tri Dao's avatar
Tri Dao committed
212
        fused_bias_fc = getattr(config, "fused_bias_fc", False)
Tri Dao's avatar
Tri Dao committed
213
        if fused_bias_fc and FusedDense is None:
Tri Dao's avatar
Tri Dao committed
214
            raise ImportError("fused_dense is not installed")
Tri Dao's avatar
Tri Dao committed
215
        linear_cls = nn.Linear if not fused_bias_fc else FusedDense
Tri Dao's avatar
Tri Dao committed
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
        self.dense = linear_cls(config.hidden_size, config.hidden_size)
        self.activation = nn.Tanh()

    def forward(self, hidden_states, pool=True):
        # We "pool" the model by simply taking the hidden state corresponding
        # to the first token.
        first_token_tensor = hidden_states[:, 0] if pool else hidden_states
        pooled_output = self.dense(first_token_tensor)
        pooled_output = self.activation(pooled_output)
        return pooled_output


class BertPredictionHeadTransform(nn.Module):
    def __init__(self, config):
        super().__init__()
Tri Dao's avatar
Tri Dao committed
231
        fused_bias_fc = getattr(config, "fused_bias_fc", False)
Tri Dao's avatar
Tri Dao committed
232
        if fused_bias_fc and FusedDense is None:
Tri Dao's avatar
Tri Dao committed
233
234
            raise ImportError("fused_dense is not installed")
        self.fused_dropout_add_ln = getattr(config, "fused_dropout_add_ln", False)
Tri Dao's avatar
Tri Dao committed
235
        if self.fused_dropout_add_ln and layer_norm is None:
Tri Dao's avatar
Tri Dao committed
236
            raise ImportError("dropout_add_layer_norm is not installed")
Tri Dao's avatar
Tri Dao committed
237
        linear_cls = nn.Linear if not fused_bias_fc else FusedDense
Tri Dao's avatar
Tri Dao committed
238
        self.dense = linear_cls(config.hidden_size, config.hidden_size)
Tri Dao's avatar
Tri Dao committed
239
        approximate = "tanh" if config.hidden_act in ["gelu_new", "gelu_fast"] else "none"
240
        self.transform_act_fn = nn.GELU(approximate=approximate)
Tri Dao's avatar
Tri Dao committed
241
242
243
244
245
246
247
248
        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.dense(hidden_states)
        hidden_states = self.transform_act_fn(hidden_states)
        if not self.fused_dropout_add_ln:
            hidden_states = self.layer_norm(hidden_states)
        else:
Tri Dao's avatar
Tri Dao committed
249
250
251
            hidden_states = layer_norm(
                hidden_states, self.layer_norm.weight, self.layer_norm.bias, self.layer_norm.eps
            )
Tri Dao's avatar
Tri Dao committed
252
253
254
255
256
257
        return hidden_states


class BertLMPredictionHead(nn.Module):
    def __init__(self, config):
        super().__init__()
Tri Dao's avatar
Tri Dao committed
258
        fused_bias_fc = getattr(config, "fused_bias_fc", False)
Tri Dao's avatar
Tri Dao committed
259
        if fused_bias_fc and FusedDense is None:
Tri Dao's avatar
Tri Dao committed
260
            raise ImportError("fused_dense is not installed")
Tri Dao's avatar
Tri Dao committed
261
        linear_cls = nn.Linear if not fused_bias_fc else FusedDense
Tri Dao's avatar
Tri Dao committed
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

        self.transform = BertPredictionHeadTransform(config)

        # The output weights are the same as the input embeddings, but there is
        # an output-only bias for each token.
        self.decoder = linear_cls(config.hidden_size, config.vocab_size, bias=True)

    def forward(self, hidden_states):
        hidden_states = self.transform(hidden_states)
        hidden_states = self.decoder(hidden_states)
        return hidden_states


class BertPreTrainingHeads(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.predictions = BertLMPredictionHead(config)
        self.seq_relationship = nn.Linear(config.hidden_size, 2)

    def forward(self, sequence_output, pooled_output):
        prediction_scores = self.predictions(sequence_output)
        seq_relationship_score = self.seq_relationship(pooled_output)
        return prediction_scores, seq_relationship_score


class BertPreTrainedModel(nn.Module):
Tri Dao's avatar
Tri Dao committed
288
289
    """An abstract class to handle weights initialization and
    a simple interface for dowloading and loading pretrained models.
Tri Dao's avatar
Tri Dao committed
290
    """
Tri Dao's avatar
Tri Dao committed
291

Tri Dao's avatar
Tri Dao committed
292
293
294
295
296
297
298
299
    def __init__(self, config, *inputs, **kwargs):
        super().__init__()
        if not isinstance(config, BertConfig):
            raise ValueError(
                "Parameter config in `{}(config)` should be an instance of class `BertConfig`. "
                "To create a model from a Google pretrained model use "
                "`model = {}.from_pretrained(PRETRAINED_MODEL_NAME)`".format(
                    self.__class__.__name__, self.__class__.__name__
Tri Dao's avatar
Tri Dao committed
300
301
                )
            )
Tri Dao's avatar
Tri Dao committed
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
        self.config = config

    @classmethod
    def from_pretrained(cls, model_name, config, *inputs, **kwargs):
        """
        Instantiate a BertPreTrainedModel from a pre-trained model file or a pytorch state dict.
        Download and cache the pre-trained model file if needed.

        Params:
            pretrained_model_name_or_path: either:
                - a path or url to a pretrained model archive containing:
                    . `bert_config.json` a configuration file for the model
                    . `pytorch_model.bin` a PyTorch dump of a BertForPretraining instance
                - a path or url to a pretrained model archive containing:
                    . `bert_config.json` a configuration file for the model
                    . `model.chkpt` a TensorFlow checkpoint
            *inputs, **kwargs: additional input for the specific Bert class
                (ex: num_labels for BertForSequenceClassification)
        """
        # Instantiate model.
        model = cls(config, *inputs, **kwargs)
Tri Dao's avatar
Tri Dao committed
323
324
325
        load_return = model.load_state_dict(
            remap_state_dict(state_dict_from_pretrained(model_name), config), strict=False
        )
Tri Dao's avatar
Tri Dao committed
326
327
328
329
330
331
332
        logger.info(load_return)
        return model


class BertModel(BertPreTrainedModel):
    def __init__(self, config: BertConfig, add_pooling_layer=True):
        super().__init__(config)
Tri Dao's avatar
Tri Dao committed
333
        self.pad_vocab_size_multiple = getattr(config, "pad_vocab_size_multiple", 1)
Tri Dao's avatar
Tri Dao committed
334
        if config.vocab_size % self.pad_vocab_size_multiple != 0:
Tri Dao's avatar
Tri Dao committed
335
336
337
338
            config.vocab_size += self.pad_vocab_size_multiple - (
                config.vocab_size % self.pad_vocab_size_multiple
            )
        self.fused_dropout_add_ln = getattr(config, "fused_dropout_add_ln", False)
339
        if self.fused_dropout_add_ln and layer_norm is None:
Tri Dao's avatar
Tri Dao committed
340
341
342
343
344
345
346
347
348
349
            raise ImportError("dropout_add_layer_norm is not installed")
        assert config.hidden_act in ["gelu", "gelu_new", "gelu_fast"]

        self.embeddings = BertEmbeddings(
            config.hidden_size,
            config.vocab_size,
            config.max_position_embeddings,
            config.type_vocab_size,
            padding_idx=config.pad_token_id,
        )
Tri Dao's avatar
Tri Dao committed
350
351
352
353
354
355
356
        self.emb_drop = nn.Dropout(config.hidden_dropout_prob)
        self.emb_ln = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.encoder = BertEncoder(config)
        self.pooler = BertPooler(config) if add_pooling_layer else None

        self.apply(partial(_init_weights, initializer_range=config.initializer_range))

Tri Dao's avatar
Tri Dao committed
357
358
359
360
361
362
363
364
    def forward(
        self,
        input_ids,
        position_ids=None,
        token_type_ids=None,
        attention_mask=None,
        masked_tokens_mask=None,
    ):
365
366
367
368
369
        """If masked_tokens_mask is not None (i.e. last_layer_subset == True in BertForPreTraining),
        we only want the output for the masked tokens. This means that we only compute the last
        layer output for these tokens.
        masked_tokens_mask: (batch, seqlen), dtype=torch.bool
        """
Tri Dao's avatar
Tri Dao committed
370
371
372
        hidden_states = self.embeddings(
            input_ids, position_ids=position_ids, token_type_ids=token_type_ids
        )
Tri Dao's avatar
Tri Dao committed
373
        # TD [2022-12:18]: Don't need to force residual in fp32
374
        # BERT puts embedding LayerNorm before embedding dropout.
Tri Dao's avatar
Tri Dao committed
375
376
377
        if not self.fused_dropout_add_ln:
            hidden_states = self.emb_ln(hidden_states)
        else:
Tri Dao's avatar
Tri Dao committed
378
379
380
            hidden_states = layer_norm(
                hidden_states, self.emb_ln.weight, self.emb_ln.bias, self.emb_ln.eps
            )
381
        hidden_states = self.emb_drop(hidden_states)
382
383
384
385

        if masked_tokens_mask is not None:
            batch_size, seqlen = input_ids.shape[:2]
            # We also need the first column for the CLS token
Tri Dao's avatar
Tri Dao committed
386
387
388
            first_col_mask = torch.zeros(
                batch_size, seqlen, dtype=torch.bool, device=input_ids.device
            )
389
390
391
392
393
            first_col_mask[:, 0] = True
            subset_mask = masked_tokens_mask | first_col_mask
        else:
            subset_mask = None

Tri Dao's avatar
Tri Dao committed
394
395
396
        sequence_output = self.encoder(
            hidden_states, key_padding_mask=attention_mask, subset_mask=subset_mask
        )
397
398
399
400
401
402
403
404
405
406
407
408

        if masked_tokens_mask is None:
            pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
        else:
            # TD [2022-03-01]: the indexing here is very tricky.
            if attention_mask is not None:
                subset_idx = subset_mask[attention_mask]
                pool_input = sequence_output[first_col_mask[attention_mask][subset_idx]]
                sequence_output = sequence_output[masked_tokens_mask[attention_mask][subset_idx]]
            else:
                pool_input = sequence_output[first_col_mask[subset_mask]]
                sequence_output = sequence_output[masked_tokens_mask[subset_mask]]
Tri Dao's avatar
Tri Dao committed
409
            pooled_output = self.pooler(pool_input, pool=False) if self.pooler is not None else None
410
411
412
413
414

        return BaseModelOutputWithPoolingAndCrossAttentions(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
        )
Tri Dao's avatar
Tri Dao committed
415
416
417
418
419
420
421


class BertForPreTraining(BertPreTrainedModel):
    def __init__(self, config: BertConfig):
        super().__init__(config)
        # If dense_seq_output, we only need to pass the hidden states for the masked out tokens
        # (around 15%) to the classifier heads.
Tri Dao's avatar
Tri Dao committed
422
        self.dense_seq_output = getattr(config, "dense_seq_output", False)
Tri Dao's avatar
Tri Dao committed
423
424
        # If last_layer_subset, we only need the compute the last layer for a subset of tokens
        # (e.g., the tokens we need to compute the masked LM loss and the next-sentence prediction).
Tri Dao's avatar
Tri Dao committed
425
        self.last_layer_subset = getattr(config, "last_layer_subset", False)
426
        if self.last_layer_subset:
Tri Dao's avatar
Tri Dao committed
427
428
            assert self.dense_seq_output, "last_layer_subset requires dense_seq_output"
        use_xentropy = getattr(config, "use_xentropy", False)
429
        if use_xentropy and CrossEntropyLoss is None:
Tri Dao's avatar
Tri Dao committed
430
431
432
433
434
435
            raise ImportError("xentropy_cuda is not installed")
        loss_cls = (
            nn.CrossEntropyLoss
            if not use_xentropy
            else partial(CrossEntropyLoss, inplace_backward=True)
        )
Tri Dao's avatar
Tri Dao committed
436
437
438
439
440
441
442
443
444
445
446
447
448

        self.bert = BertModel(config)
        self.cls = BertPreTrainingHeads(config)
        self.mlm_loss = loss_cls(ignore_index=0)
        self.nsp_loss = loss_cls(ignore_index=-1)

        # Initialize weights and apply final processing
        self.apply(partial(_init_weights, initializer_range=config.initializer_range))
        self.tie_weights()

    def tie_weights(self):
        self.cls.predictions.decoder.weight = self.bert.embeddings.word_embeddings.weight

Tri Dao's avatar
Tri Dao committed
449
450
451
452
453
454
455
456
457
    def forward(
        self,
        input_ids,
        position_ids=None,
        token_type_ids=None,
        attention_mask=None,
        labels=None,
        next_sentence_label=None,
    ):
Tri Dao's avatar
Tri Dao committed
458
        """
459
460
        If labels are provided, they must be 0 for masked out tokens (as specified in the attention
        mask).
Tri Dao's avatar
Tri Dao committed
461
462
463
464
465
466
467
468
469
470
471
        Outputs:
            if `labels` and `next_sentence_label` are not `None`:
                Outputs the total_loss which is the sum of the masked language modeling loss and the next
                sentence classification loss.
            if `labels` or `next_sentence_label` is `None`:
                Outputs a tuple comprising
                - the masked language modeling logits of shape [batch_size, sequence_length, vocab_size], and
                - the next sentence classification logits of shape [batch_size, 2].

        """
        masked_tokens_mask = labels > 0 if (self.last_layer_subset and labels is not None) else None
472
        outputs = self.bert(
Tri Dao's avatar
Tri Dao committed
473
474
475
            input_ids,
            position_ids=position_ids,
            token_type_ids=token_type_ids,
476
            attention_mask=attention_mask.bool() if attention_mask is not None else None,
Tri Dao's avatar
Tri Dao committed
477
            masked_tokens_mask=masked_tokens_mask,
Tri Dao's avatar
Tri Dao committed
478
        )
479
        sequence_output, pooled_output = outputs.last_hidden_state, outputs.pooler_output
Tri Dao's avatar
Tri Dao committed
480
481
482
        if self.dense_seq_output and labels is not None:
            masked_token_idx = torch.nonzero(labels.flatten() > 0, as_tuple=False).flatten()
            if not self.last_layer_subset:
Tri Dao's avatar
Tri Dao committed
483
484
485
                sequence_output = index_first_axis(
                    rearrange(sequence_output, "b s d -> (b s) d"), masked_token_idx
                )
Tri Dao's avatar
Tri Dao committed
486
487
        prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)

488
        total_loss = None
Tri Dao's avatar
Tri Dao committed
489
        if labels is not None and next_sentence_label is not None:
Tri Dao's avatar
Tri Dao committed
490
491
492
493
494
495
            if (
                self.dense_seq_output and labels is not None
            ):  # prediction_scores are already flattened
                masked_lm_loss = self.mlm_loss(
                    prediction_scores, labels.flatten()[masked_token_idx]
                )
Tri Dao's avatar
Tri Dao committed
496
            else:
Tri Dao's avatar
Tri Dao committed
497
498
499
500
501
502
503
504
                masked_lm_loss = self.mlm_loss(
                    rearrange(prediction_scores, "... v -> (...) v"),
                    rearrange(labels, "... -> (...)"),
                )
            next_sentence_loss = self.nsp_loss(
                rearrange(seq_relationship_score, "... t -> (...) t"),
                rearrange(next_sentence_label, "... -> (...)"),
            )
505
            total_loss = masked_lm_loss.float() + next_sentence_loss.float()
Tri Dao's avatar
Tri Dao committed
506

507
508
509
510
511
        return BertForPreTrainingOutput(
            loss=total_loss,
            prediction_logits=prediction_scores,
            seq_relationship_logits=seq_relationship_score,
        )
Tri Dao's avatar
Tri Dao committed
512
513
514
515
516


def remap_state_dict(state_dict, config):
    # LayerNorm
    def key_mapping_ln_gamma_beta(key):
Tri Dao's avatar
Tri Dao committed
517
518
        key = re.sub(r"LayerNorm.gamma$", "LayerNorm.weight", key)
        key = re.sub(r"LayerNorm.beta$", "LayerNorm.bias", key)
Tri Dao's avatar
Tri Dao committed
519
        return key
Tri Dao's avatar
Tri Dao committed
520

Tri Dao's avatar
Tri Dao committed
521
522
523
524
    state_dict = OrderedDict((key_mapping_ln_gamma_beta(k), v) for k, v in state_dict.items())

    # Layers
    def key_mapping_layers(key):
Tri Dao's avatar
Tri Dao committed
525
526
        return re.sub(r"^bert.encoder.layer.", "bert.encoder.layers.", key)

Tri Dao's avatar
Tri Dao committed
527
528
529
530
    state_dict = OrderedDict((key_mapping_layers(k), v) for k, v in state_dict.items())

    # LayerNorm
    def key_mapping_ln(key):
Tri Dao's avatar
Tri Dao committed
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
        key = re.sub(r"^bert.embeddings.LayerNorm.", "bert.emb_ln.", key)
        key = re.sub(
            r"^bert.encoder.layers.(\d+).attention.output.LayerNorm.(weight|bias)",
            r"bert.encoder.layers.\1.norm1.\2",
            key,
        )
        key = re.sub(
            r"^bert.encoder.layers.(\d+).output.LayerNorm.(weight|bias)",
            r"bert.encoder.layers.\1.norm2.\2",
            key,
        )
        key = re.sub(
            r"^cls.predictions.transform.LayerNorm.(weight|bias)",
            r"cls.predictions.transform.layer_norm.\1",
            key,
        )
Tri Dao's avatar
Tri Dao committed
547
        return key
Tri Dao's avatar
Tri Dao committed
548

Tri Dao's avatar
Tri Dao committed
549
550
551
552
    state_dict = OrderedDict((key_mapping_ln(k), v) for k, v in state_dict.items())

    # MLP
    def key_mapping_mlp(key):
Tri Dao's avatar
Tri Dao committed
553
554
555
556
557
558
559
560
561
562
        key = re.sub(
            r"^bert.encoder.layers.(\d+).intermediate.dense.(weight|bias)",
            r"bert.encoder.layers.\1.mlp.fc1.\2",
            key,
        )
        key = re.sub(
            r"^bert.encoder.layers.(\d+).output.dense.(weight|bias)",
            r"bert.encoder.layers.\1.mlp.fc2.\2",
            key,
        )
Tri Dao's avatar
Tri Dao committed
563
        return key
Tri Dao's avatar
Tri Dao committed
564

Tri Dao's avatar
Tri Dao committed
565
566
567
    state_dict = OrderedDict((key_mapping_mlp(k), v) for k, v in state_dict.items())

    # Attention
Tri Dao's avatar
Tri Dao committed
568
    last_layer_subset = getattr(config, "last_layer_subset", False)
Tri Dao's avatar
Tri Dao committed
569
    for d in range(config.num_hidden_layers):
Tri Dao's avatar
Tri Dao committed
570
571
572
573
574
575
        Wq = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.query.weight")
        Wk = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.key.weight")
        Wv = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.value.weight")
        bq = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.query.bias")
        bk = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.key.bias")
        bv = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.value.bias")
576
        if not (last_layer_subset and d == config.num_hidden_layers - 1):
Tri Dao's avatar
Tri Dao committed
577
            state_dict[f"bert.encoder.layers.{d}.mixer.Wqkv.weight"] = torch.cat(
578
579
                [Wq, Wk, Wv], dim=0
            )
Tri Dao's avatar
Tri Dao committed
580
            state_dict[f"bert.encoder.layers.{d}.mixer.Wqkv.bias"] = torch.cat([bq, bk, bv], dim=0)
581
        else:
Tri Dao's avatar
Tri Dao committed
582
583
584
585
586
            state_dict[f"bert.encoder.layers.{d}.mixer.Wq.weight"] = Wq
            state_dict[f"bert.encoder.layers.{d}.mixer.Wkv.weight"] = torch.cat([Wk, Wv], dim=0)
            state_dict[f"bert.encoder.layers.{d}.mixer.Wq.bias"] = bq
            state_dict[f"bert.encoder.layers.{d}.mixer.Wkv.bias"] = torch.cat([bk, bv], dim=0)

Tri Dao's avatar
Tri Dao committed
587
    def key_mapping_attn(key):
Tri Dao's avatar
Tri Dao committed
588
589
590
591
592
593
        return re.sub(
            r"^bert.encoder.layers.(\d+).attention.output.dense.(weight|bias)",
            r"bert.encoder.layers.\1.mixer.out_proj.\2",
            key,
        )

Tri Dao's avatar
Tri Dao committed
594
595
596
    state_dict = OrderedDict((key_mapping_attn(k), v) for k, v in state_dict.items())

    def key_mapping_decoder_bias(key):
Tri Dao's avatar
Tri Dao committed
597
598
        return re.sub(r"^cls.predictions.bias", "cls.predictions.decoder.bias", key)

Tri Dao's avatar
Tri Dao committed
599
600
    state_dict = OrderedDict((key_mapping_decoder_bias(k), v) for k, v in state_dict.items())

601
    # Word embedding
Tri Dao's avatar
Tri Dao committed
602
    pad_vocab_size_multiple = getattr(config, "pad_vocab_size_multiple", 1)
603
    if pad_vocab_size_multiple > 1:
Tri Dao's avatar
Tri Dao committed
604
605
        word_embeddings = state_dict["bert.embeddings.word_embeddings.weight"]
        state_dict["bert.embeddings.word_embeddings.weight"] = F.pad(
606
607
            word_embeddings, (0, 0, 0, config.vocab_size - word_embeddings.shape[0])
        )
Tri Dao's avatar
Tri Dao committed
608
609
        decoder_weight = state_dict["cls.predictions.decoder.weight"]
        state_dict["cls.predictions.decoder.weight"] = F.pad(
610
611
612
613
614
            decoder_weight, (0, 0, 0, config.vocab_size - decoder_weight.shape[0])
        )
        # If the vocab was padded, we want to set the decoder bias for those padded indices to be
        # strongly negative (i.e. the decoder shouldn't predict those indices).
        # TD [2022-05-09]: I don't think it affects the MLPerf training.
Tri Dao's avatar
Tri Dao committed
615
616
        decoder_bias = state_dict["cls.predictions.decoder.bias"]
        state_dict["cls.predictions.decoder.bias"] = F.pad(
617
618
619
            decoder_bias, (0, config.vocab_size - decoder_bias.shape[0]), value=-100.0
        )

Tri Dao's avatar
Tri Dao committed
620
    return state_dict