# Copyright (c) 2023, Tri Dao. import json import math import os import re from collections import OrderedDict from pathlib import Path from typing import Dict, List, Union import torch import torch.nn.functional as F from sentencepiece import SentencePieceProcessor from transformers import GPT2Config, LlamaConfig from einops import rearrange def remap_state_dict_meta_llama( state_dict: Dict[str, torch.Tensor], config: GPT2Config ) -> Dict[str, torch.Tensor]: """Convert the state_dict in Meta format to standard GPT format. This function modifies state_dict in place. """ def key_mapping_layers(key): return f"transformer.{key}" if not key.startswith("output.") else key state_dict = OrderedDict((key_mapping_layers(k), v) for k, v in state_dict.items()) # Word embedding def key_mapping_emb(key): return re.sub( r"^transformer.tok_embeddings.", "transformer.embeddings.word_embeddings.", key ) state_dict = OrderedDict((key_mapping_emb(k), v) for k, v in state_dict.items()) word_embeddings = state_dict.pop("transformer.embeddings.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]) ) if getattr(config, "tie_word_embeddings"): state_dict["lm_head.weight"] = state_dict["transformer.embeddings.word_embeddings.weight"] else: output_embeddings = state_dict.pop("output.weight") # Need to recompute vocab_size since LLaMa shards the word embeddings and output embeddings # differently. vocab_size = ( math.ceil(output_embeddings.shape[0] / pad_vocab_size_multiple) * pad_vocab_size_multiple ) # It's possible that vocab_size is padded to be a multiple of 8, for example. state_dict["lm_head.weight"] = F.pad( output_embeddings, (0, 0, 0, vocab_size - output_embeddings.shape[0]) ) # LayerNorm def key_mapping_ln(key): key = re.sub(r"^transformer.norm.", r"transformer.ln_f.", key) key = re.sub( r"^transformer.layers.(\d+).attention_norm.", r"transformer.layers.\1.norm1.", key, ) key = re.sub(r"^transformer.layers.(\d+).ffn_norm.", r"transformer.layers.\1.norm2.", key) return key state_dict = OrderedDict((key_mapping_ln(k), v) for k, v in state_dict.items()) # MLP for l in range(config.n_layer): w1 = state_dict.pop(f"transformer.layers.{l}.feed_forward.w1.weight") w3 = state_dict.pop(f"transformer.layers.{l}.feed_forward.w3.weight") # Our ordering is different state_dict[f"transformer.layers.{l}.mlp.fc1.weight"] = torch.cat([w3, w1], dim=0) def key_mapping_mlp(key): return re.sub( r"^transformer.layers.(\d+).feed_forward.w2.", r"transformer.layers.\1.mlp.fc2.", key, ) state_dict = OrderedDict((key_mapping_mlp(k), v) for k, v in state_dict.items()) # Attention for l in range(config.n_layer): Wq = state_dict.pop(f"transformer.layers.{l}.attention.wq.weight") Wk = state_dict.pop(f"transformer.layers.{l}.attention.wk.weight") Wv = state_dict.pop(f"transformer.layers.{l}.attention.wv.weight") state_dict[f"transformer.layers.{l}.mixer.Wqkv.weight"] = torch.cat([Wq, Wk, Wv], dim=0) # We don't store these state_dict.pop(f"transformer.layers.{l}.attention.inner_attention.rope.freqs", None) def key_mapping_attn(key): return re.sub( r"^transformer.layers.(\d+).attention.wo.", r"transformer.layers.\1.mixer.out_proj.", key, ) state_dict = OrderedDict((key_mapping_attn(k), v) for k, v in state_dict.items()) state_dict.pop("transformer.rope.freqs", None) return state_dict def remap_state_dict_hf_llama( state_dict: Dict[str, torch.Tensor], config: GPT2Config ) -> Dict[str, torch.Tensor]: """Convert the state_dict in Hugging Face format to standard GPT format. This function modifies state_dict in place. """ # Embedding def key_mapping_emb(key): return re.sub(r"^model.embed_tokens.", "transformer.embeddings.word_embeddings.", key) state_dict = OrderedDict((key_mapping_emb(k), v) for k, v in state_dict.items()) word_embeddings = state_dict.pop("transformer.embeddings.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]) ) # LM head if getattr(config, "tie_word_embeddings"): state_dict["lm_head.weight"] = state_dict["transformer.embeddings.word_embeddings.weight"] else: output_embeddings = state_dict.pop("lm_head.weight") # Need to recompute vocab_size since LLaMa shards the word embeddings and output embeddings # differently. vocab_size = ( math.ceil(output_embeddings.shape[0] / pad_vocab_size_multiple) * pad_vocab_size_multiple ) # It's possible that vocab_size is padded to be a multiple of 8, for example. state_dict["lm_head.weight"] = F.pad( output_embeddings, (0, 0, 0, vocab_size - output_embeddings.shape[0]) ) # MLP for l in range(config.n_layer): # Fusing weights this way based on difference in the following: # https://github.com/huggingface/transformers/blob/b42010bb1d3cbf262d27e0a328661885be46dfdb/src/transformers/models/llama/modeling_llama.py#L220 # https://github.com/Dao-AILab/flash-attention/blob/c60851a8253257eb970e06a022c82517a8033e8c/flash_attn/modules/mlp.py#L115 w1 = state_dict.pop(f"model.layers.{l}.mlp.gate_proj.weight") w3 = state_dict.pop(f"model.layers.{l}.mlp.up_proj.weight") state_dict[f"transformer.layers.{l}.mlp.fc1.weight"] = torch.cat([w3, w1], dim=0) def key_mapping_mlp(key): return re.sub( r"^model.layers.(\d+).mlp.down_proj.", r"transformer.layers.\1.mlp.fc2.", key, ) state_dict = OrderedDict((key_mapping_mlp(k), v) for k, v in state_dict.items()) # LayerNorm def key_mapping_ln(key): key = re.sub(r"^model.norm.", r"transformer.ln_f.", key) key = re.sub( r"^model.layers.(\d+).input_layernorm.", r"transformer.layers.\1.norm1.", key, ) key = re.sub( r"^model.layers.(\d+).post_attention_layernorm.", r"transformer.layers.\1.norm2.", key, ) return key state_dict = OrderedDict((key_mapping_ln(k), v) for k, v in state_dict.items()) def inv_permute(w): # Inverse of permute implemented in: # https://github.com/huggingface/transformers/blob/b42010bb1d3cbf262d27e0a328661885be46dfdb/src/transformers/models/llama/convert_llama_weights_to_hf.py#L114 return rearrange( w, "(h two d) n -> (h d two) n", d=config.n_embd // config.n_head // 2, two=2 ) # Attention for l in range(config.n_layer): Wq = state_dict.pop(f"model.layers.{l}.self_attn.q_proj.weight") Wk = state_dict.pop(f"model.layers.{l}.self_attn.k_proj.weight") Wv = state_dict.pop(f"model.layers.{l}.self_attn.v_proj.weight") state_dict[f"transformer.layers.{l}.mixer.Wqkv.weight"] = torch.cat( [inv_permute(Wq), inv_permute(Wk), Wv], dim=0 ) # We don't store these state_dict.pop(f"model.layers.{l}.self_attn.rotary_emb.inv_freq", None) def key_mapping_attn(key): return re.sub( r"^model.layers.(\d+).self_attn.o_proj.", r"transformer.layers.\1.mixer.out_proj.", key, ) state_dict = OrderedDict((key_mapping_attn(k), v) for k, v in state_dict.items()) return state_dict def inv_remap_state_dict_hf_llama( state_dict: Dict[str, torch.Tensor], config: GPT2Config ) -> Dict[str, torch.Tensor]: """Convert the state_dict in standard GPT format to Hugging Face format. This function is meant to be the inverse of remap_state_dict_hf_llama, up to a multiplier pad in the embedding and lm_head. That is if the original embedding isn't a multiple of pad_vocab_size_multiple, then inv_remap_state_dict_hf_llama(remap_state_dict_hf_llama(state_dict)) != state_dict. This function modifies state_dict in place. """ # Embedding def key_mapping_emb(key): return re.sub(r"^transformer.embeddings.word_embeddings.", "model.embed_tokens.", key) state_dict = OrderedDict((key_mapping_emb(k), v) for k, v in state_dict.items()) word_embeddings = state_dict.pop("model.embed_tokens.weight") 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["model.embed_tokens.weight"] = F.pad( word_embeddings, (0, 0, 0, vocab_size - word_embeddings.shape[0]) ) # LM head if getattr(config, "tie_word_embeddings"): state_dict["lm_head.weight"] = state_dict["model.embed_tokens.weight"] else: output_embeddings = state_dict.pop("lm_head.weight") vocab_size = ( math.ceil(output_embeddings.shape[0] / pad_vocab_size_multiple) * pad_vocab_size_multiple ) state_dict["lm_head.weight"] = F.pad( output_embeddings, (0, 0, 0, vocab_size - output_embeddings.shape[0]) ) # MLP for l in range(config.n_layer): w3, w1 = torch.chunk( state_dict.pop(f"transformer.layers.{l}.mlp.fc1.weight"), chunks=2, dim=0 ) state_dict[f"model.layers.{l}.mlp.gate_proj.weight"] = w1 state_dict[f"model.layers.{l}.mlp.up_proj.weight"] = w3 def key_mapping_mlp(key): return re.sub( r"^transformer.layers.(\d+).mlp.fc2.", r"model.layers.\1.mlp.down_proj.", key, ) state_dict = OrderedDict((key_mapping_mlp(k), v) for k, v in state_dict.items()) # LayerNorm def key_mapping_ln(key): key = re.sub(r"^transformer.ln_f.", r"model.norm.", key) key = re.sub( r"^transformer.layers.(\d+).norm1.", r"model.layers.\1.input_layernorm.", key, ) key = re.sub( r"^transformer.layers.(\d+).norm2.", r"model.layers.\1.post_attention_layernorm.", key, ) return key state_dict = OrderedDict((key_mapping_ln(k), v) for k, v in state_dict.items()) def permute(w): return rearrange( w, "(h d two) n -> (h two d) n", d=config.n_embd // config.n_head // 2, two=2 ) 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 q_dim = n_head * head_dim k_dim = v_dim = n_head_kv * head_dim # Attention for l in range(config.n_layer): Wqkv = state_dict.pop(f"transformer.layers.{l}.mixer.Wqkv.weight") Wq = Wqkv[:q_dim] Wk = Wqkv[q_dim : q_dim + k_dim] Wv = Wqkv[q_dim + k_dim : q_dim + k_dim + v_dim] state_dict[f"model.layers.{l}.self_attn.q_proj.weight"] = permute(Wq) state_dict[f"model.layers.{l}.self_attn.k_proj.weight"] = permute(Wk) state_dict[f"model.layers.{l}.self_attn.v_proj.weight"] = Wv state_dict.pop(f"transformer.layers.{l}.attention.inner_attention.rope.freqs", None) def key_mapping_attn(key): return re.sub( r"^transformer.layers.(\d+).mixer.out_proj.", r"model.layers.\1.self_attn.o_proj.", key, ) state_dict = OrderedDict((key_mapping_attn(k), v) for k, v in state_dict.items()) return state_dict def config_from_meta_checkpoint( checkpoint_path: Union[str, os.PathLike], model_name: str ) -> LlamaConfig: """Load a LlamaConfig from a checkpoint path.""" with open(Path(checkpoint_path) / model_name / "params.json") as f: params = json.load(f) config = LlamaConfig( hidden_size=params["dim"], intermediate_size=None, num_attention_heads=params["n_heads"], num_hidden_layers=params["n_layers"], rms_norm_eps=params["norm_eps"], num_key_value_heads=params.get("n_kv_heads", None), ) multiple_of = params.get("multiple_of", 1) ffn_dim_multiplier = params.get("ffn_dim_multiplier", None) # Compute the hidden dimension of the MLP # https://github.com/facebookresearch/llama/blob/1a240688810f8036049e8da36b073f63d2ac552c/llama/model.py#L224 intermediate_size = 4 * config.hidden_size # https://github.com/facebookresearch/llama/blob/1a240688810f8036049e8da36b073f63d2ac552c/llama/model.py#L195-L199 intermediate_size = int(2 * intermediate_size / 3) # custom dim factor multiplier if ffn_dim_multiplier is not None: intermediate_size = int(ffn_dim_multiplier * intermediate_size) intermediate_size = multiple_of * ((intermediate_size + multiple_of - 1) // multiple_of) config.intermediate_size = intermediate_size if "rope_theta" in params: config.rotary_emb_base = params["rope_theta"] config.vocab_size = 32000 # some CodeLLaMa have vocab_size 32000, some 32016 # Sadly it's not specified in the `params.json` file :( tokenizer = Path(checkpoint_path) / model_name / "tokenizer.model" if tokenizer.is_file(): config.vocab_size = SentencePieceProcessor(str(tokenizer)).vocab_size() return config def config_from_hf_checkpoint( checkpoint_path: Union[str, os.PathLike], model_name: str ) -> LlamaConfig: return LlamaConfig.from_pretrained(Path(checkpoint_path) / f"{model_name}-hf" / "config.json") def config_from_checkpoint( checkpoint_path: Union[str, os.PathLike], model_name: str, checkpoint_format="meta" ) -> LlamaConfig: if checkpoint_format == "meta": return config_from_meta_checkpoint(checkpoint_path, model_name) else: return config_from_hf_checkpoint(checkpoint_path, model_name) def state_dicts_from_checkpoint( checkpoint_path: Union[str, os.PathLike], model_name: str ) -> List[dict]: # Need to sort, otherwise we mess up the ordering and the weights are wrong return [ torch.load(path, map_location="cpu") for path in sorted((Path(checkpoint_path) / model_name).glob("consolidated.*.pth")) ] def llama_config_to_gpt2_config(llama_config: LlamaConfig) -> GPT2Config: return GPT2Config( vocab_size=llama_config.vocab_size, n_positions=0, # No absolute position embedding n_embd=llama_config.hidden_size, n_layer=llama_config.num_hidden_layers, n_head=llama_config.num_attention_heads, n_inner=llama_config.intermediate_size, activation_function="swiglu", # Hardcode since HF calls it 'silu' # Llama doesn't have dropout, idk if it's because they only release the inference code resid_pdrop=0.0, embd_pdrop=0.0, attn_pdrop=0.0, layer_norm_epsilon=llama_config.rms_norm_eps, initializer_range=llama_config.initializer_range, bos_token_id=llama_config.bos_token_id, eos_token_id=llama_config.eos_token_id, # These are new arguments not in the original GPT2Config pad_token_id=llama_config.pad_token_id, # Idk if this does anything rms_norm=True, rotary_emb_fraction=1.0, rotary_emb_interleaved=True, tie_word_embeddings=False, qkv_proj_bias=False, out_proj_bias=False, mlp_fc1_bias=False, mlp_fc2_bias=False, rotary_emb_base=getattr(llama_config, "rotary_emb_base", 10000.0), n_head_kv=llama_config.num_key_value_heads, )