gpt.py 44 KB
Newer Older
Tri Dao's avatar
Tri Dao committed
1
# Copyright (c) 2023, Tri Dao.
Tri Dao's avatar
Tri Dao committed
2

3
import logging
Tri Dao's avatar
Tri Dao committed
4
import math
5
import re
Tri Dao's avatar
Tri Dao committed
6
from collections import OrderedDict, namedtuple
Tri Dao's avatar
Tri Dao committed
7
from collections.abc import Sequence
Tri Dao's avatar
Tri Dao committed
8
from functools import partial
Tri Dao's avatar
Tri Dao committed
9
10
11
12

import torch
import torch.nn as nn
import torch.nn.functional as F
13
from einops import rearrange
Tri Dao's avatar
Tri Dao committed
14
15
from transformers import GPT2Config

Tri Dao's avatar
Tri Dao committed
16
17
18
19
from flash_attn.models.falcon import remap_state_dict_hf_falcon
from flash_attn.models.gpt_neox import remap_state_dict_hf_gpt_neox
from flash_attn.models.gptj import remap_state_dict_hf_gptj
from flash_attn.models.opt import remap_state_dict_hf_opt
Tri Dao's avatar
Tri Dao committed
20
from flash_attn.modules.block import Block, ParallelBlock
21
from flash_attn.modules.embedding import GPT2Embeddings, ParallelGPT2Embeddings
Tri Dao's avatar
Tri Dao committed
22
23
24
25
26
27
28
29
30
31
from flash_attn.modules.mha import MHA, ParallelMHA
from flash_attn.modules.mlp import (
    FusedMLP,
    GatedMlp,
    Mlp,
    ParallelFusedMLP,
    ParallelGatedMlp,
    ParallelMLP,
)
from flash_attn.ops.activations import sqrelu_fwd
Tri Dao's avatar
Tri Dao committed
32
from flash_attn.utils.distributed import all_gather_raw, get_dim_for_local_rank, sync_shared_params
Tri Dao's avatar
Tri Dao committed
33
from flash_attn.utils.generation import GenerationMixin
Tri Dao's avatar
Tri Dao committed
34
from flash_attn.utils.pretrained import state_dict_from_pretrained
35
36
37
38
39

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

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

46
47
48
49
50
try:
    from flash_attn.ops.layer_norm import dropout_add_layer_norm_parallel_residual
except ImportError:
    dropout_add_layer_norm_parallel_residual = None

Tri Dao's avatar
Tri Dao committed
51
52
53
try:
    from flash_attn.ops.rms_norm import RMSNorm, dropout_add_rms_norm
except ImportError:
54
    RMSNorm, dropout_add_rms_norm = None, None
Tri Dao's avatar
Tri Dao committed
55
56
57
58
59
60

try:
    from flash_attn.ops.rms_norm import dropout_add_rms_norm_parallel_residual
except ImportError:
    dropout_add_rms_norm_parallel_residual = None

Tri Dao's avatar
Tri Dao committed
61
try:
Tri Dao's avatar
Tri Dao committed
62
    from flash_attn.ops.triton.mlp import FusedDenseSqreluDense
Tri Dao's avatar
Tri Dao committed
63
64
65
except ImportError:
    FusedDenseSqreluDense = None

66
67
68
logger = logging.getLogger(__name__)


69
def create_mixer_cls(config, layer_idx=None, process_group=None, device=None, dtype=None):
Tri Dao's avatar
Tri Dao committed
70
71
    factory_kwargs = {"device": device, "dtype": dtype}
    head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
Tri Dao's avatar
Tri Dao committed
72
73
74
75
    softmax_scale = 1.0 if not config.scale_attn_weights else head_dim ** (-0.5)
    if config.scale_attn_by_inverse_layer_idx:
        assert layer_idx is not None
        softmax_scale /= float(layer_idx + 1)
Tri Dao's avatar
Tri Dao committed
76
    dwconv = getattr(config, "attn_dwconv", False)
77
    if dwconv:
Tri Dao's avatar
Tri Dao committed
78
79
80
81
82
83
84
85
86
        assert process_group is None, "TensorParallel MHA does not support dwconv yet"
    qkv_proj_bias = getattr(config, "qkv_proj_bias", True)
    out_proj_bias = getattr(config, "out_proj_bias", True)
    rotary_emb_dim = int(getattr(config, "rotary_emb_fraction", 0.0) * head_dim)
    rotary_emb_base = getattr(config, "rotary_emb_base", 10000.0)
    rotary_emb_scale_base = getattr(config, "rotary_emb_scale_base", None)
    rotary_emb_interleaved = getattr(config, "rotary_emb_interleaved", False)
    use_flash_attn = getattr(config, "use_flash_attn", False)
    fused_bias_fc = getattr(config, "fused_bias_fc", False)
87
    if not fused_bias_fc:
Tri Dao's avatar
Tri Dao committed
88
        assert process_group is None, "TensorParallel MHA requires fused_bias_fc"
89
    mha_cls = MHA if process_group is None else ParallelMHA
Tri Dao's avatar
Tri Dao committed
90
91
92
93
94
95
96
97
98
99
100
    serial_kwargs = (
        {"fused_bias_fc": fused_bias_fc, "dwconv": dwconv} if process_group is None else {}
    )
    parallel_kwargs = (
        {
            "process_group": process_group,
            "sequence_parallel": getattr(config, "sequence_parallel", True),
        }
        if process_group is not None
        else {}
    )
Tri Dao's avatar
Tri Dao committed
101
    num_heads_kv = getattr(config, "n_head_kv", None)
Tri Dao's avatar
Tri Dao committed
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
    mixer_cls = partial(
        mha_cls,
        num_heads=config.num_attention_heads,
        num_heads_kv=num_heads_kv,
        qkv_proj_bias=qkv_proj_bias,
        out_proj_bias=out_proj_bias,
        dropout=config.attn_pdrop,
        softmax_scale=softmax_scale,
        causal=True,
        layer_idx=layer_idx,
        rotary_emb_dim=rotary_emb_dim,
        rotary_emb_base=rotary_emb_base,
        rotary_emb_scale_base=rotary_emb_scale_base,
        rotary_emb_interleaved=rotary_emb_interleaved,
        use_flash_attn=use_flash_attn,
        **serial_kwargs,
        **parallel_kwargs,
        **factory_kwargs,
    )
Tri Dao's avatar
Tri Dao committed
121
122
123
    return mixer_cls


124
def create_mlp_cls(config, layer_idx=None, process_group=None, device=None, dtype=None):
Tri Dao's avatar
Tri Dao committed
125
126
127
128
    factory_kwargs = {"device": device, "dtype": dtype}
    mlp_fc1_bias = getattr(config, "mlp_fc1_bias", True)
    mlp_fc2_bias = getattr(config, "mlp_fc2_bias", True)
    fused_mlp = getattr(config, "fused_mlp", False)
129
    if fused_mlp:
Tri Dao's avatar
Tri Dao committed
130
131
132
133
134
135
136
137
        assert config.activation_function in [
            "gelu_new",
            "gelu_fast",
            "gelu_approx",
            "relu",
            "sqrelu",
        ]
    fused_dense_sqrelu_dense = getattr(config, "fused_dense_sqrelu_dense", False)
138
    if fused_dense_sqrelu_dense:
Tri Dao's avatar
Tri Dao committed
139
140
141
        assert config.activation_function == "sqrelu", (
            "fused_dense_sqrelu_dense only " "supports approximate activation_function sqrelu"
        )
142
143
    assert not (fused_dense_sqrelu_dense and fused_mlp)
    if not fused_mlp and not fused_dense_sqrelu_dense:
Tri Dao's avatar
Tri Dao committed
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
        assert config.activation_function in [
            "gelu",
            "gelu_new",
            "gelu_fast",
            "gelu_approx",
            "relu",
            "sqrelu",
            "glu",
            "swiglu",
            "geglu",
        ]
        if config.activation_function in ["glu", "swiglu", "geglu"]:
            activation = (
                F.sigmoid
                if config.activation_function == "glu"
                else (F.silu if config.activation_function == "swiglu" else F.gelu)
            )
161
            mlp_cls = GatedMlp if process_group is None else ParallelGatedMlp
Tri Dao's avatar
Tri Dao committed
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
            parallel_kwargs = (
                {
                    "process_group": process_group,
                    "sequence_parallel": getattr(config, "sequence_parallel", True),
                }
                if process_group is not None
                else {}
            )
            mlp_cls = partial(
                mlp_cls,
                hidden_features=config.n_inner,
                activation=activation,
                bias1=mlp_fc1_bias,
                bias2=mlp_fc2_bias,
                **parallel_kwargs,
                **factory_kwargs,
            )
Tri Dao's avatar
Tri Dao committed
179
        else:
Tri Dao's avatar
Tri Dao committed
180
            if config.activation_function == "relu":
Tri Dao's avatar
Tri Dao committed
181
                activation = partial(F.relu, inplace=True)
Tri Dao's avatar
Tri Dao committed
182
            elif config.activation_function == "sqrelu":
Tri Dao's avatar
Tri Dao committed
183
184
                activation = sqrelu_fwd
            else:
Tri Dao's avatar
Tri Dao committed
185
186
187
188
189
190
                approximate = (
                    "tanh"
                    if config.activation_function in ["gelu_new", "gelu_fast", "gelu_approx"]
                    else "none"
                )
                activation = partial(F.gelu, approximate=approximate)
Tri Dao's avatar
Tri Dao committed
191
            mlp_cls = Mlp if process_group is None else ParallelMLP
Tri Dao's avatar
Tri Dao committed
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
            parallel_kwargs = (
                {
                    "process_group": process_group,
                    "sequence_parallel": getattr(config, "sequence_parallel", True),
                }
                if process_group is not None
                else {}
            )
            mlp_cls = partial(
                mlp_cls,
                hidden_features=config.n_inner,
                activation=activation,
                bias1=mlp_fc1_bias,
                bias2=mlp_fc2_bias,
                **parallel_kwargs,
                **factory_kwargs,
            )
Tri Dao's avatar
Tri Dao committed
209
    else:
Tri Dao's avatar
Tri Dao committed
210
        mlp_checkpoint_lvl = getattr(config, "mlp_checkpoint_lvl", 0)
Tri Dao's avatar
Tri Dao committed
211
212
213
214
        # 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]
215
216
        if fused_mlp:
            if FusedMLP is None:
Tri Dao's avatar
Tri Dao committed
217
218
219
220
221
222
                raise ImportError("fused_dense is not installed")
            activation = (
                "gelu_approx"
                if config.activation_function in ["gelu_new", "gelu_fast", "gelu_approx"]
                else config.activation_function
            )
223
            mlp_cls = FusedMLP if process_group is None else ParallelFusedMLP
Tri Dao's avatar
Tri Dao committed
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
            parallel_kwargs = (
                {
                    "process_group": process_group,
                    "sequence_parallel": getattr(config, "sequence_parallel", True),
                }
                if process_group is not None
                else {}
            )
            mlp_cls = partial(
                mlp_cls,
                hidden_features=config.n_inner,
                activation=activation,
                checkpoint_lvl=mlp_checkpoint_lvl,
                bias1=mlp_fc1_bias,
                bias2=mlp_fc2_bias,
                **parallel_kwargs,
                **factory_kwargs,
            )
Tri Dao's avatar
Tri Dao committed
242
        elif fused_dense_sqrelu_dense:
243
            if process_group is not None:
Tri Dao's avatar
Tri Dao committed
244
                assert fused_mlp, "Tensor Parallel is not implemented for FusedDenseSqreluDense"
Tri Dao's avatar
Tri Dao committed
245
            assert FusedDenseSqreluDense is not None
Tri Dao's avatar
Tri Dao committed
246
247
248
249
250
251
            mlp_cls = partial(
                FusedDenseSqreluDense,
                hidden_features=config.n_inner,
                checkpoint_lvl=mlp_checkpoint_lvl,
                **factory_kwargs,
            )
Tri Dao's avatar
Tri Dao committed
252
        else:
Tri Dao's avatar
Tri Dao committed
253
            raise RuntimeError("MLP type not supported")
Tri Dao's avatar
Tri Dao committed
254
255
256
    return mlp_cls


257
def create_block(config, layer_idx=None, process_group=None, device=None, dtype=None):
Tri Dao's avatar
Tri Dao committed
258
259
    factory_kwargs = {"device": device, "dtype": dtype}
    sequence_parallel = getattr(config, "sequence_parallel", True)
260
261
    mixer_cls = create_mixer_cls(config, layer_idx, process_group=process_group, **factory_kwargs)
    mlp_cls = create_mlp_cls(config, layer_idx, process_group=process_group, **factory_kwargs)
Tri Dao's avatar
Tri Dao committed
262
263
264
265
266
267
    use_rms_norm = getattr(config, "rms_norm", False)
    norm_cls = partial(
        nn.LayerNorm if not use_rms_norm else RMSNorm,
        eps=config.layer_norm_epsilon,
        **factory_kwargs,
    )
Tri Dao's avatar
Tri Dao committed
268
    # TD [2022-07-30]: Force residual in fp32, seems to make fp16 training more stable
Tri Dao's avatar
Tri Dao committed
269
    residual_in_fp32 = getattr(config, "residual_in_fp32", False)
Tri Dao's avatar
Tri Dao committed
270
    resid_dropout1 = config.resid_pdrop if layer_idx is None or layer_idx > 0 else config.embd_pdrop
Tri Dao's avatar
Tri Dao committed
271
272
    prenorm = getattr(config, "prenorm", True)
    parallel_block = getattr(config, "parallel_block", False)
Tri Dao's avatar
Tri Dao committed
273
274
    if not parallel_block:
        block = Block(
Tri Dao's avatar
Tri Dao committed
275
276
277
278
279
280
281
282
            config.hidden_size,
            mixer_cls,
            mlp_cls,
            norm_cls=norm_cls,
            prenorm=prenorm,
            resid_dropout1=resid_dropout1,
            resid_dropout2=config.resid_pdrop,
            fused_dropout_add_ln=getattr(config, "fused_dropout_add_ln", False),
Tri Dao's avatar
Tri Dao committed
283
284
            residual_in_fp32=residual_in_fp32,
            sequence_parallel=sequence_parallel and process_group is not None,
Tri Dao's avatar
Tri Dao committed
285
            mark_shared_params=process_group is not None,
Tri Dao's avatar
Tri Dao committed
286
287
288
289
        )
    else:
        assert prenorm
        block = ParallelBlock(
Tri Dao's avatar
Tri Dao committed
290
291
292
293
294
295
296
297
            config.hidden_size,
            mixer_cls,
            mlp_cls,
            norm_cls=norm_cls,
            resid_dropout1=resid_dropout1,
            resid_dropout2=config.resid_pdrop,
            tied_norm=getattr(config, "parallel_block_tied_norm", False),
            fused_dropout_add_ln=getattr(config, "fused_dropout_add_ln", False),
Tri Dao's avatar
Tri Dao committed
298
299
            residual_in_fp32=residual_in_fp32,
            sequence_parallel=sequence_parallel and process_group is not None,
Tri Dao's avatar
Tri Dao committed
300
            mark_shared_params=process_group is not None,
Tri Dao's avatar
Tri Dao committed
301
        )
Tri Dao's avatar
Tri Dao committed
302
303
304
305
    block.layer_idx = layer_idx
    return block


306
class GPTPreTrainedModel(nn.Module):
Tri Dao's avatar
Tri Dao committed
307
308
    """An abstract class to handle weights initialization and
    a simple interface for dowloading and loading pretrained models.
309
    """
Tri Dao's avatar
Tri Dao committed
310

311
312
313
314
315
316
317
318
    def __init__(self, config, *inputs, **kwargs):
        super().__init__()
        if not isinstance(config, GPT2Config):
            raise ValueError(
                "Parameter config in `{}(config)` should be an instance of class `GPT2Config`. "
                "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
319
320
                )
            )
321
322
323
        self.config = config

    @classmethod
Tri Dao's avatar
Tri Dao committed
324
325
326
327
328
329
330
331
332
333
334
335
    def from_pretrained(
        cls,
        model_name,
        config,
        *args,
        strict=True,
        device=None,
        dtype=None,
        world_size=1,
        rank=0,
        **kwargs,
    ):
336
337
338
339
340
        """
        Instantiate a GPTPreTrainedModel from a pre-trained model file or a pytorch state dict.
        Download and cache the pre-trained model file if needed.
        """
        # Instantiate model.
341
        model = cls(config, *args, device=device, dtype=dtype, **kwargs)
342
343
        # Load state_dict in cpu because we already initialized the model in GPU, and we don't
        # want extra stuff taking up more GPU memory
Tri Dao's avatar
Tri Dao committed
344
345
        state_dict = state_dict_from_pretrained(model_name, device="cpu", dtype=dtype)
        if model_name.startswith("gpt2"):
Tri Dao's avatar
Tri Dao committed
346
            state_dict = remap_state_dict_hf_gpt2(state_dict, config)
Tri Dao's avatar
Tri Dao committed
347
        elif model_name.startswith("facebook/opt"):
Tri Dao's avatar
Tri Dao committed
348
            state_dict = remap_state_dict_hf_opt(state_dict, config)
Tri Dao's avatar
Tri Dao committed
349
        elif model_name.startswith("EleutherAI/gpt-j-"):
Tri Dao's avatar
Tri Dao committed
350
            state_dict = remap_state_dict_hf_gptj(state_dict, config)
Tri Dao's avatar
Tri Dao committed
351
        elif model_name.startswith("EleutherAI/gpt-neox-"):
Tri Dao's avatar
Tri Dao committed
352
            state_dict = remap_state_dict_hf_gpt_neox(state_dict, config)
Tri Dao's avatar
Tri Dao committed
353
        elif model_name.startswith("tiiuae/falcon-"):
Tri Dao's avatar
Tri Dao committed
354
            state_dict = remap_state_dict_hf_falcon(state_dict, config)
Tri Dao's avatar
Tri Dao committed
355
        else:
Tri Dao's avatar
Tri Dao committed
356
            raise NotImplementedError(f"Model {model_name} not supported")
357
358
359
        if world_size > 1:
            state_dict = shard_state_dict_tp(state_dict, config, world_size, rank)
        load_return = model.load_state_dict(state_dict, strict=strict)
360
361
362
        logger.info(load_return)
        return model

Tri Dao's avatar
Tri Dao committed
363

Tri Dao's avatar
Tri Dao committed
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
# https://github.com/huggingface/transformers/blob/c28d04e9e252a1a099944e325685f14d242ecdcd/src/transformers/models/gpt2/modeling_gpt2.py#L454
def _init_weights(module, n_layer, initializer_range=0.02, rescale_prenorm_residual=True):
    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 rescale_prenorm_residual:
        # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
        #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
        #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
        #   >   -- GPT-2 :: https://openai.com/blog/better-language-models/
        #
        # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
        for name, p in module.named_parameters():
            if name in ["out_proj.weight", "fc2.weight"]:
                # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
                nn.init.normal_(p, mean=0.0, std=initializer_range / math.sqrt(2 * n_layer))


386
class GPTModel(GPTPreTrainedModel):
387
    def __init__(self, config: GPT2Config, process_group=None, device=None, dtype=None):
388
        super().__init__(config)
Tri Dao's avatar
Tri Dao committed
389
        factory_kwargs = {"device": device, "dtype": dtype}
390
        self.process_group = process_group
Tri Dao's avatar
Tri Dao committed
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
        self.sequence_parallel = getattr(config, "sequence_parallel", True)
        assert config.activation_function in [
            "gelu",
            "gelu_new",
            "gelu_fast",
            "gelu_approx",
            "relu",
            "sqrelu",
            "glu",
            "swiglu",
            "geglu",
        ]
        pad_vocab_size_multiple = getattr(config, "pad_vocab_size_multiple", 1)
        vocab_size = (
            math.ceil(config.vocab_size / pad_vocab_size_multiple) * pad_vocab_size_multiple
        )
Tri Dao's avatar
Tri Dao committed
407
        # TD [2022-07-30]: Force residual in fp32, seems to make fp16 training more stable
Tri Dao's avatar
Tri Dao committed
408
        self.residual_in_fp32 = getattr(config, "residual_in_fp32", False)
Tri Dao's avatar
Tri Dao committed
409
        # These 2 options are for OPT-350m
Tri Dao's avatar
Tri Dao committed
410
411
412
        self.prenorm = getattr(config, "prenorm", True)
        use_rms_norm = getattr(config, "rms_norm", False)
        word_embed_proj_dim = getattr(config, "word_embed_proj_dim", None)
Tri Dao's avatar
Tri Dao committed
413
        # For GPT-J, GPT-NeoX
Tri Dao's avatar
Tri Dao committed
414
        self.parallel_block = getattr(config, "parallel_block", False)
Tri Dao's avatar
Tri Dao committed
415

416
        if process_group is None:
Tri Dao's avatar
Tri Dao committed
417
            self.embeddings = GPT2Embeddings(
Tri Dao's avatar
Tri Dao committed
418
419
420
421
422
                config.hidden_size,
                vocab_size,
                config.max_position_embeddings,
                word_embed_proj_dim=word_embed_proj_dim,
                **factory_kwargs,
Tri Dao's avatar
Tri Dao committed
423
            )
424
425
        else:
            self.embeddings = ParallelGPT2Embeddings(
Tri Dao's avatar
Tri Dao committed
426
427
428
429
430
431
                config.hidden_size,
                vocab_size,
                config.max_position_embeddings,
                process_group=process_group,
                sequence_parallel=self.sequence_parallel,
                **factory_kwargs,
432
            )
Tri Dao's avatar
Tri Dao committed
433

Tri Dao's avatar
Tri Dao committed
434
        # We change the order of dropout, residual and layer norm:
Tri Dao's avatar
Tri Dao committed
435
        # Instead of LN -> Attn / MLP -> Dropout -> Add, we do:
Tri Dao's avatar
Tri Dao committed
436
437
438
        # Dropout -> Add -> LN -> Attn / MLP, returning both the residual branch (output of Add) and
        # the main branch (output of MLP). The model definition is unchanged, but the mapping of the
        # nn.Dropout probabilities are changed.
Tri Dao's avatar
Tri Dao committed
439
        # This is for performance reason: we can fuse dropout + add + layer_norm.
Tri Dao's avatar
Tri Dao committed
440
441
442
443
444
445
        self.layers = nn.ModuleList(
            [
                create_block(config, layer_idx=i, process_group=process_group, **factory_kwargs)
                for i in range(config.num_hidden_layers)
            ]
        )
Tri Dao's avatar
Tri Dao committed
446

Tri Dao's avatar
Tri Dao committed
447
        self.fused_dropout_add_ln = getattr(config, "fused_dropout_add_ln", False)
448
        if self.fused_dropout_add_ln:
Tri Dao's avatar
Tri Dao committed
449
450
451
452
            if (not self.parallel_block and dropout_add_layer_norm is None) or (
                self.parallel_block and dropout_add_layer_norm_parallel_residual is None
            ):
                raise ImportError("dropout_layer_norm is not installed")
Tri Dao's avatar
Tri Dao committed
453
454
        if self.prenorm:
            self.drop_f = nn.Dropout(config.resid_pdrop)
Tri Dao's avatar
Tri Dao committed
455
            norm_cls = nn.LayerNorm if not use_rms_norm else RMSNorm
Tri Dao's avatar
Tri Dao committed
456
457
458
            self.ln_f = norm_cls(
                config.hidden_size, eps=config.layer_norm_epsilon, **factory_kwargs
            )
459
        if process_group is not None:
Tri Dao's avatar
Tri Dao committed
460
            for p in self.ln_f.parameters():
461
462
463
464
465
                # Mark the norm parameters as "shared_params" so that we sync their values at init.
                p._shared_params = True
                # Mark the norm params as "sequence_parallel" so we run all-reduce on their grads.
                if self.sequence_parallel:
                    p._sequence_parallel = True
466

Tri Dao's avatar
Tri Dao committed
467
468
469
470
471
472
473
        self.apply(
            partial(
                _init_weights,
                n_layer=config.num_hidden_layers,
                initializer_range=config.initializer_range,
            )
        )
474
475
476
        self.tie_weights()

    def tie_weights(self):
477
        if self.process_group is not None:
478
            sync_shared_params(self, self.process_group)
Tri Dao's avatar
Tri Dao committed
479

480
    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
Tri Dao's avatar
Tri Dao committed
481
482
483
484
        return {
            i: layer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)
            for i, layer in enumerate(self.layers)
        }
485

Tri Dao's avatar
Tri Dao committed
486
    def forward(self, input_ids, position_ids=None, inference_params=None):
487
488
489
        # If using Tensor Parallel with sequence parallel, we combine the batch and the seqlen
        # dimensions so that we can split on it easily, in case of small batch size.
        # Only the attention layers need to know the seqlen.
Tri Dao's avatar
Tri Dao committed
490
491
492
493
494
        embedding_kwargs = (
            {"combine_batch_seqlen_dim": True}
            if self.process_group is not None and self.sequence_parallel
            else {}
        )
495
        hidden_states = self.embeddings(input_ids, position_ids=position_ids, **embedding_kwargs)
Tri Dao's avatar
Tri Dao committed
496
497
        if self.parallel_block:
            hidden_states2 = None
Tri Dao's avatar
Tri Dao committed
498
        residual = None
Tri Dao's avatar
Tri Dao committed
499
500
501
502
503
        mixer_kwargs = (
            {"seqlen": input_ids.shape[1]}
            if self.process_group is not None and self.sequence_parallel
            else {}
        )
Tri Dao's avatar
Tri Dao committed
504
        if inference_params is not None:
Tri Dao's avatar
Tri Dao committed
505
            mixer_kwargs["inference_params"] = inference_params
Tri Dao's avatar
Tri Dao committed
506
        for layer in self.layers:
Tri Dao's avatar
Tri Dao committed
507
            if self.prenorm:
Tri Dao's avatar
Tri Dao committed
508
                if not self.parallel_block:
Tri Dao's avatar
Tri Dao committed
509
510
511
                    hidden_states, residual = layer(
                        hidden_states, residual, mixer_kwargs=mixer_kwargs
                    )
Tri Dao's avatar
Tri Dao committed
512
513
514
515
                else:
                    hidden_states, hidden_states2, residual = layer(
                        hidden_states, hidden_states2, residual, mixer_kwargs=mixer_kwargs
                    )
Tri Dao's avatar
Tri Dao committed
516
517
518
519
520
            else:
                hidden_states = layer(hidden_states, mixer_kwargs=mixer_kwargs)
        if self.prenorm:
            if not self.fused_dropout_add_ln:
                dropped = self.drop_f(hidden_states)
Tri Dao's avatar
Tri Dao committed
521
522
523
524
                if not self.parallel_block:
                    residual = (dropped + residual) if residual is not None else dropped
                else:
                    dropped2 = self.drop_f(hidden_states2)
Tri Dao's avatar
Tri Dao committed
525
526
527
528
529
                    residual = (
                        (residual + dropped + dropped2)
                        if residual is not None
                        else dropped + dropped2
                    )
Tri Dao's avatar
Tri Dao committed
530
531
                hidden_states = self.ln_f(residual.to(dtype=self.ln_f.weight.dtype))
            else:
Tri Dao's avatar
Tri Dao committed
532
                # Set prenorm=False here since we don't need the residual
533
                if not self.parallel_block:
Tri Dao's avatar
Tri Dao committed
534
535
536
537
538
                    fused_add_norm_fn = (
                        dropout_add_rms_norm
                        if isinstance(self.ln_f, RMSNorm)
                        else dropout_add_layer_norm
                    )
539
                    hidden_states = fused_add_norm_fn(
Tri Dao's avatar
Tri Dao committed
540
541
542
543
544
545
546
547
                        hidden_states,
                        residual,
                        self.ln_f.weight,
                        self.ln_f.bias,
                        self.drop_f.p if self.training else 0.0,
                        self.ln_f.eps,
                        prenorm=False,
                        residual_in_fp32=self.residual_in_fp32,
548
549
                    )
                else:
Tri Dao's avatar
Tri Dao committed
550
551
552
553
554
                    fused_add_norm_fn = (
                        dropout_add_rms_norm_parallel_residual
                        if isinstance(self.ln_f, RMSNorm)
                        else dropout_add_layer_norm_parallel_residual
                    )
555
                    hidden_states, _ = fused_add_norm_fn(
Tri Dao's avatar
Tri Dao committed
556
557
558
559
560
561
562
563
564
565
566
                        hidden_states,
                        hidden_states2,
                        residual,
                        self.ln_f.weight,
                        self.ln_f.bias,
                        None,
                        None,
                        self.drop_f.p if self.training else 0.0,
                        self.ln_f.eps,
                        prenorm=False,
                        residual_in_fp32=self.residual_in_fp32,
567
                    )
Tri Dao's avatar
Tri Dao committed
568
569
570
        return hidden_states


Tri Dao's avatar
Tri Dao committed
571
class GPTLMHeadModel(GPTPreTrainedModel, GenerationMixin):
572
    def __init__(self, config: GPT2Config, process_group=None, device=None, dtype=None):
Tri Dao's avatar
Tri Dao committed
573
        factory_kwargs = {"device": device, "dtype": dtype}
574
        super().__init__(config)
575
576
        self.process_group = process_group
        self.transformer = GPTModel(config, process_group=process_group, **factory_kwargs)
Tri Dao's avatar
Tri Dao committed
577
578
579
580
581
582
        self.tie_word_embeddings = getattr(config, "tie_word_embeddings", True)
        lm_head_bias = getattr(config, "lm_head_bias", False)
        pad_vocab_size_multiple = getattr(config, "pad_vocab_size_multiple", 1)
        vocab_size = (
            math.ceil(config.vocab_size / pad_vocab_size_multiple) * pad_vocab_size_multiple
        )
Tri Dao's avatar
Tri Dao committed
583
        # This option is for OPT-350m
Tri Dao's avatar
Tri Dao committed
584
        word_embed_proj_dim = getattr(config, "word_embed_proj_dim", None)
Tri Dao's avatar
Tri Dao committed
585
586
587
588
589
        embed_dim = config.n_embd if word_embed_proj_dim is None else word_embed_proj_dim
        if word_embed_proj_dim is not None:
            self.project_out = nn.Linear(config.n_embd, embed_dim, bias=False, **factory_kwargs)
        else:
            self.project_out = None
590
        if process_group is None:
Tri Dao's avatar
Tri Dao committed
591
            self.lm_head = nn.Linear(embed_dim, vocab_size, bias=lm_head_bias, **factory_kwargs)
592
593
        else:
            if ColumnParallelLinear is None:
Tri Dao's avatar
Tri Dao committed
594
                raise ImportError("fused_dense_lib is not installed")
595
            self.lm_head = ColumnParallelLinear(
Tri Dao's avatar
Tri Dao committed
596
597
598
599
600
601
                embed_dim,
                vocab_size,
                process_group,
                bias=lm_head_bias,
                sequence_parallel=getattr(config, "sequence_parallel", True),
                **factory_kwargs,
602
            )
Tri Dao's avatar
Tri Dao committed
603
        # Initialize weights and apply final processing
Tri Dao's avatar
Tri Dao committed
604
605
606
607
608
609
610
        self.apply(
            partial(
                _init_weights,
                n_layer=config.num_hidden_layers,
                initializer_range=config.initializer_range,
            )
        )
Tri Dao's avatar
Tri Dao committed
611
612
613
        self.tie_weights()

    def tie_weights(self):
Tri Dao's avatar
Tri Dao committed
614
615
        if self.tie_word_embeddings:
            self.lm_head.weight = self.transformer.embeddings.word_embeddings.weight
616
        if self.process_group is not None:
617
            sync_shared_params(self, self.process_group)
Tri Dao's avatar
Tri Dao committed
618

619
    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
Tri Dao's avatar
Tri Dao committed
620
621
622
        return self.transformer.allocate_inference_cache(
            batch_size, max_seqlen, dtype=dtype, **kwargs
        )
623

624
    def forward(self, input_ids, position_ids=None, inference_params=None, last_token_only=False):
Tri Dao's avatar
Tri Dao committed
625
        """
Tri Dao's avatar
Tri Dao committed
626
627
628
629
        inference_params: for generation. Adapted from Megatron-LM (and Apex)
        https://github.com/NVIDIA/apex/blob/3ff1a10f72ec07067c4e44759442329804ac5162/apex/transformer/testing/standalone_transformer_lm.py#L470
        last_token_only: whether to return the logit for the last token only,
            of shape (batch_size, vocab_size)
Tri Dao's avatar
Tri Dao committed
630
        """
Tri Dao's avatar
Tri Dao committed
631
632
633
        hidden_states = self.transformer(
            input_ids, position_ids=position_ids, inference_params=inference_params
        )
634
635
        if last_token_only:
            hidden_states = hidden_states[:, -1]
Tri Dao's avatar
Tri Dao committed
636
637
        if self.project_out is not None:
            hidden_states = self.project_out(hidden_states)
Tri Dao's avatar
Tri Dao committed
638
        lm_logits = self.lm_head(hidden_states)
639
640
641
        # During inference, we want the full logit for sampling
        if isinstance(self.lm_head, ColumnParallelLinear) and inference_params is not None:
            lm_logits, _ = all_gather_raw(lm_logits, self.lm_head.process_group)
Tri Dao's avatar
Tri Dao committed
642
643
            lm_logits = rearrange(lm_logits, "(n b) ... d -> b ... (n d)", b=hidden_states.shape[0])
        CausalLMOutput = namedtuple("CausalLMOutput", ["logits"])
Tri Dao's avatar
Tri Dao committed
644
        return CausalLMOutput(logits=lm_logits)
645

Tri Dao's avatar
Tri Dao committed
646
647
648
649
    def load_state_dict(self, state_dict, strict=True):
        # Remapping from our checkpoints that used a different ordering of layers in the block
        # Previous: Attn / MLP -> Dropout -> Add -> LN
        # Current: Dropout -> Add -> LN -> Attn / MLP
Tri Dao's avatar
Tri Dao committed
650
        if "transformer.ln_0.weight" in state_dict:
Tri Dao's avatar
Tri Dao committed
651
            n_layers = len(self.transformer.layers)
Tri Dao's avatar
Tri Dao committed
652
653
654
655
            ln_weight = state_dict.pop(f"transformer.layers.{n_layers - 1}.norm2.weight")
            ln_bias = state_dict.pop(f"transformer.layers.{n_layers - 1}.norm2.bias")
            state_dict["transformer.ln_f.weight"] = ln_weight
            state_dict["transformer.ln_f.bias"] = ln_bias
Tri Dao's avatar
Tri Dao committed
656
            for l in reversed(range(n_layers)):
Tri Dao's avatar
Tri Dao committed
657
658
659
660
                ln_weight = state_dict.pop(f"transformer.layers.{l}.norm1.weight")
                ln_bias = state_dict.pop(f"transformer.layers.{l}.norm1.bias")
                state_dict[f"transformer.layers.{l}.norm2.weight"] = ln_weight
                state_dict[f"transformer.layers.{l}.norm2.bias"] = ln_bias
Tri Dao's avatar
Tri Dao committed
661
                if l > 0:
Tri Dao's avatar
Tri Dao committed
662
663
664
665
666
667
668
669
                    ln_weight = state_dict.pop(f"transformer.layers.{l - 1}.norm2.weight")
                    ln_bias = state_dict.pop(f"transformer.layers.{l - 1}.norm2.bias")
                    state_dict[f"transformer.layers.{l}.norm1.weight"] = ln_weight
                    state_dict[f"transformer.layers.{l}.norm1.bias"] = ln_bias
            ln_weight = state_dict.pop("transformer.ln_0.weight")
            ln_bias = state_dict.pop("transformer.ln_0.bias")
            state_dict[f"transformer.layers.0.norm1.weight"] = ln_weight
            state_dict[f"transformer.layers.0.norm1.bias"] = ln_bias
Tri Dao's avatar
Tri Dao committed
670
671
        return super().load_state_dict(state_dict, strict=strict)

672

Tri Dao's avatar
Tri Dao committed
673
674
675
676
def shard_state_dict_tp(state_dict, config, world_size, rank):
    """Convert the state_dict of a standard GPT model to the state_dict of a GPT model
    with tensor parallel.
    """
Tri Dao's avatar
Tri Dao committed
677
678
    pad_vocab_size_multiple = getattr(config, "pad_vocab_size_multiple", 1)
    vocab_size = math.ceil(config.vocab_size / pad_vocab_size_multiple) * pad_vocab_size_multiple
Tri Dao's avatar
Tri Dao committed
679
680
681
682
683
    assert vocab_size % world_size == 0
    assert config.hidden_size % world_size == 0
    inner_dim = config.n_inner if config.n_inner is not None else 4 * config.hidden_size
    assert inner_dim % world_size == 0

684
685
686
687
688
689
    n_head = config.n_head
    n_head_kv = getattr(config, "n_head_kv", n_head)

    embed_dim = config.hidden_size
    head_dim = embed_dim // n_head

Tri Dao's avatar
Tri Dao committed
690
    def shard_first_dim(state_dict, key):
Tri Dao's avatar
Tri Dao committed
691
692
693
        if key in state_dict:
            x = state_dict[key]
            dim = x.shape[0] // world_size
Tri Dao's avatar
Tri Dao committed
694
            state_dict[key] = x[rank * dim : (rank + 1) * dim]
Tri Dao's avatar
Tri Dao committed
695

696
    def shard_last_dim(state_dict, key, multiple_of=1):
Tri Dao's avatar
Tri Dao committed
697
698
        if key in state_dict:
            x = state_dict[key]
699
700
701
702
703
704
            dim_each_rank = [
                get_dim_for_local_rank(x.size(-1), world_size, local_rank, multiple_of)
                for local_rank in range(world_size)
            ]
            beg, end = tuple(sum(dim_each_rank[:pos]) for pos in (rank, rank + 1))
            state_dict[key] = x[..., beg:end]
Tri Dao's avatar
Tri Dao committed
705

Tri Dao's avatar
Tri Dao committed
706
707
708
709
710
    def shard_gatedmlp_fc1_dim(state_dict, key):
        if key in state_dict:
            x = state_dict[key]
            dim = x.shape[0] // world_size // 2
            state_dict[key] = rearrange(
Tri Dao's avatar
Tri Dao committed
711
                rearrange(x, "(two o) ... -> two o ...", two=2)[:, rank * dim : (rank + 1) * dim],
Tri Dao's avatar
Tri Dao committed
712
                "two o ... -> (two o) ...",
Tri Dao's avatar
Tri Dao committed
713
714
            )

Tri Dao's avatar
Tri Dao committed
715
    def shard_qkv_headdim(state_dict, key):
Tri Dao's avatar
Tri Dao committed
716
        if key in state_dict:
717
            n_head_each_rank = [
Tri Dao's avatar
Tri Dao committed
718
719
                get_dim_for_local_rank(n_head, world_size, local_rank)
                for local_rank in range(world_size)
720
721
            ]
            n_head_kv_each_rank = [
Tri Dao's avatar
Tri Dao committed
722
723
                get_dim_for_local_rank(n_head_kv, world_size, local_rank)
                for local_rank in range(world_size)
724
725
726
727
728
729
730
731
            ]

            beg_n_head = sum(n_head_each_rank[:rank])
            end_n_head = sum(n_head_each_rank[: rank + 1])

            beg_n_head_kv = sum(n_head_kv_each_rank[:rank])
            end_n_head_kv = sum(n_head_kv_each_rank[: rank + 1])

Tri Dao's avatar
Tri Dao committed
732
            if n_head_kv == n_head:
Tri Dao's avatar
Tri Dao committed
733
734
                x = rearrange(state_dict[key], "(three d) ... -> three d ...", three=3)
                state_dict[key] = rearrange(
Tri Dao's avatar
Tri Dao committed
735
736
                    x[:, beg_n_head * head_dim : end_n_head * head_dim],
                    "three d ... -> (three d) ...",
Tri Dao's avatar
Tri Dao committed
737
                )
Tri Dao's avatar
Tri Dao committed
738
            else:
Tri Dao's avatar
Tri Dao committed
739
740
741
742
743
744
745
746
                x = rearrange(
                    state_dict[key],
                    "(nheadqkv headdim) ... -> nheadqkv headdim ...",
                    nheadqkv=n_head + 2 * n_head_kv,
                )
                state_dict[key] = rearrange(
                    torch.cat(
                        [
747
                            x[beg_n_head:end_n_head],
Tri Dao's avatar
Tri Dao committed
748
749
750
751
752
753
754
755
                            x[n_head + beg_n_head_kv : n_head + end_n_head_kv],
                            x[
                                n_head
                                + n_head_kv
                                + beg_n_head_kv : n_head
                                + n_head_kv
                                + end_n_head_kv
                            ],
Tri Dao's avatar
Tri Dao committed
756
757
758
759
760
761
762
763
764
765
766
                        ],
                        dim=0,
                    ),
                    "nheadqkv headdim ... -> (nheadqkv headdim) ...",
                )

    shard_first_dim(state_dict, "transformer.embeddings.word_embeddings.weight")
    if "lm_head.weight" in state_dict:
        shard_first_dim(state_dict, "lm_head.weight")
    if "transformer.embeddings.position_embeddings.weight" in state_dict:
        shard_last_dim(state_dict, "transformer.embeddings.position_embeddings.weight")
Tri Dao's avatar
Tri Dao committed
767
    for i in range(config.num_hidden_layers):
Tri Dao's avatar
Tri Dao committed
768
769
        shard_qkv_headdim(state_dict, f"transformer.layers.{i}.mixer.Wqkv.weight")
        shard_qkv_headdim(state_dict, f"transformer.layers.{i}.mixer.Wqkv.bias")
770
771
772
        shard_last_dim(
            state_dict, f"transformer.layers.{i}.mixer.out_proj.weight", multiple_of=head_dim
        )
Tri Dao's avatar
Tri Dao committed
773
        if rank != 0:
Tri Dao's avatar
Tri Dao committed
774
            state_dict.pop(f"transformer.layers.{i}.mixer.out_proj.bias", None)
Tri Dao's avatar
Tri Dao committed
775
        if config.activation_function in ["glu", "swiglu", "geglu"]:
Tri Dao's avatar
Tri Dao committed
776
777
            shard_gatedmlp_fc1_dim(state_dict, f"transformer.layers.{i}.mlp.fc1.weight")
            shard_gatedmlp_fc1_dim(state_dict, f"transformer.layers.{i}.mlp.fc1.bias")
Tri Dao's avatar
Tri Dao committed
778
        else:
Tri Dao's avatar
Tri Dao committed
779
780
781
            shard_first_dim(state_dict, f"transformer.layers.{i}.mlp.fc1.weight")
            shard_first_dim(state_dict, f"transformer.layers.{i}.mlp.fc1.bias")
        shard_last_dim(state_dict, f"transformer.layers.{i}.mlp.fc2.weight")
Tri Dao's avatar
Tri Dao committed
782
        if rank != 0:
Tri Dao's avatar
Tri Dao committed
783
            state_dict.pop(f"transformer.layers.{i}.mlp.fc2.bias", None)
Tri Dao's avatar
Tri Dao committed
784
785
786
787
    return state_dict


def combine_state_dicts_tp(state_dicts, config):
788
789
790
791
    """Convert the state_dict of a GPT model with tensor parallel to the state_dict of a
    standard GPT model.

    This function is meant to be the "reverse" of shard_state_dict_tp.
Tri Dao's avatar
Tri Dao committed
792
793
794
    """
    world_size = len(state_dicts)
    keys = state_dicts[0].keys()
Tri Dao's avatar
Tri Dao committed
795
796
    pad_vocab_size_multiple = getattr(config, "pad_vocab_size_multiple", 1)
    vocab_size = math.ceil(config.vocab_size / pad_vocab_size_multiple) * pad_vocab_size_multiple
Tri Dao's avatar
Tri Dao committed
797
798
799
800
801
    assert vocab_size % world_size == 0
    assert config.hidden_size % world_size == 0
    inner_dim = config.n_inner if config.n_inner is not None else 4 * config.hidden_size
    assert inner_dim % world_size == 0

Tri Dao's avatar
Tri Dao committed
802
    # Sometimes the word embeddings are sharded on the 0th dim, sometimes on the 1st dim.
Tri Dao's avatar
Tri Dao committed
803
804
    # vocab_size // world_size coordinates are nonzero.
    def combine_word_embeddings(state_dicts, state_dict, key):
Tri Dao's avatar
Tri Dao committed
805
806
        dim = 0 if state_dicts[0][key].shape[0] == vocab_size // world_size else 1
        state_dict[key] = torch.cat([s[key] for s in state_dicts], dim=dim)
Tri Dao's avatar
Tri Dao committed
807
808

    def combine_dim(state_dicts, state_dict, key, dim=-1):
Tri Dao's avatar
Tri Dao committed
809
810
        if key in state_dict:
            state_dict[key] = torch.cat([s[key] for s in state_dicts], dim=dim)
Tri Dao's avatar
Tri Dao committed
811
812

    def combine_qkv_headdim(state_dicts, state_dict, key):
Tri Dao's avatar
Tri Dao committed
813
        n_head = config.n_head
Tri Dao's avatar
Tri Dao committed
814
        n_head_kv = getattr(config, "n_head_kv", n_head)
Tri Dao's avatar
Tri Dao committed
815
816
817
        assert n_head % world_size == 0 and n_head_kv % world_size == 0
        n_head_per_rank = n_head // world_size
        n_head_kv_per_rank = n_head_kv // world_size
Tri Dao's avatar
Tri Dao committed
818
        if key in state_dict:
Tri Dao's avatar
Tri Dao committed
819
            if n_head_kv == n_head:
Tri Dao's avatar
Tri Dao committed
820
821
822
823
                xs = [
                    rearrange(s[key], "(three d) ... -> three d ...", three=3) for s in state_dicts
                ]
                state_dict[key] = rearrange(torch.cat(xs, dim=1), "three d ... -> (three d) ...")
Tri Dao's avatar
Tri Dao committed
824
            else:
Tri Dao's avatar
Tri Dao committed
825
826
827
828
829
830
831
832
833
834
835
836
837
838
                xs = [
                    rearrange(
                        s[key],
                        "(nheadqkv headdim) ... -> nheadqkv headdim ...",
                        nheadqkv=n_head + 2 * n_head_kv,
                    )
                    for s in state_dicts
                ]
                state_dict[key] = rearrange(
                    torch.cat(
                        [
                            torch.cat([x[:n_head_per_rank] for x in xs], dim=0),
                            torch.cat(
                                [
Tri Dao's avatar
Tri Dao committed
839
                                    x[n_head_per_rank : n_head_per_rank + n_head_kv_per_rank]
Tri Dao's avatar
Tri Dao committed
840
841
842
843
844
845
846
847
848
849
                                    for x in xs
                                ],
                                dim=0,
                            ),
                            torch.cat([x[-n_head_kv_per_rank:] for x in xs], dim=0),
                        ],
                        dim=0,
                    ),
                    "nheadqkv headdim ... -> (nheadqkv headdim) ...",
                )
Tri Dao's avatar
Tri Dao committed
850
851
852

    def combine_gated_mlp(state_dicts, state_dict, key):
        if key in state_dict:
Tri Dao's avatar
Tri Dao committed
853
854
            xs = [rearrange(s[key], "(two d) ... -> two d ...", two=2) for s in state_dicts]
            state_dict[key] = rearrange(torch.cat(xs, dim=1), "two d ... -> (two d) ...")
Tri Dao's avatar
Tri Dao committed
855
856

    state_dict = state_dicts[0].copy()  # don't modify state_dict[0] inplace
Tri Dao's avatar
Tri Dao committed
857
858
859
860
861
862
863
864
865
866
867
868
869
870
    combine_word_embeddings(
        state_dicts, state_dict, "transformer.embeddings.word_embeddings.weight"
    )
    if "lm_head.weight" in state_dict:
        combine_word_embeddings(state_dicts, state_dict, "lm_head.weight")
    if "transformer.embeddings.position_embeddings.weight" in state_dict:
        combine_dim(
            state_dicts, state_dict, "transformer.embeddings.position_embeddings.weight", -1
        )
    mlp_combine_fn = (
        combine_gated_mlp
        if config.activation_function in ["glu", "swiglu", "geglu"]
        else partial(combine_dim, dim=0)
    )
Tri Dao's avatar
Tri Dao committed
871
    for i in range(config.num_hidden_layers):
Tri Dao's avatar
Tri Dao committed
872
873
874
875
876
877
        combine_qkv_headdim(state_dicts, state_dict, f"transformer.layers.{i}.mixer.Wqkv.weight")
        combine_qkv_headdim(state_dicts, state_dict, f"transformer.layers.{i}.mixer.Wqkv.bias")
        combine_dim(state_dicts, state_dict, f"transformer.layers.{i}.mixer.out_proj.weight", -1)
        mlp_combine_fn(state_dicts, state_dict, f"transformer.layers.{i}.mlp.fc1.weight")
        combine_dim(state_dicts, state_dict, f"transformer.layers.{i}.mlp.fc1.bias", 0)
        combine_dim(state_dicts, state_dict, f"transformer.layers.{i}.mlp.fc2.weight", -1)
Tri Dao's avatar
Tri Dao committed
878
879
880
881
    return state_dict


def remap_state_dict_hf_gpt2(state_dict, config):
882
883
    # Word embedding and position embedding
    def key_mapping_pos_emb(key):
Tri Dao's avatar
Tri Dao committed
884
885
        return re.sub(r"^wpe.", "transformer.embeddings.position_embeddings.", key)

886
    state_dict = OrderedDict((key_mapping_pos_emb(k), v) for k, v in state_dict.items())
Tri Dao's avatar
Tri Dao committed
887
    word_embeddings = state_dict.pop("wte.weight")
888
    # It's possible that vocab_size is padded to be a multiple of 8, for example.
Tri Dao's avatar
Tri Dao committed
889
890
891
    pad_vocab_size_multiple = getattr(config, "pad_vocab_size_multiple", 1)
    vocab_size = math.ceil(config.vocab_size / pad_vocab_size_multiple) * pad_vocab_size_multiple
    state_dict["transformer.embeddings.word_embeddings.weight"] = F.pad(
892
        word_embeddings, (0, 0, 0, vocab_size - word_embeddings.shape[0])
893
    )
Tri Dao's avatar
Tri Dao committed
894
    state_dict["lm_head.weight"] = state_dict["transformer.embeddings.word_embeddings.weight"]
895
896

    # LayerNorm
Tri Dao's avatar
Tri Dao committed
897
    def key_mapping_ln(key):
Tri Dao's avatar
Tri Dao committed
898
899
        key = re.sub(r"^ln_f.(weight|bias)", r"transformer.ln_f.\1", key)
        key = re.sub(r"^h.(\d+).ln_(1|2).(weight|bias)", r"transformer.layers.\1.norm\2.\3", key)
Tri Dao's avatar
Tri Dao committed
900
        return key
Tri Dao's avatar
Tri Dao committed
901

Tri Dao's avatar
Tri Dao committed
902
    state_dict = OrderedDict((key_mapping_ln(k), v) for k, v in state_dict.items())
903
904
905

    # MLP
    for d in range(config.num_hidden_layers):
Tri Dao's avatar
Tri Dao committed
906
907
908
909
910
        W1 = state_dict.pop(f"h.{d}.mlp.c_fc.weight")
        state_dict[f"transformer.layers.{d}.mlp.fc1.weight"] = W1.t()
        W2 = state_dict.pop(f"h.{d}.mlp.c_proj.weight")
        state_dict[f"transformer.layers.{d}.mlp.fc2.weight"] = W2.t()

911
    def key_mapping_mlp(key):
Tri Dao's avatar
Tri Dao committed
912
913
        key = re.sub(r"^h.(\d+).mlp.c_fc.bias", r"transformer.layers.\1.mlp.fc1.bias", key)
        key = re.sub(r"^h.(\d+).mlp.c_proj.bias", r"transformer.layers.\1.mlp.fc2.bias", key)
914
        return key
Tri Dao's avatar
Tri Dao committed
915

916
917
918
919
    state_dict = OrderedDict((key_mapping_mlp(k), v) for k, v in state_dict.items())

    # Attention
    for d in range(config.num_hidden_layers):
Tri Dao's avatar
Tri Dao committed
920
921
922
923
924
925
        state_dict.pop(f"h.{d}.attn.bias")  # We don't store this bias
        Wqkv = state_dict.pop(f"h.{d}.attn.c_attn.weight")
        state_dict[f"transformer.layers.{d}.mixer.Wqkv.weight"] = Wqkv.t()
        Wout = state_dict.pop(f"h.{d}.attn.c_proj.weight")
        state_dict[f"transformer.layers.{d}.mixer.out_proj.weight"] = Wout.t()

926
    def key_mapping_attn(key):
Tri Dao's avatar
Tri Dao committed
927
928
929
930
        key = re.sub(r"^h.(\d+).attn.c_attn.bias", r"transformer.layers.\1.mixer.Wqkv.bias", key)
        key = re.sub(
            r"^h.(\d+).attn.c_proj.bias", r"transformer.layers.\1.mixer.out_proj.bias", key
        )
931
        return key
Tri Dao's avatar
Tri Dao committed
932

933
934
935
    state_dict = OrderedDict((key_mapping_attn(k), v) for k, v in state_dict.items())

    return state_dict
936
937


Tri Dao's avatar
Tri Dao committed
938
939
def remap_state_dict_megatron(state_dict, config):
    def key_mapping_transformer(key):
Tri Dao's avatar
Tri Dao committed
940
941
        key = re.sub(r"^language_model.encoder.", "transformer.", key)
        key = re.sub(r"^language_model.", "transformer.", key)
Tri Dao's avatar
Tri Dao committed
942
        return key
Tri Dao's avatar
Tri Dao committed
943

Tri Dao's avatar
Tri Dao committed
944
    state_dict = OrderedDict((key_mapping_transformer(k), v) for k, v in state_dict.items())
945

Tri Dao's avatar
Tri Dao committed
946
947
    # Word embedding and position embedding
    def key_mapping_pos_emb(key):
Tri Dao's avatar
Tri Dao committed
948
949
        return re.sub(r"^wpe.", "transformer.embeddings.position_embeddings.", key)

Tri Dao's avatar
Tri Dao committed
950
    state_dict = OrderedDict((key_mapping_pos_emb(k), v) for k, v in state_dict.items())
Tri Dao's avatar
Tri Dao committed
951
    word_embeddings = state_dict.pop("transformer.embedding.word_embeddings.weight")
Tri Dao's avatar
Tri Dao committed
952
    # It's possible that vocab_size is padded to be a multiple of 8, for example.
Tri Dao's avatar
Tri Dao committed
953
954
955
956
957
    pad_vocab_size_multiple = getattr(config, "pad_vocab_size_multiple", 1)
    vocab_size = (
        math.ceil(word_embeddings.shape[0] / pad_vocab_size_multiple) * pad_vocab_size_multiple
    )
    state_dict["transformer.embeddings.word_embeddings.weight"] = F.pad(
Tri Dao's avatar
Tri Dao committed
958
959
        word_embeddings, (0, 0, 0, vocab_size - word_embeddings.shape[0])
    )
Tri Dao's avatar
Tri Dao committed
960
    state_dict["lm_head.weight"] = state_dict["transformer.embeddings.word_embeddings.weight"]
961

Tri Dao's avatar
Tri Dao committed
962
963
    # LayerNorm
    def key_mapping_ln(key):
Tri Dao's avatar
Tri Dao committed
964
965
966
967
968
969
970
971
972
973
974
        key = re.sub(r"^transformer.final_layernorm.(weight|bias)", r"transformer.ln_f.\1", key)
        key = re.sub(
            r"^transformer.layers.(\d+).input_layernorm.(weight|bias)",
            r"transformer.layers.\1.norm1.\2",
            key,
        )
        key = re.sub(
            r"^transformer.layers.(\d+).post_attention_layernorm.(weight|bias)",
            r"transformer.layers.\1.norm2.\2",
            key,
        )
Tri Dao's avatar
Tri Dao committed
975
        return key
Tri Dao's avatar
Tri Dao committed
976

Tri Dao's avatar
Tri Dao committed
977
    state_dict = OrderedDict((key_mapping_ln(k), v) for k, v in state_dict.items())
978

Tri Dao's avatar
Tri Dao committed
979
980
    # MLP
    def key_mapping_mlp(key):
Tri Dao's avatar
Tri Dao committed
981
982
983
984
985
986
987
988
989
990
        key = re.sub(
            r"^transformer.layers.(\d+).mlp.dense_h_to_4h.(weight|bias)",
            r"transformer.layers.\1.mlp.fc1.\2",
            key,
        )
        key = re.sub(
            r"^transformer.layers.(\d+).mlp.dense_4h_to_h.(weight|bias)",
            r"transformer.layers.\1.mlp.fc2.\2",
            key,
        )
Tri Dao's avatar
Tri Dao committed
991
        return key
Tri Dao's avatar
Tri Dao committed
992

Tri Dao's avatar
Tri Dao committed
993
    state_dict = OrderedDict((key_mapping_mlp(k), v) for k, v in state_dict.items())
994

Tri Dao's avatar
Tri Dao committed
995
996
    # Attention
    def key_mapping_attn(key):
Tri Dao's avatar
Tri Dao committed
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
        key = re.sub(
            r"^transformer.layers.(\d+).self_attention.rotary_emb.inv_freq",
            r"transformer.layers.\1.mixer.rotary_emb.inv_freq",
            key,
        )
        key = re.sub(
            r"^transformer.layers.(\d+).self_attention.query_key_value.(weight|bias)",
            r"transformer.layers.\1.mixer.Wqkv.\2",
            key,
        )
        key = re.sub(
            r"^transformer.layers.(\d+).self_attention.dense.(weight|bias)",
            r"transformer.layers.\1.mixer.out_proj.\2",
            key,
        )
Tri Dao's avatar
Tri Dao committed
1012
        return key
Tri Dao's avatar
Tri Dao committed
1013

Tri Dao's avatar
Tri Dao committed
1014
1015
1016
1017
1018
    state_dict = OrderedDict((key_mapping_attn(k), v) for k, v in state_dict.items())
    # Megatron stores Wqkv as ((nheads 3 headdim), hidden_dim)
    # while we store Wqkv as ((3 nheads headdim), hidden_dim)
    headdim = config.hidden_size // config.num_attention_heads
    for d in range(config.num_hidden_layers):
Tri Dao's avatar
Tri Dao committed
1019
1020
1021
1022
1023
1024
        Wqkv = state_dict.pop(f"transformer.layers.{d}.mixer.Wqkv.weight")
        state_dict[f"transformer.layers.{d}.mixer.Wqkv.weight"] = rearrange(
            Wqkv,
            "(nheads three headdim) ... -> (three nheads headdim) ...",
            three=3,
            headdim=headdim,
Tri Dao's avatar
Tri Dao committed
1025
        )
Tri Dao's avatar
Tri Dao committed
1026
1027
1028
        bqkv = state_dict.pop(f"transformer.layers.{d}.mixer.Wqkv.bias")
        state_dict[f"transformer.layers.{d}.mixer.Wqkv.bias"] = rearrange(
            bqkv, "(nheads three headdim) -> (three nheads headdim)", three=3, headdim=headdim
Tri Dao's avatar
Tri Dao committed
1029
        )
1030
1031

    return state_dict