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

3
import logging
Tri Dao's avatar
Tri Dao committed
4
import math
5
import re
Tri Dao's avatar
Tri Dao committed
6
from collections import OrderedDict, namedtuple
Tri Dao's avatar
Tri Dao committed
7
from collections.abc import Sequence
Tri Dao's avatar
Tri Dao committed
8
from functools import partial
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
50
51
52
53

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

54
try:
Kevin Hu's avatar
Kevin Hu committed
55
    from flash_attn.ops.layer_norm import dropout_add_layer_norm_parallel_residual
56
57
58
except ImportError:
    dropout_add_layer_norm_parallel_residual = None

Tri Dao's avatar
Tri Dao committed
59
60
61
try:
    from flash_attn.ops.rms_norm import RMSNorm, dropout_add_rms_norm
except ImportError:
62
    RMSNorm, dropout_add_rms_norm = None, None
Tri Dao's avatar
Tri Dao committed
63
64
65
66
67
68

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

Tri Dao's avatar
Tri Dao committed
69
try:
Tri Dao's avatar
Tri Dao committed
70
    from flash_attn.ops.triton.mlp import FusedDenseSqreluDense
Tri Dao's avatar
Tri Dao committed
71
72
73
except ImportError:
    FusedDenseSqreluDense = None

74
75
76
logger = logging.getLogger(__name__)


77
def create_mixer_cls(config, layer_idx=None, process_group=None, device=None, dtype=None):
Tri Dao's avatar
Tri Dao committed
78
79
    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
80
81
82
83
    softmax_scale = 1.0 if not config.scale_attn_weights else head_dim ** (-0.5)
    if config.scale_attn_by_inverse_layer_idx:
        assert layer_idx is not None
        softmax_scale /= float(layer_idx + 1)
Tri Dao's avatar
Tri Dao committed
84
    dwconv = getattr(config, "attn_dwconv", False)
85
    if dwconv:
Tri Dao's avatar
Tri Dao committed
86
87
88
89
90
91
92
        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)
93
    use_alibi = getattr(config, "use_alibi", False)
Tri Dao's avatar
Tri Dao committed
94
95
    use_flash_attn = getattr(config, "use_flash_attn", False)
    fused_bias_fc = getattr(config, "fused_bias_fc", False)
96
    if not fused_bias_fc:
Tri Dao's avatar
Tri Dao committed
97
        assert process_group is None, "TensorParallel MHA requires fused_bias_fc"
98
    mha_cls = MHA if process_group is None else ParallelMHA
Tri Dao's avatar
Tri Dao committed
99
100
101
102
103
104
105
106
107
108
109
    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
110
    num_heads_kv = getattr(config, "n_head_kv", None)
Tri Dao's avatar
Tri Dao committed
111
112
113
114
115
116
117
118
119
120
121
122
123
124
    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,
125
        use_alibi=use_alibi,
Tri Dao's avatar
Tri Dao committed
126
127
128
129
130
        use_flash_attn=use_flash_attn,
        **serial_kwargs,
        **parallel_kwargs,
        **factory_kwargs,
    )
Tri Dao's avatar
Tri Dao committed
131
132
133
    return mixer_cls


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


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


320
class GPTPreTrainedModel(nn.Module):
Tri Dao's avatar
Tri Dao committed
321
322
    """An abstract class to handle weights initialization and
    a simple interface for dowloading and loading pretrained models.
323
    """
Tri Dao's avatar
Tri Dao committed
324

325
326
327
328
329
330
331
332
    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
333
334
                )
            )
335
336
337
        self.config = config

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

Tri Dao's avatar
Tri Dao committed
387

Tri Dao's avatar
Tri Dao committed
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
# https://github.com/huggingface/transformers/blob/c28d04e9e252a1a099944e325685f14d242ecdcd/src/transformers/models/gpt2/modeling_gpt2.py#L454
def _init_weights(module, n_layer, initializer_range=0.02, rescale_prenorm_residual=True):
    if isinstance(module, nn.Linear):
        nn.init.normal_(module.weight, std=initializer_range)
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Embedding):
        nn.init.normal_(module.weight, std=initializer_range)

    if rescale_prenorm_residual:
        # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
        #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
        #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
        #   >   -- GPT-2 :: https://openai.com/blog/better-language-models/
        #
        # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
        for name, p in module.named_parameters():
            if name in ["out_proj.weight", "fc2.weight"]:
                # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
                nn.init.normal_(p, mean=0.0, std=initializer_range / math.sqrt(2 * n_layer))


410
class GPTModel(GPTPreTrainedModel):
411
    def __init__(self, config: GPT2Config, process_group=None, device=None, dtype=None):
412
        super().__init__(config)
Tri Dao's avatar
Tri Dao committed
413
        factory_kwargs = {"device": device, "dtype": dtype}
414
        self.process_group = process_group
Tri Dao's avatar
Tri Dao committed
415
416
417
418
419
420
        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
421
            "gelu_pytorch_tanh",
Tri Dao's avatar
Tri Dao committed
422
423
424
425
426
427
428
429
430
431
            "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
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)
            ]
        )
Tri Dao's avatar
Tri Dao committed
471

Tri Dao's avatar
Tri Dao committed
472
        self.fused_dropout_add_ln = getattr(config, "fused_dropout_add_ln", False)
473
        if self.fused_dropout_add_ln:
Tri Dao's avatar
Tri Dao committed
474
475
476
477
            if (not self.parallel_block and dropout_add_layer_norm is None) or (
                self.parallel_block and dropout_add_layer_norm_parallel_residual is None
            ):
                raise ImportError("dropout_layer_norm is not installed")
Tri Dao's avatar
Tri Dao committed
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
497
498
        self.apply(
            partial(
                _init_weights,
                n_layer=config.num_hidden_layers,
                initializer_range=config.initializer_range,
            )
        )
499
500
501
        self.tie_weights()

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

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

Tri Dao's avatar
Tri Dao committed
511
    def forward(self, input_ids, position_ids=None, inference_params=None):
512
513
514
        # 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
515
516
517
518
519
        embedding_kwargs = (
            {"combine_batch_seqlen_dim": True}
            if self.process_group is not None and self.sequence_parallel
            else {}
        )
520
        hidden_states = self.embeddings(input_ids, position_ids=position_ids, **embedding_kwargs)
Tri Dao's avatar
Tri Dao committed
521
522
        if self.parallel_block:
            hidden_states2 = None
Tri Dao's avatar
Tri Dao committed
523
        residual = None
Tri Dao's avatar
Tri Dao committed
524
525
526
527
528
        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
529
        if inference_params is not None:
Tri Dao's avatar
Tri Dao committed
530
            mixer_kwargs["inference_params"] = inference_params
Tri Dao's avatar
Tri Dao committed
531
        for layer in self.layers:
Tri Dao's avatar
Tri Dao committed
532
            if self.prenorm:
Tri Dao's avatar
Tri Dao committed
533
                if not self.parallel_block:
Tri Dao's avatar
Tri Dao committed
534
535
536
                    hidden_states, residual = layer(
                        hidden_states, residual, mixer_kwargs=mixer_kwargs
                    )
Tri Dao's avatar
Tri Dao committed
537
538
539
540
                else:
                    hidden_states, hidden_states2, residual = layer(
                        hidden_states, hidden_states2, residual, mixer_kwargs=mixer_kwargs
                    )
Tri Dao's avatar
Tri Dao committed
541
542
543
544
545
            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
546
547
548
549
                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
550
551
552
553
554
                    residual = (
                        (residual + dropped + dropped2)
                        if residual is not None
                        else dropped + dropped2
                    )
Tri Dao's avatar
Tri Dao committed
555
556
                hidden_states = self.ln_f(residual.to(dtype=self.ln_f.weight.dtype))
            else:
Tri Dao's avatar
Tri Dao committed
557
                # Set prenorm=False here since we don't need the residual
558
                if not self.parallel_block:
Tri Dao's avatar
Tri Dao committed
559
560
561
562
563
                    fused_add_norm_fn = (
                        dropout_add_rms_norm
                        if isinstance(self.ln_f, RMSNorm)
                        else dropout_add_layer_norm
                    )
564
                    hidden_states = fused_add_norm_fn(
Tri Dao's avatar
Tri Dao committed
565
566
567
568
569
570
571
572
                        hidden_states,
                        residual,
                        self.ln_f.weight,
                        self.ln_f.bias,
                        self.drop_f.p if self.training else 0.0,
                        self.ln_f.eps,
                        prenorm=False,
                        residual_in_fp32=self.residual_in_fp32,
573
574
                    )
                else:
Tri Dao's avatar
Tri Dao committed
575
576
577
578
579
                    fused_add_norm_fn = (
                        dropout_add_rms_norm_parallel_residual
                        if isinstance(self.ln_f, RMSNorm)
                        else dropout_add_layer_norm_parallel_residual
                    )
580
                    hidden_states, _ = fused_add_norm_fn(
Tri Dao's avatar
Tri Dao committed
581
582
583
584
585
586
587
588
589
590
591
                        hidden_states,
                        hidden_states2,
                        residual,
                        self.ln_f.weight,
                        self.ln_f.bias,
                        None,
                        None,
                        self.drop_f.p if self.training else 0.0,
                        self.ln_f.eps,
                        prenorm=False,
                        residual_in_fp32=self.residual_in_fp32,
592
                    )
Tri Dao's avatar
Tri Dao committed
593
594
595
        return hidden_states


Tri Dao's avatar
Tri Dao committed
596
class GPTLMHeadModel(GPTPreTrainedModel, GenerationMixin):
597
    def __init__(self, config: GPT2Config, process_group=None, device=None, dtype=None):
Tri Dao's avatar
Tri Dao committed
598
        factory_kwargs = {"device": device, "dtype": dtype}
599
        super().__init__(config)
600
601
        self.process_group = process_group
        self.transformer = GPTModel(config, process_group=process_group, **factory_kwargs)
Tri Dao's avatar
Tri Dao committed
602
603
604
605
606
607
        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
608
        # This option is for OPT-350m
Tri Dao's avatar
Tri Dao committed
609
        word_embed_proj_dim = getattr(config, "word_embed_proj_dim", None)
Tri Dao's avatar
Tri Dao committed
610
611
612
613
614
        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
615
        if process_group is None:
Tri Dao's avatar
Tri Dao committed
616
            self.lm_head = nn.Linear(embed_dim, vocab_size, bias=lm_head_bias, **factory_kwargs)
617
618
        else:
            if ColumnParallelLinear is None:
Tri Dao's avatar
Tri Dao committed
619
                raise ImportError("fused_dense_lib is not installed")
620
            self.lm_head = ColumnParallelLinear(
Tri Dao's avatar
Tri Dao committed
621
622
623
624
625
626
                embed_dim,
                vocab_size,
                process_group,
                bias=lm_head_bias,
                sequence_parallel=getattr(config, "sequence_parallel", True),
                **factory_kwargs,
627
            )
Tri Dao's avatar
Tri Dao committed
628
        self.norm_head = getattr(config, "norm_head", False)
Tri Dao's avatar
Tri Dao committed
629
        # Initialize weights and apply final processing
Tri Dao's avatar
Tri Dao committed
630
631
632
633
634
635
636
        self.apply(
            partial(
                _init_weights,
                n_layer=config.num_hidden_layers,
                initializer_range=config.initializer_range,
            )
        )
Tri Dao's avatar
Tri Dao committed
637
638
639
        self.tie_weights()

    def tie_weights(self):
Tri Dao's avatar
Tri Dao committed
640
641
        if self.tie_word_embeddings:
            self.lm_head.weight = self.transformer.embeddings.word_embeddings.weight
642
        if self.process_group is not None:
643
            sync_shared_params(self, self.process_group)
Tri Dao's avatar
Tri Dao committed
644

645
    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
Tri Dao's avatar
Tri Dao committed
646
647
648
        return self.transformer.allocate_inference_cache(
            batch_size, max_seqlen, dtype=dtype, **kwargs
        )
649

650
    def forward(self, input_ids, position_ids=None, inference_params=None, num_last_tokens=0):
Tri Dao's avatar
Tri Dao committed
651
        """
652
        input_ids: (batch, seqlen) int tensor
Tri Dao's avatar
Tri Dao committed
653
654
        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
655
        num_last_tokens: if > 0, only return the logits for the last n tokens
Tri Dao's avatar
Tri Dao committed
656
        """
Kevin Hu's avatar
Kevin Hu committed
657
658
659
        assert (
            input_ids.ndim == 2
        ), f"Expected `input_ids` to have shape [b, slen], but got shape {input_ids.shape}"
660
        b, slen = input_ids.shape
Tri Dao's avatar
Tri Dao committed
661
662
663
        hidden_states = self.transformer(
            input_ids, position_ids=position_ids, inference_params=inference_params
        )
Tri Dao's avatar
Tri Dao committed
664
665
        if inference_params is not None:
            assert hidden_states.ndim == 3, "sequence_parallel is not supported in generation mode"
666
667
        if num_last_tokens > 0:
            hidden_states = hidden_states[:, -num_last_tokens:]
Tri Dao's avatar
Tri Dao committed
668
669
        if self.project_out is not None:
            hidden_states = self.project_out(hidden_states)
Tri Dao's avatar
Tri Dao committed
670
671
672
673
674
675
676
        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)
677
678
679
        # 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)
680
            lm_logits = rearrange(lm_logits, "(n b) ... d -> b ... (n d)", b=b)
Tri Dao's avatar
Tri Dao committed
681
        CausalLMOutput = namedtuple("CausalLMOutput", ["logits"])
Tri Dao's avatar
Tri Dao committed
682
        return CausalLMOutput(logits=lm_logits)
683

Tri Dao's avatar
Tri Dao committed
684
685
686
687
    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
688
        if "transformer.ln_0.weight" in state_dict:
Tri Dao's avatar
Tri Dao committed
689
            n_layers = len(self.transformer.layers)
Tri Dao's avatar
Tri Dao committed
690
691
692
693
            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
694
            for l in reversed(range(n_layers)):
Tri Dao's avatar
Tri Dao committed
695
696
697
698
                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
699
                if l > 0:
Tri Dao's avatar
Tri Dao committed
700
701
702
703
704
705
706
707
                    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
708
709
        return super().load_state_dict(state_dict, strict=strict)

710

Tri Dao's avatar
Tri Dao committed
711
712
713
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.
714
715

    This function modifies state_dict in place.
Tri Dao's avatar
Tri Dao committed
716
    """
Tri Dao's avatar
Tri Dao committed
717
718
    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
719
720
721
722
723
    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

724
725
726
727
728
729
    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
730
    def shard_first_dim(state_dict, key):
Tri Dao's avatar
Tri Dao committed
731
732
733
        if key in state_dict:
            x = state_dict[key]
            dim = x.shape[0] // world_size
Tri Dao's avatar
Tri Dao committed
734
            state_dict[key] = x[rank * dim : (rank + 1) * dim]
Tri Dao's avatar
Tri Dao committed
735

736
    def shard_last_dim(state_dict, key, multiple_of=1):
Tri Dao's avatar
Tri Dao committed
737
738
        if key in state_dict:
            x = state_dict[key]
739
740
741
742
743
744
            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
745

Tri Dao's avatar
Tri Dao committed
746
747
748
749
750
    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
751
                rearrange(x, "(two o) ... -> two o ...", two=2)[:, rank * dim : (rank + 1) * dim],
Tri Dao's avatar
Tri Dao committed
752
                "two o ... -> (two o) ...",
Tri Dao's avatar
Tri Dao committed
753
754
            )

Tri Dao's avatar
Tri Dao committed
755
    def shard_qkv_headdim(state_dict, key):
Tri Dao's avatar
Tri Dao committed
756
        if key in state_dict:
757
            n_head_each_rank = [
Tri Dao's avatar
Tri Dao committed
758
759
                get_dim_for_local_rank(n_head, world_size, local_rank)
                for local_rank in range(world_size)
760
761
            ]
            n_head_kv_each_rank = [
Tri Dao's avatar
Tri Dao committed
762
763
                get_dim_for_local_rank(n_head_kv, world_size, local_rank)
                for local_rank in range(world_size)
764
765
766
767
768
769
770
771
            ]

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


Yuchao Dai's avatar
Yuchao Dai committed
827
def combine_state_dicts_tp(state_dicts: List[Dict[str, torch.Tensor]], config: GPT2Config):
828
829
    """Convert the list of sharded state_dict of a GPT model with tensor parallel to
    the state_dict of a standard GPT model.
830
831

    This function is meant to be the "reverse" of shard_state_dict_tp.
832
833
834

    Precondition:
        - state_dicts should be ordered in the same way as the shards were created.
Tri Dao's avatar
Tri Dao committed
835
836
837
    """
    world_size = len(state_dicts)
    keys = state_dicts[0].keys()
Tri Dao's avatar
Tri Dao committed
838
839
    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
840
841
842
843
    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
844
845
    assert config.hidden_size % config.n_head == 0
    headdim = config.hidden_size // config.n_head
Tri Dao's avatar
Tri Dao committed
846

Tri Dao's avatar
Tri Dao committed
847
    # Sometimes the word embeddings are sharded on the 0th dim, sometimes on the 1st dim.
Tri Dao's avatar
Tri Dao committed
848
849
    # vocab_size // world_size coordinates are nonzero.
    def combine_word_embeddings(state_dicts, state_dict, key):
Tri Dao's avatar
Tri Dao committed
850
851
        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
852
853

    def combine_dim(state_dicts, state_dict, key, dim=-1):
Tri Dao's avatar
Tri Dao committed
854
855
        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
856
857

    def combine_qkv_headdim(state_dicts, state_dict, key):
Tri Dao's avatar
Tri Dao committed
858
        n_head = config.n_head
Tri Dao's avatar
Tri Dao committed
859
        n_head_kv = getattr(config, "n_head_kv", n_head)
Tri Dao's avatar
Tri Dao committed
860
        if key in state_dict:
Tri Dao's avatar
Tri Dao committed
861
            if n_head_kv == n_head:
Tri Dao's avatar
Tri Dao committed
862
863
864
865
                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
866
            else:
867
868
869
870
871
872
873
874
                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)
                ]
875
876
877
878
879
880
881
                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
882
883
884
                    for s, rank_n_head, rank_n_head_kv in zip(
                        state_dicts, n_head_each_rank, n_head_kv_each_rank
                    )
885
                ]
Kevin Hu's avatar
Kevin Hu committed
886
                wq = torch.cat([x[: n_head_each_rank[rank]] for rank, x in enumerate(xs)], dim=0)
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
                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
908
                state_dict[key] = rearrange(
909
                    wqkv,
Tri Dao's avatar
Tri Dao committed
910
911
                    "nheadqkv headdim ... -> (nheadqkv headdim) ...",
                )
Tri Dao's avatar
Tri Dao committed
912
913
914

    def combine_gated_mlp(state_dicts, state_dict, key):
        if key in state_dict:
Tri Dao's avatar
Tri Dao committed
915
916
            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
917
918

    state_dict = state_dicts[0].copy()  # don't modify state_dict[0] inplace
Tri Dao's avatar
Tri Dao committed
919
920
921
922
923
924
925
926
927
928
929
930
931
932
    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
933
    for i in range(config.num_hidden_layers):
Tri Dao's avatar
Tri Dao committed
934
935
936
937
938
939
        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
940
941
942
943
    return state_dict


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

948
    state_dict = OrderedDict((key_mapping_pos_emb(k), v) for k, v in state_dict.items())
Tri Dao's avatar
Tri Dao committed
949
    word_embeddings = state_dict.pop("wte.weight")
950
    # It's possible that vocab_size is padded to be a multiple of 8, for example.
Tri Dao's avatar
Tri Dao committed
951
952
953
    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(
954
        word_embeddings, (0, 0, 0, vocab_size - word_embeddings.shape[0])
955
    )
Tri Dao's avatar
Tri Dao committed
956
    state_dict["lm_head.weight"] = state_dict["transformer.embeddings.word_embeddings.weight"]
957
958

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

Tri Dao's avatar
Tri Dao committed
964
    state_dict = OrderedDict((key_mapping_ln(k), v) for k, v in state_dict.items())
965
966
967

    # MLP
    for d in range(config.num_hidden_layers):
Tri Dao's avatar
Tri Dao committed
968
969
970
971
972
        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()

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

978
979
980
981
    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
982
983
984
985
986
987
        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()

988
    def key_mapping_attn(key):
Tri Dao's avatar
Tri Dao committed
989
990
991
992
        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
        )
993
        return key
Tri Dao's avatar
Tri Dao committed
994

995
996
997
    state_dict = OrderedDict((key_mapping_attn(k), v) for k, v in state_dict.items())

    return state_dict
998
999


Tri Dao's avatar
Tri Dao committed
1000
1001
def remap_state_dict_megatron(state_dict, config):
    def key_mapping_transformer(key):
Tri Dao's avatar
Tri Dao committed
1002
1003
        key = re.sub(r"^language_model.encoder.", "transformer.", key)
        key = re.sub(r"^language_model.", "transformer.", key)
Tri Dao's avatar
Tri Dao committed
1004
        return key
Tri Dao's avatar
Tri Dao committed
1005

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

Tri Dao's avatar
Tri Dao committed
1008
1009
    # Word embedding and position embedding
    def key_mapping_pos_emb(key):
Tri Dao's avatar
Tri Dao committed
1010
1011
        return re.sub(r"^wpe.", "transformer.embeddings.position_embeddings.", key)

Tri Dao's avatar
Tri Dao committed
1012
    state_dict = OrderedDict((key_mapping_pos_emb(k), v) for k, v in state_dict.items())
Tri Dao's avatar
Tri Dao committed
1013
    word_embeddings = state_dict.pop("transformer.embedding.word_embeddings.weight")
Tri Dao's avatar
Tri Dao committed
1014
    # It's possible that vocab_size is padded to be a multiple of 8, for example.
Tri Dao's avatar
Tri Dao committed
1015
1016
1017
1018
1019
    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
1020
1021
        word_embeddings, (0, 0, 0, vocab_size - word_embeddings.shape[0])
    )
Tri Dao's avatar
Tri Dao committed
1022
    state_dict["lm_head.weight"] = state_dict["transformer.embeddings.word_embeddings.weight"]
1023

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

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

Tri Dao's avatar
Tri Dao committed
1041
1042
    # MLP
    def key_mapping_mlp(key):
Tri Dao's avatar
Tri Dao committed
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
        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
1053
        return key
Tri Dao's avatar
Tri Dao committed
1054

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

Tri Dao's avatar
Tri Dao committed
1057
1058
    # Attention
    def key_mapping_attn(key):
Tri Dao's avatar
Tri Dao committed
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
        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
1074
        return key
Tri Dao's avatar
Tri Dao committed
1075

Tri Dao's avatar
Tri Dao committed
1076
1077
1078
1079
1080
    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
1081
1082
1083
1084
1085
1086
        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
1087
        )
Tri Dao's avatar
Tri Dao committed
1088
1089
1090
        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
1091
        )
1092
1093

    return state_dict