gpt.py 46.5 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)
Tri Dao's avatar
Tri Dao committed
81
82
    use_flash_attn = getattr(config, "use_flash_attn", False)
    fused_bias_fc = getattr(config, "fused_bias_fc", False)
83
    if not fused_bias_fc:
Tri Dao's avatar
Tri Dao committed
84
        assert process_group is None, "TensorParallel MHA requires fused_bias_fc"
85
    mha_cls = MHA if process_group is None else ParallelMHA
Tri Dao's avatar
Tri Dao committed
86
87
88
89
90
91
92
93
94
95
96
    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
97
    num_heads_kv = getattr(config, "n_head_kv", None)
Tri Dao's avatar
Tri Dao committed
98
99
100
101
102
103
104
105
106
107
108
109
110
111
    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,
112
        use_alibi=use_alibi,
Tri Dao's avatar
Tri Dao committed
113
114
115
116
117
        use_flash_attn=use_flash_attn,
        **serial_kwargs,
        **parallel_kwargs,
        **factory_kwargs,
    )
Tri Dao's avatar
Tri Dao committed
118
119
120
    return mixer_cls


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


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


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

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

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

Tri Dao's avatar
Tri Dao committed
376

Tri Dao's avatar
Tri Dao committed
377
# https://github.com/huggingface/transformers/blob/c28d04e9e252a1a099944e325685f14d242ecdcd/src/transformers/models/gpt2/modeling_gpt2.py#L454
Tri Dao's avatar
Tri Dao committed
378
379
380
381
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
382
    if isinstance(module, nn.Linear):
Tri Dao's avatar
Tri Dao committed
383
        nn.init.normal_(module.weight, std=initializer_range * mup_init_scale)
Tri Dao's avatar
Tri Dao committed
384
385
386
        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
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
        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
402
403
404
                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
405
406


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

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

Tri Dao's avatar
Tri Dao committed
457
        # We change the order of dropout, residual and layer norm:
Tri Dao's avatar
Tri Dao committed
458
        # Instead of LN -> Attn / MLP -> Dropout -> Add, we do:
Tri Dao's avatar
Tri Dao committed
459
460
461
        # 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
462
        # This is for performance reason: we can fuse dropout + add + layer_norm.
Tri Dao's avatar
Tri Dao committed
463
464
465
466
467
468
        self.layers = nn.ModuleList(
            [
                create_block(config, layer_idx=i, process_group=process_group, **factory_kwargs)
                for i in range(config.num_hidden_layers)
            ]
        )
469
470
471
472
        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
473

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

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

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

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

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


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

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

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

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

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

695

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

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

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

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

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

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

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


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

    This function is meant to be the "reverse" of shard_state_dict_tp.
817
818
819

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

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

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

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

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

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


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

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

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

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

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

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

963
964
965
966
    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
967
968
969
970
971
972
        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()

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

980
981
982
    state_dict = OrderedDict((key_mapping_attn(k), v) for k, v in state_dict.items())

    return state_dict
983
984


Tri Dao's avatar
Tri Dao committed
985
986
def remap_state_dict_megatron(state_dict, config):
    def key_mapping_transformer(key):
Tri Dao's avatar
Tri Dao committed
987
988
        key = re.sub(r"^language_model.encoder.", "transformer.", key)
        key = re.sub(r"^language_model.", "transformer.", key)
Tri Dao's avatar
Tri Dao committed
989
        return key
Tri Dao's avatar
Tri Dao committed
990

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

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

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

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

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

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

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

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

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

    return state_dict