gpt.py 46.6 KB
Newer Older
1
# Copyright (c) 2024, 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
Yuchao Dai's avatar
Yuchao Dai committed
9
from typing import Dict, List
Tri Dao's avatar
Tri Dao committed
10
11
12
13

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

Kevin Hu's avatar
Kevin Hu committed
17
from flash_attn.models.bigcode import remap_state_dict_hf_bigcode
Tri Dao's avatar
Tri Dao committed
18
19
20
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
21
from flash_attn.models.llama import remap_state_dict_hf_llama
Tri Dao's avatar
Tri Dao committed
22
from flash_attn.models.opt import remap_state_dict_hf_opt
Tri Dao's avatar
Tri Dao committed
23
from flash_attn.modules.block import Block, ParallelBlock
24
from flash_attn.modules.embedding import GPT2Embeddings, ParallelGPT2Embeddings
Tri Dao's avatar
Tri Dao committed
25
from flash_attn.modules.mha import MHA, ParallelMHA
Kevin Hu's avatar
Kevin Hu committed
26
27
28
29
30
31
32
33
from flash_attn.modules.mlp import (
    FusedMLP,
    GatedMlp,
    Mlp,
    ParallelFusedMLP,
    ParallelGatedMlp,
    ParallelMLP,
)
Tri Dao's avatar
Tri Dao committed
34
from flash_attn.ops.activations import sqrelu_fwd
Tri Dao's avatar
Tri Dao committed
35
36
37
38
39
40
from flash_attn.utils.distributed import (
    all_gather,
    all_gather_raw,
    get_dim_for_local_rank,
    sync_shared_params,
)
Tri Dao's avatar
Tri Dao committed
41
from flash_attn.utils.generation import GenerationMixin
Tri Dao's avatar
Tri Dao committed
42
from flash_attn.utils.pretrained import state_dict_from_pretrained
43
44
45
46
47

try:
    from flash_attn.ops.fused_dense import ColumnParallelLinear
except ImportError:
    ColumnParallelLinear = None
Tri Dao's avatar
Tri Dao committed
48
49

try:
50
    from flash_attn.ops.triton.mlp import FusedDenseSqreluDense
Tri Dao's avatar
Tri Dao committed
51
except ImportError:
52
    FusedDenseSqreluDense = None
Tri Dao's avatar
Tri Dao committed
53

Tri Dao's avatar
Tri Dao committed
54
try:
55
    from flash_attn.ops.triton.layer_norm import layer_norm_fn, RMSNorm
Tri Dao's avatar
Tri Dao committed
56
except ImportError:
57
    layer_norm_fn, RMSNorm = None, None
Tri Dao's avatar
Tri Dao committed
58

59
60
61
logger = logging.getLogger(__name__)


62
def create_mixer_cls(config, layer_idx=None, process_group=None, device=None, dtype=None):
Tri Dao's avatar
Tri Dao committed
63
64
    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
65
66
67
    attn_scale_power = 0.5 if not getattr(config, "mup_scale_qk_dot_by_d", False) else 1.0
    softmax_scale = 1.0 if not config.scale_attn_weights else (head_dim ** (-attn_scale_power))
    softmax_scale *= getattr(config, "mup_attn_multiplier", 1.0)
Tri Dao's avatar
Tri Dao committed
68
69
70
    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
71
    dwconv = getattr(config, "attn_dwconv", False)
72
    if dwconv:
Tri Dao's avatar
Tri Dao committed
73
74
75
76
77
78
79
        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)
80
    use_alibi = getattr(config, "use_alibi", False)
81
    window_size = getattr(config, "window_size", (-1, -1))
Tri Dao's avatar
Tri Dao committed
82
83
    use_flash_attn = getattr(config, "use_flash_attn", False)
    fused_bias_fc = getattr(config, "fused_bias_fc", False)
84
    if not fused_bias_fc:
Tri Dao's avatar
Tri Dao committed
85
        assert process_group is None, "TensorParallel MHA requires fused_bias_fc"
86
    mha_cls = MHA if process_group is None else ParallelMHA
Tri Dao's avatar
Tri Dao committed
87
88
89
90
91
92
93
94
95
96
97
    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
98
    num_heads_kv = getattr(config, "n_head_kv", None)
Tri Dao's avatar
Tri Dao committed
99
100
101
102
103
104
105
106
107
108
109
110
111
112
    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,
113
        use_alibi=use_alibi,
114
        window_size=window_size,
Tri Dao's avatar
Tri Dao committed
115
116
117
118
119
        use_flash_attn=use_flash_attn,
        **serial_kwargs,
        **parallel_kwargs,
        **factory_kwargs,
    )
Tri Dao's avatar
Tri Dao committed
120
121
122
    return mixer_cls


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


262
def create_block(config, layer_idx=None, process_group=None, device=None, dtype=None):
Tri Dao's avatar
Tri Dao committed
263
264
    factory_kwargs = {"device": device, "dtype": dtype}
    sequence_parallel = getattr(config, "sequence_parallel", True)
265
266
    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
267
268
269
270
271
272
    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
273
    # TD [2022-07-30]: Force residual in fp32, seems to make fp16 training more stable
Tri Dao's avatar
Tri Dao committed
274
    residual_in_fp32 = getattr(config, "residual_in_fp32", False)
Tri Dao's avatar
Tri Dao committed
275
    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
276
277
    prenorm = getattr(config, "prenorm", True)
    parallel_block = getattr(config, "parallel_block", False)
Tri Dao's avatar
Tri Dao committed
278
279
    if not parallel_block:
        block = Block(
Tri Dao's avatar
Tri Dao committed
280
281
282
283
284
285
286
287
            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
288
289
            residual_in_fp32=residual_in_fp32,
            sequence_parallel=sequence_parallel and process_group is not None,
Tri Dao's avatar
Tri Dao committed
290
            mark_shared_params=process_group is not None,
Tri Dao's avatar
Tri Dao committed
291
292
293
294
        )
    else:
        assert prenorm
        block = ParallelBlock(
Tri Dao's avatar
Tri Dao committed
295
296
297
298
299
300
301
302
            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
303
304
            residual_in_fp32=residual_in_fp32,
            sequence_parallel=sequence_parallel and process_group is not None,
Tri Dao's avatar
Tri Dao committed
305
            mark_shared_params=process_group is not None,
Tri Dao's avatar
Tri Dao committed
306
        )
Tri Dao's avatar
Tri Dao committed
307
308
309
310
    block.layer_idx = layer_idx
    return block


311
class GPTPreTrainedModel(nn.Module):
Tri Dao's avatar
Tri Dao committed
312
313
    """An abstract class to handle weights initialization and
    a simple interface for dowloading and loading pretrained models.
314
    """
Tri Dao's avatar
Tri Dao committed
315

316
317
318
319
320
321
322
323
    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
324
325
                )
            )
326
327
328
        self.config = config

    @classmethod
Tri Dao's avatar
Tri Dao committed
329
330
331
332
333
334
335
336
337
338
339
340
    def from_pretrained(
        cls,
        model_name,
        config,
        *args,
        strict=True,
        device=None,
        dtype=None,
        world_size=1,
        rank=0,
        **kwargs,
    ):
341
342
343
344
345
        """
        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.
346
        model = cls(config, *args, device=device, dtype=dtype, **kwargs)
347
348
        # 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
349
350
        state_dict = state_dict_from_pretrained(model_name, device="cpu", dtype=dtype)
        if model_name.startswith("gpt2"):
Tri Dao's avatar
Tri Dao committed
351
            state_dict = remap_state_dict_hf_gpt2(state_dict, config)
Tri Dao's avatar
Tri Dao committed
352
        elif model_name.startswith("facebook/opt"):
Tri Dao's avatar
Tri Dao committed
353
            state_dict = remap_state_dict_hf_opt(state_dict, config)
Tri Dao's avatar
Tri Dao committed
354
355
        elif model_name.startswith("EleutherAI/gpt-j-") or model_name.startswith(
            "togethercomputer/GPT-JT-"
356
        ):
Tri Dao's avatar
Tri Dao committed
357
            state_dict = remap_state_dict_hf_gptj(state_dict, config)
358
359
360
361
362
        elif (
            model_name.startswith("EleutherAI/gpt-neox-")
            or model_name.startswith("EleutherAI/pythia-")
            or model_name.startswith("togethercomputer/RedPajama-INCITE-")
        ):
Tri Dao's avatar
Tri Dao committed
363
            state_dict = remap_state_dict_hf_gpt_neox(state_dict, config)
Tri Dao's avatar
Tri Dao committed
364
        elif model_name.startswith("tiiuae/falcon-"):
Tri Dao's avatar
Tri Dao committed
365
            state_dict = remap_state_dict_hf_falcon(state_dict, config)
366
367
        elif model_name.startswith("meta-llama/Llama-"):
            state_dict = remap_state_dict_hf_llama(state_dict, config)
Kevin Hu's avatar
Kevin Hu committed
368
369
        elif model_name.startswith("bigcode/") or model_name.startswith("WizardLM/"):
            state_dict = remap_state_dict_hf_bigcode(state_dict, config)
Tri Dao's avatar
Tri Dao committed
370
        else:
Tri Dao's avatar
Tri Dao committed
371
            raise NotImplementedError(f"Model {model_name} not supported")
372
373
374
        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)
375
376
377
        logger.info(load_return)
        return model

Tri Dao's avatar
Tri Dao committed
378

Tri Dao's avatar
Tri Dao committed
379
# https://github.com/huggingface/transformers/blob/c28d04e9e252a1a099944e325685f14d242ecdcd/src/transformers/models/gpt2/modeling_gpt2.py#L454
Tri Dao's avatar
Tri Dao committed
380
381
382
383
def _init_weights(
    module, n_layer, initializer_range=0.02, mup_width_scale=1.0, rescale_prenorm_residual=True
):
    mup_init_scale = math.sqrt(mup_width_scale)
Tri Dao's avatar
Tri Dao committed
384
    if isinstance(module, nn.Linear):
Tri Dao's avatar
Tri Dao committed
385
        nn.init.normal_(module.weight, std=initializer_range * mup_init_scale)
Tri Dao's avatar
Tri Dao committed
386
387
388
        optim_cfg = getattr(module.weight, "_optim", {})
        optim_cfg.update({"lr_multiplier": mup_width_scale})
        setattr(module.weight, "_optim", optim_cfg)
Tri Dao's avatar
Tri Dao committed
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
        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
Tri Dao's avatar
Tri Dao committed
404
405
406
                nn.init.normal_(
                    p, mean=0.0, std=initializer_range * mup_init_scale / math.sqrt(2 * n_layer)
                )
Tri Dao's avatar
Tri Dao committed
407
408


409
class GPTModel(GPTPreTrainedModel):
410
    def __init__(self, config: GPT2Config, process_group=None, device=None, dtype=None):
411
        super().__init__(config)
Tri Dao's avatar
Tri Dao committed
412
        factory_kwargs = {"device": device, "dtype": dtype}
413
        self.process_group = process_group
Tri Dao's avatar
Tri Dao committed
414
415
416
417
418
419
        self.sequence_parallel = getattr(config, "sequence_parallel", True)
        assert config.activation_function in [
            "gelu",
            "gelu_new",
            "gelu_fast",
            "gelu_approx",
Kevin Hu's avatar
Kevin Hu committed
420
            "gelu_pytorch_tanh",
Tri Dao's avatar
Tri Dao committed
421
422
423
424
425
426
427
428
429
430
            "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
431
        self.embeddings_multiplier = getattr(config, "mup_embeddings_multiplier", 1.0)
Tri Dao's avatar
Tri Dao committed
432
        # TD [2022-07-30]: Force residual in fp32, seems to make fp16 training more stable
Tri Dao's avatar
Tri Dao committed
433
        self.residual_in_fp32 = getattr(config, "residual_in_fp32", False)
Tri Dao's avatar
Tri Dao committed
434
        # These 2 options are for OPT-350m
Tri Dao's avatar
Tri Dao committed
435
436
437
        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
438
        # For GPT-J, GPT-NeoX
Tri Dao's avatar
Tri Dao committed
439
        self.parallel_block = getattr(config, "parallel_block", False)
Tri Dao's avatar
Tri Dao committed
440

441
        if process_group is None:
Tri Dao's avatar
Tri Dao committed
442
            self.embeddings = GPT2Embeddings(
Tri Dao's avatar
Tri Dao committed
443
444
445
446
447
                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
448
            )
449
450
        else:
            self.embeddings = ParallelGPT2Embeddings(
Tri Dao's avatar
Tri Dao committed
451
452
453
454
455
456
                config.hidden_size,
                vocab_size,
                config.max_position_embeddings,
                process_group=process_group,
                sequence_parallel=self.sequence_parallel,
                **factory_kwargs,
457
            )
Tri Dao's avatar
Tri Dao committed
458

Tri Dao's avatar
Tri Dao committed
459
        # We change the order of dropout, residual and layer norm:
Tri Dao's avatar
Tri Dao committed
460
        # Instead of LN -> Attn / MLP -> Dropout -> Add, we do:
Tri Dao's avatar
Tri Dao committed
461
462
463
        # 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
464
        # This is for performance reason: we can fuse dropout + add + layer_norm.
Tri Dao's avatar
Tri Dao committed
465
466
467
468
469
470
        self.layers = nn.ModuleList(
            [
                create_block(config, layer_idx=i, process_group=process_group, **factory_kwargs)
                for i in range(config.num_hidden_layers)
            ]
        )
471
472
473
474
        rotary_emb_fraction = getattr(config, "rotary_emb_fraction", 0.0)
        if rotary_emb_fraction > 0.0:  # Tie all the RotaryEmbedding modules to share the same cos/sin cache
            for layer in self.layers[1:]:
                layer.mixer.rotary_emb = self.layers[0].mixer.rotary_emb
Tri Dao's avatar
Tri Dao committed
475

Tri Dao's avatar
Tri Dao committed
476
        self.fused_dropout_add_ln = getattr(config, "fused_dropout_add_ln", False)
477
        if self.fused_dropout_add_ln:
478
479
            if layer_norm_fn is None:
                raise ImportError("Triton is not installed")
Tri Dao's avatar
Tri Dao committed
480
481
        if self.prenorm:
            self.drop_f = nn.Dropout(config.resid_pdrop)
Tri Dao's avatar
Tri Dao committed
482
            norm_cls = nn.LayerNorm if not use_rms_norm else RMSNorm
Tri Dao's avatar
Tri Dao committed
483
484
485
            self.ln_f = norm_cls(
                config.hidden_size, eps=config.layer_norm_epsilon, **factory_kwargs
            )
486
        if process_group is not None:
Tri Dao's avatar
Tri Dao committed
487
            for p in self.ln_f.parameters():
488
489
490
491
492
                # 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
493

Tri Dao's avatar
Tri Dao committed
494
495
496
497
498
        self.apply(
            partial(
                _init_weights,
                n_layer=config.num_hidden_layers,
                initializer_range=config.initializer_range,
Tri Dao's avatar
Tri Dao committed
499
                mup_width_scale=getattr(config, "mup_width_scale", 1.0),
Tri Dao's avatar
Tri Dao committed
500
501
            )
        )
502
503
504
        self.tie_weights()

    def tie_weights(self):
505
        if self.process_group is not None:
506
            sync_shared_params(self, self.process_group)
Tri Dao's avatar
Tri Dao committed
507

508
    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
Tri Dao's avatar
Tri Dao committed
509
510
511
512
        return {
            i: layer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)
            for i, layer in enumerate(self.layers)
        }
513

Tri Dao's avatar
Tri Dao committed
514
    def forward(self, input_ids, position_ids=None, inference_params=None):
515
516
517
        # 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
518
519
520
521
522
        embedding_kwargs = (
            {"combine_batch_seqlen_dim": True}
            if self.process_group is not None and self.sequence_parallel
            else {}
        )
523
        hidden_states = self.embeddings(input_ids, position_ids=position_ids, **embedding_kwargs)
Tri Dao's avatar
Tri Dao committed
524
525
        if self.embeddings_multiplier != 1.0:
            hidden_states = hidden_states * self.embeddings_multiplier
Tri Dao's avatar
Tri Dao committed
526
527
        if self.parallel_block:
            hidden_states2 = None
Tri Dao's avatar
Tri Dao committed
528
        residual = None
Tri Dao's avatar
Tri Dao committed
529
530
531
532
533
        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
534
        if inference_params is not None:
Tri Dao's avatar
Tri Dao committed
535
            mixer_kwargs["inference_params"] = inference_params
Tri Dao's avatar
Tri Dao committed
536
        for layer in self.layers:
Tri Dao's avatar
Tri Dao committed
537
            if self.prenorm:
Tri Dao's avatar
Tri Dao committed
538
                if not self.parallel_block:
Tri Dao's avatar
Tri Dao committed
539
540
541
                    hidden_states, residual = layer(
                        hidden_states, residual, mixer_kwargs=mixer_kwargs
                    )
Tri Dao's avatar
Tri Dao committed
542
543
544
545
                else:
                    hidden_states, hidden_states2, residual = layer(
                        hidden_states, hidden_states2, residual, mixer_kwargs=mixer_kwargs
                    )
Tri Dao's avatar
Tri Dao committed
546
547
548
549
550
            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
551
552
553
554
                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
555
556
557
558
559
                    residual = (
                        (residual + dropped + dropped2)
                        if residual is not None
                        else dropped + dropped2
                    )
Tri Dao's avatar
Tri Dao committed
560
561
                hidden_states = self.ln_f(residual.to(dtype=self.ln_f.weight.dtype))
            else:
Tri Dao's avatar
Tri Dao committed
562
                # Set prenorm=False here since we don't need the residual
563
564
565
566
567
568
569
570
571
572
573
                hidden_states = layer_norm_fn(
                    hidden_states,
                    self.ln_f.weight,
                    self.ln_f.bias,
                    residual=residual,
                    x1=None if not self.parallel_block else hidden_states2,
                    eps=self.ln_f.eps,
                    dropout_p=self.drop_f.p if self.training else 0.0,
                    prenorm=False,
                    is_rms_norm=isinstance(self.ln_f, RMSNorm)
                )
Tri Dao's avatar
Tri Dao committed
574
575
576
        return hidden_states


Tri Dao's avatar
Tri Dao committed
577
class GPTLMHeadModel(GPTPreTrainedModel, GenerationMixin):
578
    def __init__(self, config: GPT2Config, process_group=None, device=None, dtype=None):
Tri Dao's avatar
Tri Dao committed
579
        factory_kwargs = {"device": device, "dtype": dtype}
580
        super().__init__(config)
581
582
        self.process_group = process_group
        self.transformer = GPTModel(config, process_group=process_group, **factory_kwargs)
Tri Dao's avatar
Tri Dao committed
583
584
585
586
587
588
        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
589
        # This option is for OPT-350m
Tri Dao's avatar
Tri Dao committed
590
        word_embed_proj_dim = getattr(config, "word_embed_proj_dim", None)
Tri Dao's avatar
Tri Dao committed
591
592
593
594
595
        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
Tri Dao's avatar
Tri Dao committed
596
597
598
        mup_width_scale = getattr(config, "mup_width_scale", 1.0)
        mup_output_multiplier = getattr(config, "mup_output_multiplier", 1.0)
        self.output_scale = mup_output_multiplier * mup_width_scale
599
        if process_group is None:
Tri Dao's avatar
Tri Dao committed
600
            self.lm_head = nn.Linear(embed_dim, vocab_size, bias=lm_head_bias, **factory_kwargs)
601
602
        else:
            if ColumnParallelLinear is None:
Tri Dao's avatar
Tri Dao committed
603
                raise ImportError("fused_dense_lib is not installed")
604
            self.lm_head = ColumnParallelLinear(
Tri Dao's avatar
Tri Dao committed
605
606
607
608
609
610
                embed_dim,
                vocab_size,
                process_group,
                bias=lm_head_bias,
                sequence_parallel=getattr(config, "sequence_parallel", True),
                **factory_kwargs,
611
            )
Tri Dao's avatar
Tri Dao committed
612
        self.norm_head = getattr(config, "norm_head", False)
Tri Dao's avatar
Tri Dao committed
613
        # Initialize weights and apply final processing
Tri Dao's avatar
Tri Dao committed
614
615
616
617
618
        self.apply(
            partial(
                _init_weights,
                n_layer=config.num_hidden_layers,
                initializer_range=config.initializer_range,
Tri Dao's avatar
Tri Dao committed
619
                mup_width_scale=mup_width_scale,
Tri Dao's avatar
Tri Dao committed
620
621
            )
        )
Tri Dao's avatar
Tri Dao committed
622
623
624
        self.tie_weights()

    def tie_weights(self):
Tri Dao's avatar
Tri Dao committed
625
626
        if self.tie_word_embeddings:
            self.lm_head.weight = self.transformer.embeddings.word_embeddings.weight
627
        if self.process_group is not None:
628
            sync_shared_params(self, self.process_group)
Tri Dao's avatar
Tri Dao committed
629

630
    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
Tri Dao's avatar
Tri Dao committed
631
632
633
        return self.transformer.allocate_inference_cache(
            batch_size, max_seqlen, dtype=dtype, **kwargs
        )
634

635
    def forward(self, input_ids, position_ids=None, inference_params=None, num_last_tokens=0):
Tri Dao's avatar
Tri Dao committed
636
        """
637
        input_ids: (batch, seqlen) int tensor
Tri Dao's avatar
Tri Dao committed
638
639
        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
640
        num_last_tokens: if > 0, only return the logits for the last n tokens
Tri Dao's avatar
Tri Dao committed
641
        """
Kevin Hu's avatar
Kevin Hu committed
642
643
644
        assert (
            input_ids.ndim == 2
        ), f"Expected `input_ids` to have shape [b, slen], but got shape {input_ids.shape}"
645
        b, slen = input_ids.shape
Tri Dao's avatar
Tri Dao committed
646
647
648
        hidden_states = self.transformer(
            input_ids, position_ids=position_ids, inference_params=inference_params
        )
Tri Dao's avatar
Tri Dao committed
649
650
        if inference_params is not None:
            assert hidden_states.ndim == 3, "sequence_parallel is not supported in generation mode"
651
652
        if num_last_tokens > 0:
            hidden_states = hidden_states[:, -num_last_tokens:]
Tri Dao's avatar
Tri Dao committed
653
654
        if self.project_out is not None:
            hidden_states = self.project_out(hidden_states)
Tri Dao's avatar
Tri Dao committed
655
656
        if self.output_scale != 1.0:
            hidden_states = hidden_states * self.output_scale
Tri Dao's avatar
Tri Dao committed
657
658
659
660
661
662
663
        if not self.norm_head:
            lm_logits = self.lm_head(hidden_states)
        else:
            lm_head_weight = F.normalize(self.lm_head.weight)
            if isinstance(self.lm_head, ColumnParallelLinear) and self.lm_head.sequence_parallel:
                hidden_states = all_gather(hidden_states, self.lm_head.process_group)
            lm_logits = F.linear(hidden_states, lm_head_weight, bias=self.lm_head.bias)
664
665
666
        # 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)
667
            lm_logits = rearrange(lm_logits, "(n b) ... d -> b ... (n d)", b=b)
Tri Dao's avatar
Tri Dao committed
668
        CausalLMOutput = namedtuple("CausalLMOutput", ["logits"])
Tri Dao's avatar
Tri Dao committed
669
        return CausalLMOutput(logits=lm_logits)
670

Tri Dao's avatar
Tri Dao committed
671
672
673
674
    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
675
        if "transformer.ln_0.weight" in state_dict:
Tri Dao's avatar
Tri Dao committed
676
            n_layers = len(self.transformer.layers)
Tri Dao's avatar
Tri Dao committed
677
678
679
680
            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
681
            for l in reversed(range(n_layers)):
Tri Dao's avatar
Tri Dao committed
682
683
684
685
                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
686
                if l > 0:
Tri Dao's avatar
Tri Dao committed
687
688
689
690
691
692
693
694
                    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
695
696
        return super().load_state_dict(state_dict, strict=strict)

697

Tri Dao's avatar
Tri Dao committed
698
699
700
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.
701
702

    This function modifies state_dict in place.
Tri Dao's avatar
Tri Dao committed
703
    """
Tri Dao's avatar
Tri Dao committed
704
705
    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
706
707
708
709
710
    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

711
712
713
714
715
716
    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
717
    def shard_first_dim(state_dict, key):
Tri Dao's avatar
Tri Dao committed
718
719
720
        if key in state_dict:
            x = state_dict[key]
            dim = x.shape[0] // world_size
Tri Dao's avatar
Tri Dao committed
721
            state_dict[key] = x[rank * dim : (rank + 1) * dim]
Tri Dao's avatar
Tri Dao committed
722

723
    def shard_last_dim(state_dict, key, multiple_of=1):
Tri Dao's avatar
Tri Dao committed
724
725
        if key in state_dict:
            x = state_dict[key]
726
727
728
729
730
731
            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
732

Tri Dao's avatar
Tri Dao committed
733
734
735
736
737
    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
738
                rearrange(x, "(two o) ... -> two o ...", two=2)[:, rank * dim : (rank + 1) * dim],
Tri Dao's avatar
Tri Dao committed
739
                "two o ... -> (two o) ...",
Tri Dao's avatar
Tri Dao committed
740
741
            )

Tri Dao's avatar
Tri Dao committed
742
    def shard_qkv_headdim(state_dict, key):
Tri Dao's avatar
Tri Dao committed
743
        if key in state_dict:
744
            n_head_each_rank = [
Tri Dao's avatar
Tri Dao committed
745
746
                get_dim_for_local_rank(n_head, world_size, local_rank)
                for local_rank in range(world_size)
747
748
            ]
            n_head_kv_each_rank = [
Tri Dao's avatar
Tri Dao committed
749
750
                get_dim_for_local_rank(n_head_kv, world_size, local_rank)
                for local_rank in range(world_size)
751
752
753
754
755
756
757
758
            ]

            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
759
            if n_head_kv == n_head:
Tri Dao's avatar
Tri Dao committed
760
761
                x = rearrange(state_dict[key], "(three d) ... -> three d ...", three=3)
                state_dict[key] = rearrange(
Tri Dao's avatar
Tri Dao committed
762
763
                    x[:, beg_n_head * head_dim : end_n_head * head_dim],
                    "three d ... -> (three d) ...",
Tri Dao's avatar
Tri Dao committed
764
                )
Tri Dao's avatar
Tri Dao committed
765
            else:
Tri Dao's avatar
Tri Dao committed
766
767
768
769
770
771
772
773
                x = rearrange(
                    state_dict[key],
                    "(nheadqkv headdim) ... -> nheadqkv headdim ...",
                    nheadqkv=n_head + 2 * n_head_kv,
                )
                state_dict[key] = rearrange(
                    torch.cat(
                        [
774
                            x[beg_n_head:end_n_head],
Tri Dao's avatar
Tri Dao committed
775
776
777
778
779
780
781
782
                            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
783
784
785
786
787
788
789
790
791
792
793
                        ],
                        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
794
    for i in range(config.num_hidden_layers):
Tri Dao's avatar
Tri Dao committed
795
796
        shard_qkv_headdim(state_dict, f"transformer.layers.{i}.mixer.Wqkv.weight")
        shard_qkv_headdim(state_dict, f"transformer.layers.{i}.mixer.Wqkv.bias")
797
798
799
        shard_last_dim(
            state_dict, f"transformer.layers.{i}.mixer.out_proj.weight", multiple_of=head_dim
        )
Tri Dao's avatar
Tri Dao committed
800
        if rank != 0:
Tri Dao's avatar
Tri Dao committed
801
            state_dict.pop(f"transformer.layers.{i}.mixer.out_proj.bias", None)
Tri Dao's avatar
Tri Dao committed
802
        if config.activation_function in ["glu", "swiglu", "geglu"]:
Tri Dao's avatar
Tri Dao committed
803
804
            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
805
        else:
Tri Dao's avatar
Tri Dao committed
806
807
808
            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
809
        if rank != 0:
Tri Dao's avatar
Tri Dao committed
810
            state_dict.pop(f"transformer.layers.{i}.mlp.fc2.bias", None)
Tri Dao's avatar
Tri Dao committed
811
812
813
    return state_dict


Yuchao Dai's avatar
Yuchao Dai committed
814
def combine_state_dicts_tp(state_dicts: List[Dict[str, torch.Tensor]], config: GPT2Config):
815
816
    """Convert the list of sharded state_dict of a GPT model with tensor parallel to
    the state_dict of a standard GPT model.
817
818

    This function is meant to be the "reverse" of shard_state_dict_tp.
819
820
821

    Precondition:
        - state_dicts should be ordered in the same way as the shards were created.
Tri Dao's avatar
Tri Dao committed
822
823
824
    """
    world_size = len(state_dicts)
    keys = state_dicts[0].keys()
Tri Dao's avatar
Tri Dao committed
825
826
    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
827
828
829
830
    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
831
832
    assert config.hidden_size % config.n_head == 0
    headdim = config.hidden_size // config.n_head
Tri Dao's avatar
Tri Dao committed
833

Tri Dao's avatar
Tri Dao committed
834
    # Sometimes the word embeddings are sharded on the 0th dim, sometimes on the 1st dim.
Tri Dao's avatar
Tri Dao committed
835
836
    # vocab_size // world_size coordinates are nonzero.
    def combine_word_embeddings(state_dicts, state_dict, key):
Tri Dao's avatar
Tri Dao committed
837
838
        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
839
840

    def combine_dim(state_dicts, state_dict, key, dim=-1):
Tri Dao's avatar
Tri Dao committed
841
842
        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
843
844

    def combine_qkv_headdim(state_dicts, state_dict, key):
Tri Dao's avatar
Tri Dao committed
845
        n_head = config.n_head
Tri Dao's avatar
Tri Dao committed
846
        n_head_kv = getattr(config, "n_head_kv", n_head)
Tri Dao's avatar
Tri Dao committed
847
        if key in state_dict:
Tri Dao's avatar
Tri Dao committed
848
            if n_head_kv == n_head:
Tri Dao's avatar
Tri Dao committed
849
850
851
852
                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
853
            else:
854
855
856
857
858
859
860
861
                n_head_each_rank = [
                    get_dim_for_local_rank(n_head, world_size, local_rank)
                    for local_rank in range(world_size)
                ]
                n_head_kv_each_rank = [
                    get_dim_for_local_rank(n_head_kv, world_size, local_rank)
                    for local_rank in range(world_size)
                ]
862
863
864
865
866
867
868
                xs = [
                    rearrange(
                        s[key],
                        "(nheadqkv headdim) ... -> nheadqkv headdim ...",
                        nheadqkv=rank_n_head + 2 * rank_n_head_kv,
                        headdim=headdim,
                    )
Kevin Hu's avatar
Kevin Hu committed
869
870
871
                    for s, rank_n_head, rank_n_head_kv in zip(
                        state_dicts, n_head_each_rank, n_head_kv_each_rank
                    )
872
                ]
Kevin Hu's avatar
Kevin Hu committed
873
                wq = torch.cat([x[: n_head_each_rank[rank]] for rank, x in enumerate(xs)], dim=0)
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
                wk = torch.cat(
                    [
                        x[
                            n_head_each_rank[rank] : n_head_each_rank[rank]
                            + n_head_kv_each_rank[rank]
                        ]
                        for rank, x in enumerate(xs)
                    ],
                    dim=0,
                )
                wv = torch.cat(
                    [
                        x[n_head_each_rank[rank] + n_head_kv_each_rank[rank] :]
                        for rank, x in enumerate(xs)
                    ],
                    dim=0,
                )
                wqkv = torch.cat(
                    [wq, wk, wv],
                    dim=0,
                )
Tri Dao's avatar
Tri Dao committed
895
                state_dict[key] = rearrange(
896
                    wqkv,
Tri Dao's avatar
Tri Dao committed
897
898
                    "nheadqkv headdim ... -> (nheadqkv headdim) ...",
                )
Tri Dao's avatar
Tri Dao committed
899
900
901

    def combine_gated_mlp(state_dicts, state_dict, key):
        if key in state_dict:
Tri Dao's avatar
Tri Dao committed
902
903
            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
904
905

    state_dict = state_dicts[0].copy()  # don't modify state_dict[0] inplace
Tri Dao's avatar
Tri Dao committed
906
907
908
909
910
911
912
913
914
915
916
917
918
919
    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
920
    for i in range(config.num_hidden_layers):
Tri Dao's avatar
Tri Dao committed
921
922
923
924
925
926
        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
927
928
929
930
    return state_dict


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

935
    state_dict = OrderedDict((key_mapping_pos_emb(k), v) for k, v in state_dict.items())
Tri Dao's avatar
Tri Dao committed
936
    word_embeddings = state_dict.pop("wte.weight")
937
    # It's possible that vocab_size is padded to be a multiple of 8, for example.
Tri Dao's avatar
Tri Dao committed
938
939
940
    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(
941
        word_embeddings, (0, 0, 0, vocab_size - word_embeddings.shape[0])
942
    )
Tri Dao's avatar
Tri Dao committed
943
    state_dict["lm_head.weight"] = state_dict["transformer.embeddings.word_embeddings.weight"]
944
945

    # LayerNorm
Tri Dao's avatar
Tri Dao committed
946
    def key_mapping_ln(key):
Tri Dao's avatar
Tri Dao committed
947
948
        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
949
        return key
Tri Dao's avatar
Tri Dao committed
950

Tri Dao's avatar
Tri Dao committed
951
    state_dict = OrderedDict((key_mapping_ln(k), v) for k, v in state_dict.items())
952
953
954

    # MLP
    for d in range(config.num_hidden_layers):
Tri Dao's avatar
Tri Dao committed
955
956
957
958
959
        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()

960
    def key_mapping_mlp(key):
Tri Dao's avatar
Tri Dao committed
961
962
        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)
963
        return key
Tri Dao's avatar
Tri Dao committed
964

965
966
967
968
    state_dict = OrderedDict((key_mapping_mlp(k), v) for k, v in state_dict.items())

    # Attention
    for d in range(config.num_hidden_layers):
969
        state_dict.pop(f"h.{d}.attn.bias", None)  # We don't store this bias
Tri Dao's avatar
Tri Dao committed
970
971
972
973
974
        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()

975
    def key_mapping_attn(key):
Tri Dao's avatar
Tri Dao committed
976
977
978
979
        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
        )
980
        return key
Tri Dao's avatar
Tri Dao committed
981

982
983
984
    state_dict = OrderedDict((key_mapping_attn(k), v) for k, v in state_dict.items())

    return state_dict
985
986


Tri Dao's avatar
Tri Dao committed
987
988
def remap_state_dict_megatron(state_dict, config):
    def key_mapping_transformer(key):
Tri Dao's avatar
Tri Dao committed
989
990
        key = re.sub(r"^language_model.encoder.", "transformer.", key)
        key = re.sub(r"^language_model.", "transformer.", 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_transformer(k), v) for k, v in state_dict.items())
994

Tri Dao's avatar
Tri Dao committed
995
996
    # Word embedding and position embedding
    def key_mapping_pos_emb(key):
Tri Dao's avatar
Tri Dao committed
997
998
        return re.sub(r"^wpe.", "transformer.embeddings.position_embeddings.", key)

Tri Dao's avatar
Tri Dao committed
999
    state_dict = OrderedDict((key_mapping_pos_emb(k), v) for k, v in state_dict.items())
Tri Dao's avatar
Tri Dao committed
1000
    word_embeddings = state_dict.pop("transformer.embedding.word_embeddings.weight")
Tri Dao's avatar
Tri Dao committed
1001
    # It's possible that vocab_size is padded to be a multiple of 8, for example.
Tri Dao's avatar
Tri Dao committed
1002
1003
1004
1005
1006
    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
1007
1008
        word_embeddings, (0, 0, 0, vocab_size - word_embeddings.shape[0])
    )
Tri Dao's avatar
Tri Dao committed
1009
    state_dict["lm_head.weight"] = state_dict["transformer.embeddings.word_embeddings.weight"]
1010

Tri Dao's avatar
Tri Dao committed
1011
1012
    # LayerNorm
    def key_mapping_ln(key):
Tri Dao's avatar
Tri Dao committed
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
        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
1024
        return key
Tri Dao's avatar
Tri Dao committed
1025

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

Tri Dao's avatar
Tri Dao committed
1028
1029
    # MLP
    def key_mapping_mlp(key):
Tri Dao's avatar
Tri Dao committed
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
        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
1040
        return key
Tri Dao's avatar
Tri Dao committed
1041

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

Tri Dao's avatar
Tri Dao committed
1044
1045
    # Attention
    def key_mapping_attn(key):
Tri Dao's avatar
Tri Dao committed
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
        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
1061
        return key
Tri Dao's avatar
Tri Dao committed
1062

Tri Dao's avatar
Tri Dao committed
1063
1064
1065
1066
1067
    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
1068
1069
1070
1071
1072
1073
        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
1074
        )
Tri Dao's avatar
Tri Dao committed
1075
1076
1077
        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
1078
        )
1079
1080

    return state_dict