0d810cfb73
Update handling for KeyError in state_dict.pop() for non-existing keys.
Changed state_dict.pop(f"h.{d}.attn.bias") to state_dict.pop(f"h.{d}.attn.bias", None) to prevent KeyError exceptions.
The following code can re-produce the issue
```
from transformers import AutoTokenizer, GPT2Model, GPT2Config
from flash_attn.models.gpt import GPTLMHeadModel, GPTModel
# >>> transformers.__version__
# '4.38.2'
model_path = 'gpt2'
output_model_path = 'gpt2_model'
config = GPT2Config.from_pretrained(model_path, output_hidden_states=True)
model = GPT2Model.from_pretrained(model_path, from_tf=False, config=config)
'''
model fine-tuning here
'''
# dump the fine-tuned model
model.save_pretrained(output_model_path)
# load the fine-tuned model
config = GPT2Config.from_pretrained(output_model_path, output_hidden_states=True)
model = GPTModel.from_pretrained(output_model_path, config=config, strict=True) # failed due to KeyError: 'h.0.attn.bias'
model = GPTLMHeadModel.from_pretrained(output_model_path, config=config, strict=True) # failed due to KeyError: 'h.0.attn.bias'
```
1081 lines
47 KiB
Python
1081 lines
47 KiB
Python
# Copyright (c) 2024, Tri Dao.
|
|
|
|
import logging
|
|
import math
|
|
import re
|
|
from collections import OrderedDict, namedtuple
|
|
from collections.abc import Sequence
|
|
from functools import partial
|
|
from typing import Dict, List
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from einops import rearrange
|
|
from transformers import GPT2Config
|
|
|
|
from flash_attn.models.bigcode import remap_state_dict_hf_bigcode
|
|
from flash_attn.models.falcon import remap_state_dict_hf_falcon
|
|
from flash_attn.models.gpt_neox import remap_state_dict_hf_gpt_neox
|
|
from flash_attn.models.gptj import remap_state_dict_hf_gptj
|
|
from flash_attn.models.llama import remap_state_dict_hf_llama
|
|
from flash_attn.models.opt import remap_state_dict_hf_opt
|
|
from flash_attn.modules.block import Block, ParallelBlock
|
|
from flash_attn.modules.embedding import GPT2Embeddings, ParallelGPT2Embeddings
|
|
from flash_attn.modules.mha import MHA, ParallelMHA
|
|
from flash_attn.modules.mlp import (
|
|
FusedMLP,
|
|
GatedMlp,
|
|
Mlp,
|
|
ParallelFusedMLP,
|
|
ParallelGatedMlp,
|
|
ParallelMLP,
|
|
)
|
|
from flash_attn.ops.activations import sqrelu_fwd
|
|
from flash_attn.utils.distributed import (
|
|
all_gather,
|
|
all_gather_raw,
|
|
get_dim_for_local_rank,
|
|
sync_shared_params,
|
|
)
|
|
from flash_attn.utils.generation import GenerationMixin
|
|
from flash_attn.utils.pretrained import state_dict_from_pretrained
|
|
|
|
try:
|
|
from flash_attn.ops.fused_dense import ColumnParallelLinear
|
|
except ImportError:
|
|
ColumnParallelLinear = None
|
|
|
|
try:
|
|
from flash_attn.ops.triton.mlp import FusedDenseSqreluDense
|
|
except ImportError:
|
|
FusedDenseSqreluDense = None
|
|
|
|
try:
|
|
from flash_attn.ops.triton.layer_norm import layer_norm_fn, RMSNorm
|
|
except ImportError:
|
|
layer_norm_fn, RMSNorm = None, None
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def create_mixer_cls(config, layer_idx=None, process_group=None, device=None, dtype=None):
|
|
factory_kwargs = {"device": device, "dtype": dtype}
|
|
head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
|
|
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)
|
|
if config.scale_attn_by_inverse_layer_idx:
|
|
assert layer_idx is not None
|
|
softmax_scale /= float(layer_idx + 1)
|
|
dwconv = getattr(config, "attn_dwconv", False)
|
|
if dwconv:
|
|
assert process_group is None, "TensorParallel MHA does not support dwconv yet"
|
|
qkv_proj_bias = getattr(config, "qkv_proj_bias", True)
|
|
out_proj_bias = getattr(config, "out_proj_bias", True)
|
|
rotary_emb_dim = int(getattr(config, "rotary_emb_fraction", 0.0) * head_dim)
|
|
rotary_emb_base = getattr(config, "rotary_emb_base", 10000.0)
|
|
rotary_emb_scale_base = getattr(config, "rotary_emb_scale_base", None)
|
|
rotary_emb_interleaved = getattr(config, "rotary_emb_interleaved", False)
|
|
use_alibi = getattr(config, "use_alibi", False)
|
|
window_size = getattr(config, "window_size", (-1, -1))
|
|
use_flash_attn = getattr(config, "use_flash_attn", False)
|
|
fused_bias_fc = getattr(config, "fused_bias_fc", False)
|
|
if not fused_bias_fc:
|
|
assert process_group is None, "TensorParallel MHA requires fused_bias_fc"
|
|
mha_cls = MHA if process_group is None else ParallelMHA
|
|
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 {}
|
|
)
|
|
num_heads_kv = getattr(config, "n_head_kv", None)
|
|
mixer_cls = partial(
|
|
mha_cls,
|
|
num_heads=config.num_attention_heads,
|
|
num_heads_kv=num_heads_kv,
|
|
qkv_proj_bias=qkv_proj_bias,
|
|
out_proj_bias=out_proj_bias,
|
|
dropout=config.attn_pdrop,
|
|
softmax_scale=softmax_scale,
|
|
causal=True,
|
|
layer_idx=layer_idx,
|
|
rotary_emb_dim=rotary_emb_dim,
|
|
rotary_emb_base=rotary_emb_base,
|
|
rotary_emb_scale_base=rotary_emb_scale_base,
|
|
rotary_emb_interleaved=rotary_emb_interleaved,
|
|
use_alibi=use_alibi,
|
|
window_size=window_size,
|
|
use_flash_attn=use_flash_attn,
|
|
**serial_kwargs,
|
|
**parallel_kwargs,
|
|
**factory_kwargs,
|
|
)
|
|
return mixer_cls
|
|
|
|
|
|
def create_mlp_cls(config, layer_idx=None, process_group=None, device=None, dtype=None):
|
|
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)
|
|
if fused_mlp:
|
|
assert config.activation_function in [
|
|
"gelu_new",
|
|
"gelu_fast",
|
|
"gelu_approx",
|
|
"gelu_pytorch_tanh",
|
|
"relu",
|
|
"sqrelu",
|
|
]
|
|
fused_dense_sqrelu_dense = getattr(config, "fused_dense_sqrelu_dense", False)
|
|
if fused_dense_sqrelu_dense:
|
|
assert config.activation_function == "sqrelu", (
|
|
"fused_dense_sqrelu_dense only " "supports approximate activation_function sqrelu"
|
|
)
|
|
assert not (fused_dense_sqrelu_dense and fused_mlp)
|
|
if not fused_mlp and not fused_dense_sqrelu_dense:
|
|
assert config.activation_function in [
|
|
"gelu",
|
|
"gelu_new",
|
|
"gelu_fast",
|
|
"gelu_approx",
|
|
"gelu_pytorch_tanh",
|
|
"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)
|
|
)
|
|
mlp_cls = GatedMlp if process_group is None else ParallelGatedMlp
|
|
parallel_kwargs = (
|
|
{
|
|
"process_group": process_group,
|
|
"sequence_parallel": getattr(config, "sequence_parallel", True),
|
|
}
|
|
if process_group is not None
|
|
else {}
|
|
)
|
|
mlp_multiple_of = getattr(config, "mlp_multiple_of", 128)
|
|
mlp_cls = partial(
|
|
mlp_cls,
|
|
hidden_features=config.n_inner,
|
|
activation=activation,
|
|
bias1=mlp_fc1_bias,
|
|
bias2=mlp_fc2_bias,
|
|
multiple_of=mlp_multiple_of,
|
|
**parallel_kwargs,
|
|
**factory_kwargs,
|
|
)
|
|
else:
|
|
if config.activation_function == "relu":
|
|
activation = partial(F.relu, inplace=True)
|
|
elif config.activation_function == "sqrelu":
|
|
activation = sqrelu_fwd
|
|
else:
|
|
approximate = (
|
|
"tanh"
|
|
if config.activation_function
|
|
in ["gelu_new", "gelu_fast", "gelu_approx", "gelu_pytorch_tanh"]
|
|
else "none"
|
|
)
|
|
activation = partial(F.gelu, approximate=approximate)
|
|
mlp_cls = Mlp if process_group is None else ParallelMLP
|
|
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,
|
|
)
|
|
else:
|
|
mlp_checkpoint_lvl = getattr(config, "mlp_checkpoint_lvl", 0)
|
|
# 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]
|
|
if fused_mlp:
|
|
if FusedMLP is None:
|
|
raise ImportError("fused_dense is not installed")
|
|
activation = (
|
|
"gelu_approx"
|
|
if config.activation_function
|
|
in ["gelu_new", "gelu_fast", "gelu_approx", "gelu_pytorch_tanh"]
|
|
else config.activation_function
|
|
)
|
|
mlp_cls = FusedMLP if process_group is None else ParallelFusedMLP
|
|
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,
|
|
)
|
|
elif fused_dense_sqrelu_dense:
|
|
if process_group is not None:
|
|
assert fused_mlp, "Tensor Parallel is not implemented for FusedDenseSqreluDense"
|
|
assert FusedDenseSqreluDense is not None
|
|
mlp_cls = partial(
|
|
FusedDenseSqreluDense,
|
|
hidden_features=config.n_inner,
|
|
checkpoint_lvl=mlp_checkpoint_lvl,
|
|
**factory_kwargs,
|
|
)
|
|
else:
|
|
raise RuntimeError("MLP type not supported")
|
|
return mlp_cls
|
|
|
|
|
|
def create_block(config, layer_idx=None, process_group=None, device=None, dtype=None):
|
|
factory_kwargs = {"device": device, "dtype": dtype}
|
|
sequence_parallel = getattr(config, "sequence_parallel", True)
|
|
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)
|
|
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,
|
|
)
|
|
# TD [2022-07-30]: Force residual in fp32, seems to make fp16 training more stable
|
|
residual_in_fp32 = getattr(config, "residual_in_fp32", False)
|
|
resid_dropout1 = config.resid_pdrop if layer_idx is None or layer_idx > 0 else config.embd_pdrop
|
|
prenorm = getattr(config, "prenorm", True)
|
|
parallel_block = getattr(config, "parallel_block", False)
|
|
if not parallel_block:
|
|
block = Block(
|
|
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),
|
|
residual_in_fp32=residual_in_fp32,
|
|
sequence_parallel=sequence_parallel and process_group is not None,
|
|
mark_shared_params=process_group is not None,
|
|
)
|
|
else:
|
|
assert prenorm
|
|
block = ParallelBlock(
|
|
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),
|
|
residual_in_fp32=residual_in_fp32,
|
|
sequence_parallel=sequence_parallel and process_group is not None,
|
|
mark_shared_params=process_group is not None,
|
|
)
|
|
block.layer_idx = layer_idx
|
|
return block
|
|
|
|
|
|
class GPTPreTrainedModel(nn.Module):
|
|
"""An abstract class to handle weights initialization and
|
|
a simple interface for dowloading and loading pretrained models.
|
|
"""
|
|
|
|
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__
|
|
)
|
|
)
|
|
self.config = config
|
|
|
|
@classmethod
|
|
def from_pretrained(
|
|
cls,
|
|
model_name,
|
|
config,
|
|
*args,
|
|
strict=True,
|
|
device=None,
|
|
dtype=None,
|
|
world_size=1,
|
|
rank=0,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
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.
|
|
model = cls(config, *args, device=device, dtype=dtype, **kwargs)
|
|
# 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
|
|
state_dict = state_dict_from_pretrained(model_name, device="cpu", dtype=dtype)
|
|
if model_name.startswith("gpt2"):
|
|
state_dict = remap_state_dict_hf_gpt2(state_dict, config)
|
|
elif model_name.startswith("facebook/opt"):
|
|
state_dict = remap_state_dict_hf_opt(state_dict, config)
|
|
elif model_name.startswith("EleutherAI/gpt-j-") or model_name.startswith(
|
|
"togethercomputer/GPT-JT-"
|
|
):
|
|
state_dict = remap_state_dict_hf_gptj(state_dict, config)
|
|
elif (
|
|
model_name.startswith("EleutherAI/gpt-neox-")
|
|
or model_name.startswith("EleutherAI/pythia-")
|
|
or model_name.startswith("togethercomputer/RedPajama-INCITE-")
|
|
):
|
|
state_dict = remap_state_dict_hf_gpt_neox(state_dict, config)
|
|
elif model_name.startswith("tiiuae/falcon-"):
|
|
state_dict = remap_state_dict_hf_falcon(state_dict, config)
|
|
elif model_name.startswith("meta-llama/Llama-"):
|
|
state_dict = remap_state_dict_hf_llama(state_dict, config)
|
|
elif model_name.startswith("bigcode/") or model_name.startswith("WizardLM/"):
|
|
state_dict = remap_state_dict_hf_bigcode(state_dict, config)
|
|
else:
|
|
raise NotImplementedError(f"Model {model_name} not supported")
|
|
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)
|
|
logger.info(load_return)
|
|
return model
|
|
|
|
|
|
# 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, mup_width_scale=1.0, rescale_prenorm_residual=True
|
|
):
|
|
mup_init_scale = math.sqrt(mup_width_scale)
|
|
if isinstance(module, nn.Linear):
|
|
nn.init.normal_(module.weight, std=initializer_range * mup_init_scale)
|
|
optim_cfg = getattr(module.weight, "_optim", {})
|
|
optim_cfg.update({"lr_multiplier": mup_width_scale})
|
|
setattr(module.weight, "_optim", optim_cfg)
|
|
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 * mup_init_scale / math.sqrt(2 * n_layer)
|
|
)
|
|
|
|
|
|
class GPTModel(GPTPreTrainedModel):
|
|
def __init__(self, config: GPT2Config, process_group=None, device=None, dtype=None):
|
|
super().__init__(config)
|
|
factory_kwargs = {"device": device, "dtype": dtype}
|
|
self.process_group = process_group
|
|
self.sequence_parallel = getattr(config, "sequence_parallel", True)
|
|
assert config.activation_function in [
|
|
"gelu",
|
|
"gelu_new",
|
|
"gelu_fast",
|
|
"gelu_approx",
|
|
"gelu_pytorch_tanh",
|
|
"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
|
|
)
|
|
self.embeddings_multiplier = getattr(config, "mup_embeddings_multiplier", 1.0)
|
|
# TD [2022-07-30]: Force residual in fp32, seems to make fp16 training more stable
|
|
self.residual_in_fp32 = getattr(config, "residual_in_fp32", False)
|
|
# These 2 options are for OPT-350m
|
|
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)
|
|
# For GPT-J, GPT-NeoX
|
|
self.parallel_block = getattr(config, "parallel_block", False)
|
|
|
|
if process_group is None:
|
|
self.embeddings = GPT2Embeddings(
|
|
config.hidden_size,
|
|
vocab_size,
|
|
config.max_position_embeddings,
|
|
word_embed_proj_dim=word_embed_proj_dim,
|
|
**factory_kwargs,
|
|
)
|
|
else:
|
|
self.embeddings = ParallelGPT2Embeddings(
|
|
config.hidden_size,
|
|
vocab_size,
|
|
config.max_position_embeddings,
|
|
process_group=process_group,
|
|
sequence_parallel=self.sequence_parallel,
|
|
**factory_kwargs,
|
|
)
|
|
|
|
# We change the order of dropout, residual and layer norm:
|
|
# Instead of LN -> Attn / MLP -> Dropout -> Add, we do:
|
|
# 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.
|
|
# This is for performance reason: we can fuse dropout + add + layer_norm.
|
|
self.layers = nn.ModuleList(
|
|
[
|
|
create_block(config, layer_idx=i, process_group=process_group, **factory_kwargs)
|
|
for i in range(config.num_hidden_layers)
|
|
]
|
|
)
|
|
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
|
|
|
|
self.fused_dropout_add_ln = getattr(config, "fused_dropout_add_ln", False)
|
|
if self.fused_dropout_add_ln:
|
|
if layer_norm_fn is None:
|
|
raise ImportError("Triton is not installed")
|
|
if self.prenorm:
|
|
self.drop_f = nn.Dropout(config.resid_pdrop)
|
|
norm_cls = nn.LayerNorm if not use_rms_norm else RMSNorm
|
|
self.ln_f = norm_cls(
|
|
config.hidden_size, eps=config.layer_norm_epsilon, **factory_kwargs
|
|
)
|
|
if process_group is not None:
|
|
for p in self.ln_f.parameters():
|
|
# 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
|
|
|
|
self.apply(
|
|
partial(
|
|
_init_weights,
|
|
n_layer=config.num_hidden_layers,
|
|
initializer_range=config.initializer_range,
|
|
mup_width_scale=getattr(config, "mup_width_scale", 1.0),
|
|
)
|
|
)
|
|
self.tie_weights()
|
|
|
|
def tie_weights(self):
|
|
if self.process_group is not None:
|
|
sync_shared_params(self, self.process_group)
|
|
|
|
def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
|
|
return {
|
|
i: layer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)
|
|
for i, layer in enumerate(self.layers)
|
|
}
|
|
|
|
def forward(self, input_ids, position_ids=None, inference_params=None):
|
|
# 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.
|
|
embedding_kwargs = (
|
|
{"combine_batch_seqlen_dim": True}
|
|
if self.process_group is not None and self.sequence_parallel
|
|
else {}
|
|
)
|
|
hidden_states = self.embeddings(input_ids, position_ids=position_ids, **embedding_kwargs)
|
|
if self.embeddings_multiplier != 1.0:
|
|
hidden_states = hidden_states * self.embeddings_multiplier
|
|
if self.parallel_block:
|
|
hidden_states2 = None
|
|
residual = None
|
|
mixer_kwargs = (
|
|
{"seqlen": input_ids.shape[1]}
|
|
if self.process_group is not None and self.sequence_parallel
|
|
else {}
|
|
)
|
|
if inference_params is not None:
|
|
mixer_kwargs["inference_params"] = inference_params
|
|
for layer in self.layers:
|
|
if self.prenorm:
|
|
if not self.parallel_block:
|
|
hidden_states, residual = layer(
|
|
hidden_states, residual, mixer_kwargs=mixer_kwargs
|
|
)
|
|
else:
|
|
hidden_states, hidden_states2, residual = layer(
|
|
hidden_states, hidden_states2, residual, mixer_kwargs=mixer_kwargs
|
|
)
|
|
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)
|
|
if not self.parallel_block:
|
|
residual = (dropped + residual) if residual is not None else dropped
|
|
else:
|
|
dropped2 = self.drop_f(hidden_states2)
|
|
residual = (
|
|
(residual + dropped + dropped2)
|
|
if residual is not None
|
|
else dropped + dropped2
|
|
)
|
|
hidden_states = self.ln_f(residual.to(dtype=self.ln_f.weight.dtype))
|
|
else:
|
|
# Set prenorm=False here since we don't need the residual
|
|
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)
|
|
)
|
|
return hidden_states
|
|
|
|
|
|
class GPTLMHeadModel(GPTPreTrainedModel, GenerationMixin):
|
|
def __init__(self, config: GPT2Config, process_group=None, device=None, dtype=None):
|
|
factory_kwargs = {"device": device, "dtype": dtype}
|
|
super().__init__(config)
|
|
self.process_group = process_group
|
|
self.transformer = GPTModel(config, process_group=process_group, **factory_kwargs)
|
|
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
|
|
)
|
|
# This option is for OPT-350m
|
|
word_embed_proj_dim = getattr(config, "word_embed_proj_dim", None)
|
|
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
|
|
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
|
|
if process_group is None:
|
|
self.lm_head = nn.Linear(embed_dim, vocab_size, bias=lm_head_bias, **factory_kwargs)
|
|
else:
|
|
if ColumnParallelLinear is None:
|
|
raise ImportError("fused_dense_lib is not installed")
|
|
self.lm_head = ColumnParallelLinear(
|
|
embed_dim,
|
|
vocab_size,
|
|
process_group,
|
|
bias=lm_head_bias,
|
|
sequence_parallel=getattr(config, "sequence_parallel", True),
|
|
**factory_kwargs,
|
|
)
|
|
self.norm_head = getattr(config, "norm_head", False)
|
|
# Initialize weights and apply final processing
|
|
self.apply(
|
|
partial(
|
|
_init_weights,
|
|
n_layer=config.num_hidden_layers,
|
|
initializer_range=config.initializer_range,
|
|
mup_width_scale=mup_width_scale,
|
|
)
|
|
)
|
|
self.tie_weights()
|
|
|
|
def tie_weights(self):
|
|
if self.tie_word_embeddings:
|
|
self.lm_head.weight = self.transformer.embeddings.word_embeddings.weight
|
|
if self.process_group is not None:
|
|
sync_shared_params(self, self.process_group)
|
|
|
|
def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
|
|
return self.transformer.allocate_inference_cache(
|
|
batch_size, max_seqlen, dtype=dtype, **kwargs
|
|
)
|
|
|
|
def forward(self, input_ids, position_ids=None, inference_params=None, num_last_tokens=0):
|
|
"""
|
|
input_ids: (batch, seqlen) int tensor
|
|
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
|
|
num_last_tokens: if > 0, only return the logits for the last n tokens
|
|
"""
|
|
assert (
|
|
input_ids.ndim == 2
|
|
), f"Expected `input_ids` to have shape [b, slen], but got shape {input_ids.shape}"
|
|
b, slen = input_ids.shape
|
|
hidden_states = self.transformer(
|
|
input_ids, position_ids=position_ids, inference_params=inference_params
|
|
)
|
|
if inference_params is not None:
|
|
assert hidden_states.ndim == 3, "sequence_parallel is not supported in generation mode"
|
|
if num_last_tokens > 0:
|
|
hidden_states = hidden_states[:, -num_last_tokens:]
|
|
if self.project_out is not None:
|
|
hidden_states = self.project_out(hidden_states)
|
|
if self.output_scale != 1.0:
|
|
hidden_states = hidden_states * self.output_scale
|
|
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)
|
|
# 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)
|
|
lm_logits = rearrange(lm_logits, "(n b) ... d -> b ... (n d)", b=b)
|
|
CausalLMOutput = namedtuple("CausalLMOutput", ["logits"])
|
|
return CausalLMOutput(logits=lm_logits)
|
|
|
|
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
|
|
if "transformer.ln_0.weight" in state_dict:
|
|
n_layers = len(self.transformer.layers)
|
|
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
|
|
for l in reversed(range(n_layers)):
|
|
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
|
|
if l > 0:
|
|
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
|
|
return super().load_state_dict(state_dict, strict=strict)
|
|
|
|
|
|
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.
|
|
|
|
This function modifies state_dict in place.
|
|
"""
|
|
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
|
|
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
|
|
|
|
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
|
|
|
|
def shard_first_dim(state_dict, key):
|
|
if key in state_dict:
|
|
x = state_dict[key]
|
|
dim = x.shape[0] // world_size
|
|
state_dict[key] = x[rank * dim : (rank + 1) * dim]
|
|
|
|
def shard_last_dim(state_dict, key, multiple_of=1):
|
|
if key in state_dict:
|
|
x = state_dict[key]
|
|
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]
|
|
|
|
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(
|
|
rearrange(x, "(two o) ... -> two o ...", two=2)[:, rank * dim : (rank + 1) * dim],
|
|
"two o ... -> (two o) ...",
|
|
)
|
|
|
|
def shard_qkv_headdim(state_dict, key):
|
|
if key in state_dict:
|
|
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)
|
|
]
|
|
|
|
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])
|
|
|
|
if n_head_kv == n_head:
|
|
x = rearrange(state_dict[key], "(three d) ... -> three d ...", three=3)
|
|
state_dict[key] = rearrange(
|
|
x[:, beg_n_head * head_dim : end_n_head * head_dim],
|
|
"three d ... -> (three d) ...",
|
|
)
|
|
else:
|
|
x = rearrange(
|
|
state_dict[key],
|
|
"(nheadqkv headdim) ... -> nheadqkv headdim ...",
|
|
nheadqkv=n_head + 2 * n_head_kv,
|
|
)
|
|
state_dict[key] = rearrange(
|
|
torch.cat(
|
|
[
|
|
x[beg_n_head:end_n_head],
|
|
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
|
|
],
|
|
],
|
|
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")
|
|
for i in range(config.num_hidden_layers):
|
|
shard_qkv_headdim(state_dict, f"transformer.layers.{i}.mixer.Wqkv.weight")
|
|
shard_qkv_headdim(state_dict, f"transformer.layers.{i}.mixer.Wqkv.bias")
|
|
shard_last_dim(
|
|
state_dict, f"transformer.layers.{i}.mixer.out_proj.weight", multiple_of=head_dim
|
|
)
|
|
if rank != 0:
|
|
state_dict.pop(f"transformer.layers.{i}.mixer.out_proj.bias", None)
|
|
if config.activation_function in ["glu", "swiglu", "geglu"]:
|
|
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")
|
|
else:
|
|
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")
|
|
if rank != 0:
|
|
state_dict.pop(f"transformer.layers.{i}.mlp.fc2.bias", None)
|
|
return state_dict
|
|
|
|
|
|
def combine_state_dicts_tp(state_dicts: List[Dict[str, torch.Tensor]], config: GPT2Config):
|
|
"""Convert the list of sharded state_dict of a GPT model with tensor parallel to
|
|
the state_dict of a standard GPT model.
|
|
|
|
This function is meant to be the "reverse" of shard_state_dict_tp.
|
|
|
|
Precondition:
|
|
- state_dicts should be ordered in the same way as the shards were created.
|
|
"""
|
|
world_size = len(state_dicts)
|
|
keys = state_dicts[0].keys()
|
|
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
|
|
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
|
|
assert config.hidden_size % config.n_head == 0
|
|
headdim = config.hidden_size // config.n_head
|
|
|
|
# Sometimes the word embeddings are sharded on the 0th dim, sometimes on the 1st dim.
|
|
# vocab_size // world_size coordinates are nonzero.
|
|
def combine_word_embeddings(state_dicts, state_dict, key):
|
|
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)
|
|
|
|
def combine_dim(state_dicts, state_dict, key, dim=-1):
|
|
if key in state_dict:
|
|
state_dict[key] = torch.cat([s[key] for s in state_dicts], dim=dim)
|
|
|
|
def combine_qkv_headdim(state_dicts, state_dict, key):
|
|
n_head = config.n_head
|
|
n_head_kv = getattr(config, "n_head_kv", n_head)
|
|
if key in state_dict:
|
|
if n_head_kv == n_head:
|
|
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) ...")
|
|
else:
|
|
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)
|
|
]
|
|
xs = [
|
|
rearrange(
|
|
s[key],
|
|
"(nheadqkv headdim) ... -> nheadqkv headdim ...",
|
|
nheadqkv=rank_n_head + 2 * rank_n_head_kv,
|
|
headdim=headdim,
|
|
)
|
|
for s, rank_n_head, rank_n_head_kv in zip(
|
|
state_dicts, n_head_each_rank, n_head_kv_each_rank
|
|
)
|
|
]
|
|
wq = torch.cat([x[: n_head_each_rank[rank]] for rank, x in enumerate(xs)], dim=0)
|
|
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,
|
|
)
|
|
state_dict[key] = rearrange(
|
|
wqkv,
|
|
"nheadqkv headdim ... -> (nheadqkv headdim) ...",
|
|
)
|
|
|
|
def combine_gated_mlp(state_dicts, state_dict, key):
|
|
if key in state_dict:
|
|
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) ...")
|
|
|
|
state_dict = state_dicts[0].copy() # don't modify state_dict[0] inplace
|
|
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)
|
|
)
|
|
for i in range(config.num_hidden_layers):
|
|
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)
|
|
return state_dict
|
|
|
|
|
|
def remap_state_dict_hf_gpt2(state_dict, config):
|
|
# Word embedding and position embedding
|
|
def key_mapping_pos_emb(key):
|
|
return re.sub(r"^wpe.", "transformer.embeddings.position_embeddings.", key)
|
|
|
|
state_dict = OrderedDict((key_mapping_pos_emb(k), v) for k, v in state_dict.items())
|
|
word_embeddings = state_dict.pop("wte.weight")
|
|
# It's possible that vocab_size is padded to be a multiple of 8, for example.
|
|
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(
|
|
word_embeddings, (0, 0, 0, vocab_size - word_embeddings.shape[0])
|
|
)
|
|
state_dict["lm_head.weight"] = state_dict["transformer.embeddings.word_embeddings.weight"]
|
|
|
|
# LayerNorm
|
|
def key_mapping_ln(key):
|
|
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)
|
|
return key
|
|
|
|
state_dict = OrderedDict((key_mapping_ln(k), v) for k, v in state_dict.items())
|
|
|
|
# MLP
|
|
for d in range(config.num_hidden_layers):
|
|
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()
|
|
|
|
def key_mapping_mlp(key):
|
|
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)
|
|
return key
|
|
|
|
state_dict = OrderedDict((key_mapping_mlp(k), v) for k, v in state_dict.items())
|
|
|
|
# Attention
|
|
for d in range(config.num_hidden_layers):
|
|
state_dict.pop(f"h.{d}.attn.bias", None) # 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()
|
|
|
|
def key_mapping_attn(key):
|
|
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
|
|
)
|
|
return key
|
|
|
|
state_dict = OrderedDict((key_mapping_attn(k), v) for k, v in state_dict.items())
|
|
|
|
return state_dict
|
|
|
|
|
|
def remap_state_dict_megatron(state_dict, config):
|
|
def key_mapping_transformer(key):
|
|
key = re.sub(r"^language_model.encoder.", "transformer.", key)
|
|
key = re.sub(r"^language_model.", "transformer.", key)
|
|
return key
|
|
|
|
state_dict = OrderedDict((key_mapping_transformer(k), v) for k, v in state_dict.items())
|
|
|
|
# Word embedding and position embedding
|
|
def key_mapping_pos_emb(key):
|
|
return re.sub(r"^wpe.", "transformer.embeddings.position_embeddings.", key)
|
|
|
|
state_dict = OrderedDict((key_mapping_pos_emb(k), v) for k, v in state_dict.items())
|
|
word_embeddings = state_dict.pop("transformer.embedding.word_embeddings.weight")
|
|
# It's possible that vocab_size is padded to be a multiple of 8, for example.
|
|
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(
|
|
word_embeddings, (0, 0, 0, vocab_size - word_embeddings.shape[0])
|
|
)
|
|
state_dict["lm_head.weight"] = state_dict["transformer.embeddings.word_embeddings.weight"]
|
|
|
|
# LayerNorm
|
|
def key_mapping_ln(key):
|
|
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,
|
|
)
|
|
return key
|
|
|
|
state_dict = OrderedDict((key_mapping_ln(k), v) for k, v in state_dict.items())
|
|
|
|
# MLP
|
|
def key_mapping_mlp(key):
|
|
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,
|
|
)
|
|
return key
|
|
|
|
state_dict = OrderedDict((key_mapping_mlp(k), v) for k, v in state_dict.items())
|
|
|
|
# Attention
|
|
def key_mapping_attn(key):
|
|
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,
|
|
)
|
|
return key
|
|
|
|
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):
|
|
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,
|
|
)
|
|
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
|
|
)
|
|
|
|
return state_dict
|