bert.py 32.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
Kevin Hu's avatar
Kevin Hu committed
13
from typing import Any, Mapping
Tri Dao's avatar
Tri Dao committed
14
15
16
17
18

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
Kevin Hu's avatar
Kevin Hu committed
19
from transformers import BertConfig, PretrainedConfig
Tri Dao's avatar
Tri Dao committed
20
from transformers.models.bert.modeling_bert import (
Kevin Hu's avatar
Kevin Hu committed
21
22
23
24
25
26
27
28
29
30
    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
31
32
from flash_attn.modules.block import Block
from flash_attn.modules.embedding import BertEmbeddings
Tri Dao's avatar
Tri Dao committed
33
34
from flash_attn.modules.mha import MHA
from flash_attn.modules.mlp import FusedMLP, Mlp
35
from flash_attn.utils.pretrained import state_dict_from_pretrained
Tri Dao's avatar
Tri Dao committed
36
37

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

try:
43
    from flash_attn.ops.triton.layer_norm import layer_norm_fn
Tri Dao's avatar
Tri Dao committed
44
except ImportError:
45
46
    layer_norm_fn = None

Tri Dao's avatar
Tri Dao committed
47
48

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


logger = logging.getLogger(__name__)


57
def create_mixer_cls(config, cross_attn=False, return_residual=False):
Tri Dao's avatar
Tri Dao committed
58
59
    use_flash_attn = getattr(config, "use_flash_attn", False)
    fused_bias_fc = getattr(config, "fused_bias_fc", False)
60
61
62
63
64
65
    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
66
67
68
69
70
71
72
73
74
75
76
    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
77
78
79
    return mixer_cls


80
def create_mlp_cls(config, layer_idx=None, return_residual=False):
Tri Dao's avatar
Tri Dao committed
81
    inner_dim = config.intermediate_size
Tri Dao's avatar
Tri Dao committed
82
    fused_mlp = getattr(config, "fused_mlp", False)
83
    if fused_mlp:
Kevin Hu's avatar
Kevin Hu committed
84
        assert config.hidden_act in ["gelu_new", "gelu_fast", "gelu_pytorch_tanh"], (
Tri Dao's avatar
Tri Dao committed
85
86
            "fused_mlp only " "supports approximate gelu"
        )
87
    if not fused_mlp:
Kevin Hu's avatar
Kevin Hu committed
88
89
90
91
92
        approximate = (
            "tanh"
            if config.hidden_act in ["gelu_new", "gelu_fast", "gelu_pytorch_tanh"]
            else "none"
        )
Tri Dao's avatar
Tri Dao committed
93
94
95
96
97
98
        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
99
    else:
100
        if FusedMLP is None:
Tri Dao's avatar
Tri Dao committed
101
102
            raise ImportError("fused_dense is not installed")
        mlp_checkpoint_lvl = getattr(config, "mlp_checkpoint_lvl", 0)
Tri Dao's avatar
Tri Dao committed
103
104
105
106
        # 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
107
108
109
110
111
112
        mlp_cls = partial(
            FusedMLP,
            hidden_features=inner_dim,
            checkpoint_lvl=mlp_checkpoint_lvl,
            return_residual=return_residual,
        )
Tri Dao's avatar
Tri Dao committed
113
114
115
116
    return mlp_cls


def create_block(config, layer_idx=None):
Tri Dao's avatar
Tri Dao committed
117
118
    last_layer_subset = getattr(config, "last_layer_subset", False)
    cross_attn = last_layer_subset and layer_idx == config.num_hidden_layers - 1
119
120
121
122
123
124
    # 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
125
    norm_cls = partial(nn.LayerNorm, eps=config.layer_norm_eps)
Tri Dao's avatar
Tri Dao committed
126
127
128
129
130
131
132
133
134
135
136
    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
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
    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
155
156
157
158
        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
159

160
161
162
163
164
    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
165
        if key_padding_mask is None or not self.use_flash_attn:
Tri Dao's avatar
Tri Dao committed
166
167
168
            mixer_kwargs = (
                {"key_padding_mask": key_padding_mask} if key_padding_mask is not None else None
            )
Tri Dao's avatar
Tri Dao committed
169
170
            for layer in self.layers:
                hidden_states = layer(hidden_states, mixer_kwargs=mixer_kwargs)
171
172
            if subset_mask is not None:
                hidden_states = hidden_states[subset_mask]
Tri Dao's avatar
Tri Dao committed
173
174
175
176
177
        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
178
            mixer_kwargs = {"cu_seqlens": cu_seqlens, "max_seqlen": max_seqlen_in_batch}
179
180
181
182
183
184
185
186
            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
187
188
189
                    subset_idx = torch.nonzero(
                        subset_mask[key_padding_mask], as_tuple=False
                    ).flatten()
190
                    subset_seqlens = (subset_mask & key_padding_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
                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
197
198
199
                    subset_cu_seqlens = F.pad(
                        torch.cumsum(subset_seqlens, dim=0, dtype=torch.torch.int32), (1, 0)
                    )
200
201
202
203
                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
204
205
206
207
208
209
210
                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,
                }
211
                hidden_states = self.layers[-1](hidden_states_subset, mixer_kwargs=mixer_kwargs)
Tri Dao's avatar
Tri Dao committed
212
213
214
215
216
217
        return hidden_states


class BertPooler(nn.Module):
    def __init__(self, config):
        super().__init__()
Tri Dao's avatar
Tri Dao committed
218
        fused_bias_fc = getattr(config, "fused_bias_fc", False)
Tri Dao's avatar
Tri Dao committed
219
        if fused_bias_fc and FusedDense is None:
Tri Dao's avatar
Tri Dao committed
220
            raise ImportError("fused_dense is not installed")
Tri Dao's avatar
Tri Dao committed
221
        linear_cls = nn.Linear if not fused_bias_fc else FusedDense
Tri Dao's avatar
Tri Dao committed
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
        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
237
        fused_bias_fc = getattr(config, "fused_bias_fc", False)
Tri Dao's avatar
Tri Dao committed
238
        if fused_bias_fc and FusedDense is None:
Tri Dao's avatar
Tri Dao committed
239
240
            raise ImportError("fused_dense is not installed")
        self.fused_dropout_add_ln = getattr(config, "fused_dropout_add_ln", False)
241
242
        if self.fused_dropout_add_ln and layer_norm_fn is None:
            raise ImportError("Triton is not installed")
Tri Dao's avatar
Tri Dao committed
243
        linear_cls = nn.Linear if not fused_bias_fc else FusedDense
Tri Dao's avatar
Tri Dao committed
244
        self.dense = linear_cls(config.hidden_size, config.hidden_size)
Kevin Hu's avatar
Kevin Hu committed
245
246
247
248
249
        approximate = (
            "tanh"
            if config.hidden_act in ["gelu_new", "gelu_fast", "gelu_pytorch_tanh"]
            else "none"
        )
250
        self.transform_act_fn = nn.GELU(approximate=approximate)
Tri Dao's avatar
Tri Dao committed
251
252
253
254
255
256
257
258
        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:
259
260
            hidden_states = layer_norm_fn(
                hidden_states, self.layer_norm.weight, self.layer_norm.bias, eps=self.layer_norm.eps
Tri Dao's avatar
Tri Dao committed
261
            )
Tri Dao's avatar
Tri Dao committed
262
263
264
265
266
267
        return hidden_states


class BertLMPredictionHead(nn.Module):
    def __init__(self, config):
        super().__init__()
Tri Dao's avatar
Tri Dao committed
268
        fused_bias_fc = getattr(config, "fused_bias_fc", False)
Tri Dao's avatar
Tri Dao committed
269
        if fused_bias_fc and FusedDense is None:
Tri Dao's avatar
Tri Dao committed
270
            raise ImportError("fused_dense is not installed")
Tri Dao's avatar
Tri Dao committed
271
        linear_cls = nn.Linear if not fused_bias_fc else FusedDense
Tri Dao's avatar
Tri Dao committed
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297

        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
298
299
    """An abstract class to handle weights initialization and
    a simple interface for dowloading and loading pretrained models.
Tri Dao's avatar
Tri Dao committed
300
    """
Tri Dao's avatar
Tri Dao committed
301

Tri Dao's avatar
Tri Dao committed
302
303
304
305
306
307
308
309
    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
310
311
                )
            )
Tri Dao's avatar
Tri Dao committed
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
        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
333
334
335
        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
336
337
338
339
340
341
342
        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
343
        self.pad_vocab_size_multiple = getattr(config, "pad_vocab_size_multiple", 1)
Tri Dao's avatar
Tri Dao committed
344
        if config.vocab_size % self.pad_vocab_size_multiple != 0:
Tri Dao's avatar
Tri Dao committed
345
346
347
348
            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)
349
350
        if self.fused_dropout_add_ln and layer_norm_fn is None:
            raise ImportError("Triton is not installed")
Kevin Hu's avatar
Kevin Hu committed
351
        assert config.hidden_act in ["gelu", "gelu_new", "gelu_fast", "gelu_pytorch_tanh"]
Tri Dao's avatar
Tri Dao committed
352
353
354
355
356
357
358
359

        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
360
361
362
363
364
365
366
        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
367
368
369
370
371
372
373
374
    def forward(
        self,
        input_ids,
        position_ids=None,
        token_type_ids=None,
        attention_mask=None,
        masked_tokens_mask=None,
    ):
375
376
377
378
379
        """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
380
381
382
        hidden_states = self.embeddings(
            input_ids, position_ids=position_ids, token_type_ids=token_type_ids
        )
Tri Dao's avatar
Tri Dao committed
383
        # TD [2022-12:18]: Don't need to force residual in fp32
384
        # BERT puts embedding LayerNorm before embedding dropout.
Tri Dao's avatar
Tri Dao committed
385
386
387
        if not self.fused_dropout_add_ln:
            hidden_states = self.emb_ln(hidden_states)
        else:
388
389
            hidden_states = layer_norm_fn(
                hidden_states, self.emb_ln.weight, self.emb_ln.bias, eps=self.emb_ln.eps
Tri Dao's avatar
Tri Dao committed
390
            )
391
        hidden_states = self.emb_drop(hidden_states)
392
393
394
395

        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
396
397
398
            first_col_mask = torch.zeros(
                batch_size, seqlen, dtype=torch.bool, device=input_ids.device
            )
399
400
401
402
403
            first_col_mask[:, 0] = True
            subset_mask = masked_tokens_mask | first_col_mask
        else:
            subset_mask = None

Tri Dao's avatar
Tri Dao committed
404
405
406
        sequence_output = self.encoder(
            hidden_states, key_padding_mask=attention_mask, subset_mask=subset_mask
        )
407
408
409
410
411
412
413
414
415
416
417
418

        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
419
            pooled_output = self.pooler(pool_input, pool=False) if self.pooler is not None else None
420
421
422
423
424

        return BaseModelOutputWithPoolingAndCrossAttentions(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
        )
Tri Dao's avatar
Tri Dao committed
425
426
427
428
429
430
431


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
432
        self.dense_seq_output = getattr(config, "dense_seq_output", False)
Tri Dao's avatar
Tri Dao committed
433
434
        # 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
435
        self.last_layer_subset = getattr(config, "last_layer_subset", False)
436
        if self.last_layer_subset:
Tri Dao's avatar
Tri Dao committed
437
438
            assert self.dense_seq_output, "last_layer_subset requires dense_seq_output"
        use_xentropy = getattr(config, "use_xentropy", False)
439
        if use_xentropy and CrossEntropyLoss is None:
Tri Dao's avatar
Tri Dao committed
440
441
442
443
444
445
            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
446
447
448
449
450
451
452
453
454
455
456
457
458

        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
459
460
461
462
463
464
465
466
467
    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
468
        """
469
470
        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
471
472
473
474
475
476
477
478
479
480
481
        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
482
        outputs = self.bert(
Tri Dao's avatar
Tri Dao committed
483
484
485
            input_ids,
            position_ids=position_ids,
            token_type_ids=token_type_ids,
486
            attention_mask=attention_mask.bool() if attention_mask is not None else None,
Tri Dao's avatar
Tri Dao committed
487
            masked_tokens_mask=masked_tokens_mask,
Tri Dao's avatar
Tri Dao committed
488
        )
489
        sequence_output, pooled_output = outputs.last_hidden_state, outputs.pooler_output
Tri Dao's avatar
Tri Dao committed
490
491
492
        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
493
494
495
                sequence_output = index_first_axis(
                    rearrange(sequence_output, "b s d -> (b s) d"), masked_token_idx
                )
Tri Dao's avatar
Tri Dao committed
496
497
        prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)

498
        total_loss = None
Tri Dao's avatar
Tri Dao committed
499
        if labels is not None and next_sentence_label is not None:
Tri Dao's avatar
Tri Dao committed
500
501
502
503
504
505
            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
506
            else:
Tri Dao's avatar
Tri Dao committed
507
508
509
510
511
512
513
514
                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, "... -> (...)"),
            )
515
            total_loss = masked_lm_loss.float() + next_sentence_loss.float()
Tri Dao's avatar
Tri Dao committed
516

517
518
519
520
521
        return BertForPreTrainingOutput(
            loss=total_loss,
            prediction_logits=prediction_scores,
            seq_relationship_logits=seq_relationship_score,
        )
Tri Dao's avatar
Tri Dao committed
522
523


Kevin Hu's avatar
Kevin Hu committed
524
525
526
527
528
def remap_state_dict(state_dict, config: PretrainedConfig):
    """
    Map the state_dict of a Huggingface BERT model to be flash_attn compatible.
    """

Tri Dao's avatar
Tri Dao committed
529
530
    # LayerNorm
    def key_mapping_ln_gamma_beta(key):
Tri Dao's avatar
Tri Dao committed
531
532
        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
533
        return key
Tri Dao's avatar
Tri Dao committed
534

Tri Dao's avatar
Tri Dao committed
535
536
537
538
    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
539
540
        return re.sub(r"^bert.encoder.layer.", "bert.encoder.layers.", key)

Tri Dao's avatar
Tri Dao committed
541
542
543
544
    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
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
        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
561
        return key
Tri Dao's avatar
Tri Dao committed
562

Tri Dao's avatar
Tri Dao committed
563
564
565
566
    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
567
568
569
570
571
572
573
574
575
576
        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
577
        return key
Tri Dao's avatar
Tri Dao committed
578

Tri Dao's avatar
Tri Dao committed
579
580
581
    state_dict = OrderedDict((key_mapping_mlp(k), v) for k, v in state_dict.items())

    # Attention
Tri Dao's avatar
Tri Dao committed
582
    last_layer_subset = getattr(config, "last_layer_subset", False)
Tri Dao's avatar
Tri Dao committed
583
    for d in range(config.num_hidden_layers):
Tri Dao's avatar
Tri Dao committed
584
585
586
587
588
589
        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")
590
        if not (last_layer_subset and d == config.num_hidden_layers - 1):
Tri Dao's avatar
Tri Dao committed
591
            state_dict[f"bert.encoder.layers.{d}.mixer.Wqkv.weight"] = torch.cat(
592
593
                [Wq, Wk, Wv], dim=0
            )
Tri Dao's avatar
Tri Dao committed
594
            state_dict[f"bert.encoder.layers.{d}.mixer.Wqkv.bias"] = torch.cat([bq, bk, bv], dim=0)
595
        else:
Tri Dao's avatar
Tri Dao committed
596
597
598
599
600
            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
601
    def key_mapping_attn(key):
Tri Dao's avatar
Tri Dao committed
602
603
604
605
606
607
        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
608
609
610
    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
611
612
        return re.sub(r"^cls.predictions.bias", "cls.predictions.decoder.bias", key)

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

615
    # Word embedding
Tri Dao's avatar
Tri Dao committed
616
    pad_vocab_size_multiple = getattr(config, "pad_vocab_size_multiple", 1)
617
    if pad_vocab_size_multiple > 1:
Tri Dao's avatar
Tri Dao committed
618
619
        word_embeddings = state_dict["bert.embeddings.word_embeddings.weight"]
        state_dict["bert.embeddings.word_embeddings.weight"] = F.pad(
620
621
            word_embeddings, (0, 0, 0, config.vocab_size - word_embeddings.shape[0])
        )
Tri Dao's avatar
Tri Dao committed
622
623
        decoder_weight = state_dict["cls.predictions.decoder.weight"]
        state_dict["cls.predictions.decoder.weight"] = F.pad(
624
625
626
627
628
            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
629
630
        decoder_bias = state_dict["cls.predictions.decoder.bias"]
        state_dict["cls.predictions.decoder.bias"] = F.pad(
631
632
633
            decoder_bias, (0, config.vocab_size - decoder_bias.shape[0]), value=-100.0
        )

Tri Dao's avatar
Tri Dao committed
634
    return state_dict
Kevin Hu's avatar
Kevin Hu committed
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764


def inv_remap_state_dict(state_dict, config: PretrainedConfig):
    """
    Map the state_dict of a flash_attn model to be Huggingface BERT compatible.

    This function is meant to be the inverse of remap_state_dict.
    """
    # Word embedding
    pad_vocab_size_multiple = getattr(config, "pad_vocab_size_multiple", 1)
    if pad_vocab_size_multiple > 1:
        word_embeddings = state_dict["bert.embeddings.word_embeddings.weight"]
        decoder_weight = state_dict["cls.predictions.decoder.weight"]
        decoder_bias = state_dict["cls.predictions.decoder.bias"]
        # unpad embeddings
        state_dict["bert.embeddings.word_embeddings.weight"] = word_embeddings[
            : config.orig_vocab_size, :
        ]
        state_dict["cls.predictions.decoder.weight"] = decoder_weight[: config.orig_vocab_size, :]
        state_dict["cls.predictions.decoder.bias"] = decoder_bias[: config.orig_vocab_size]

    for d in range(config.num_hidden_layers):
        last_layer_subset = getattr(config, "last_layer_subset", False)
        if not last_layer_subset or d != (config.num_hidden_layers - 1):
            Wqkv_weights = state_dict.pop(f"bert.encoder.layers.{d}.mixer.Wqkv.weight")
            Wqkv_biases = state_dict.pop(f"bert.encoder.layers.{d}.mixer.Wqkv.bias")
            state_dict[f"bert.encoder.layers.{d}.attention.self.query.weight"] = Wqkv_weights[
                : Wqkv_weights.shape[0] // 3, :
            ]
            state_dict[f"bert.encoder.layers.{d}.attention.self.key.weight"] = Wqkv_weights[
                Wqkv_weights.shape[0] // 3 : 2 * Wqkv_weights.shape[0] // 3, :
            ]
            state_dict[f"bert.encoder.layers.{d}.attention.self.value.weight"] = Wqkv_weights[
                2 * Wqkv_weights.shape[0] // 3 :, :
            ]
            state_dict[f"bert.encoder.layers.{d}.attention.self.query.bias"] = Wqkv_biases[
                : Wqkv_biases.shape[0] // 3
            ]
            state_dict[f"bert.encoder.layers.{d}.attention.self.key.bias"] = Wqkv_biases[
                Wqkv_biases.shape[0] // 3 : 2 * Wqkv_biases.shape[0] // 3
            ]
            state_dict[f"bert.encoder.layers.{d}.attention.self.value.bias"] = Wqkv_biases[
                2 * Wqkv_biases.shape[0] // 3 :
            ]
        else:
            Wq_weight = state_dict.pop(f"bert.encoder.layers.{d}.mixer.Wq.weight")
            Wkv_weights = state_dict.pop(f"bert.encoder.layers.{d}.mixer.Wkv.weight")
            Wq_bias = state_dict.pop(f"bert.encoder.layers.{d}.mixer.Wq.bias")
            Wkv_biases = state_dict.pop(f"bert.encoder.layers.{d}.mixer.Wkv.bias")
            state_dict[f"bert.encoder.layers.{d}.attention.self.query.weight"] = Wq_weight
            state_dict[f"bert.encoder.layers.{d}.attention.self.key.weight"] = Wkv_weights[
                : Wkv_weights.shape[0] // 2, :
            ]
            state_dict[f"bert.encoder.layers.{d}.attention.self.value.weight"] = Wkv_weights[
                Wkv_weights.shape[0] // 2 :, :
            ]
            state_dict[f"bert.encoder.layers.{d}.attention.self.query.bias"] = Wq_bias
            state_dict[f"bert.encoder.layers.{d}.attention.self.key.bias"] = Wkv_biases[
                : Wkv_biases.shape[0] // 2
            ]
            state_dict[f"bert.encoder.layers.{d}.attention.self.value.bias"] = Wkv_biases[
                Wkv_biases.shape[0] // 2 :
            ]

    def inv_key_mapping_ln(key):
        key = re.sub(r"bert.emb_ln.", "bert.embeddings.LayerNorm.", key)
        key = re.sub(
            r"bert.encoder.layers.(\d+).norm1.(weight|bias)",
            r"bert.encoder.layers.\1.attention.output.LayerNorm.\2",
            key,
        )
        key = re.sub(
            r"bert.encoder.layers.(\d+).norm2.(weight|bias)",
            r"bert.encoder.layers.\1.output.LayerNorm.\2",
            key,
        )
        key = re.sub(
            r"cls.predictions.transform.layer_norm.(weight|bias)",
            r"cls.predictions.transform.LayerNorm.\1",
            key,
        )
        return key

    def inv_key_mapping_ln_gamma_beta(key):
        key = re.sub(r"LayerNorm.weight$", "LayerNorm.gamma", key)
        key = re.sub(r"LayerNorm.bias$", "LayerNorm.beta", key)
        return key

    def inv_key_mapping_layers(key):
        return re.sub(r"bert.encoder.layers.", "bert.encoder.layer.", key)

    def inv_key_mapping_mlp(key):
        key = re.sub(
            r"bert.encoder.layer.(\d+).mlp.fc1.(weight|bias)",
            r"bert.encoder.layer.\1.intermediate.dense.\2",
            key,
        )
        key = re.sub(
            r"bert.encoder.layer.(\d+).mlp.fc2.(weight|bias)",
            r"bert.encoder.layer.\1.output.dense.\2",
            key,
        )
        return key

    def inv_key_mapping_attn(key):
        return re.sub(
            r"bert.encoder.layer.(\d+).mixer.out_proj.(weight|bias)",
            r"bert.encoder.layer.\1.attention.output.dense.\2",
            key,
        )

    def inv_key_mapping_decoder_bias(key):
        return re.sub(r"cls.predictions.decoder.bias", "cls.predictions.bias", key)

    state_dict = OrderedDict((inv_key_mapping_ln(key), value) for key, value in state_dict.items())
    state_dict = OrderedDict(
        (inv_key_mapping_ln_gamma_beta(key), value) for key, value in state_dict.items()
    )
    state_dict = OrderedDict(
        (inv_key_mapping_layers(key), value) for key, value in state_dict.items()
    )
    state_dict = OrderedDict((inv_key_mapping_mlp(key), value) for key, value in state_dict.items())
    state_dict = OrderedDict(
        (inv_key_mapping_attn(key), value) for key, value in state_dict.items()
    )
    state_dict = OrderedDict(
        (inv_key_mapping_decoder_bias(key), value) for key, value in state_dict.items()
    )

    return state_dict