## This script is adapted from: https://github.com/Lightning-AI/lit-gpt

## This script is used to convert the HF checkpoint to the LIT checkpoint

import gc

import json

import sys

from dataclasses import asdict

from functools import partial

from pathlib import Path

from typing import Dict, List, Optional, Tuple, Union

import torch

from lightning.fabric.utilities.load import _NotYetLoadedTensor as NotYetLoadedTensor

# support running without installing as a package

wd = Path(__file__).parent.parent.resolve()

sys.path.append(str(wd))

from lit_gpt import Config

from lit_gpt.utils import incremental_save, lazy_load

def copy_weights_gpt_neox(

    state_dict: Dict[str, torch.Tensor],

    hf_weights: Dict[str, Union[torch.Tensor, NotYetLoadedTensor]],

    saver: Optional[incremental_save] = None,

    dtype: Optional[torch.dtype] = None,

) -> None:

    weight_map = {

        "gpt_neox.embed_in.weight": "transformer.wte.weight",

        "gpt_neox.layers.{}.input_layernorm.bias": "transformer.h.{}.norm_1.bias",

        "gpt_neox.layers.{}.input_layernorm.weight": "transformer.h.{}.norm_1.weight",

        "gpt_neox.layers.{}.attention.query_key_value.bias": "transformer.h.{}.attn.attn.bias",

        "gpt_neox.layers.{}.attention.query_key_value.weight": "transformer.h.{}.attn.attn.weight",

        "gpt_neox.layers.{}.attention.dense.bias": "transformer.h.{}.attn.proj.bias",

        "gpt_neox.layers.{}.attention.dense.weight": "transformer.h.{}.attn.proj.weight",

        "gpt_neox.layers.{}.attention.rotary_emb.inv_freq": None,

        "gpt_neox.layers.{}.attention.bias": None,

        "gpt_neox.layers.{}.attention.masked_bias": None,

        "gpt_neox.layers.{}.post_attention_layernorm.bias": "transformer.h.{}.norm_2.bias",

        "gpt_neox.layers.{}.post_attention_layernorm.weight": "transformer.h.{}.norm_2.weight",

        "gpt_neox.layers.{}.mlp.dense_h_to_4h.bias": "transformer.h.{}.mlp.fc.bias",

        "gpt_neox.layers.{}.mlp.dense_h_to_4h.weight": "transformer.h.{}.mlp.fc.weight",

        "gpt_neox.layers.{}.mlp.dense_4h_to_h.bias": "transformer.h.{}.mlp.proj.bias",

        "gpt_neox.layers.{}.mlp.dense_4h_to_h.weight": "transformer.h.{}.mlp.proj.weight",

        "gpt_neox.final_layer_norm.bias": "transformer.ln_f.bias",

        "gpt_neox.final_layer_norm.weight": "transformer.ln_f.weight",

        "embed_out.weight": "lm_head.weight",

    }

    for name, param in hf_weights.items():

        if "gpt_neox.layers" in name:

            from_name, number = layer_template(name, 2)

            to_name = weight_map[from_name]

            if to_name is None:

                continue

            to_name = to_name.format(number)

        else:

            to_name = weight_map[name]

        param = load_param(param, name, dtype)

        if saver is not None:

            param = saver.store_early(param)

        state_dict[to_name] = param

def copy_weights_falcon(

    model_name: str,

    state_dict: Dict[str, torch.Tensor],

    hf_weights: Dict[str, Union[torch.Tensor, NotYetLoadedTensor]],

    saver: Optional[incremental_save] = None,

    dtype: Optional[torch.dtype] = None,

) -> None:

    weight_map = {

        "transformer.word_embeddings.weight": "transformer.wte.weight",

        "transformer.h.{}.self_attention.query_key_value.weight": "transformer.h.{}.attn.attn.weight",

        "transformer.h.{}.self_attention.dense.weight": "transformer.h.{}.attn.proj.weight",

        "transformer.h.{}.mlp.dense_h_to_4h.weight": "transformer.h.{}.mlp.fc.weight",

        "transformer.h.{}.mlp.dense_4h_to_h.weight": "transformer.h.{}.mlp.proj.weight",

        "transformer.ln_f.bias": "transformer.ln_f.bias",

        "transformer.ln_f.weight": "transformer.ln_f.weight",

        "lm_head.weight": "lm_head.weight",

    }

    # the original model definition is different for each size

    if "7b" in model_name:

        weight_map.update(

            {

                "transformer.h.{}.input_layernorm.bias": "transformer.h.{}.norm_1.bias",

                "transformer.h.{}.input_layernorm.weight": "transformer.h.{}.norm_1.weight",

            }

        )

    elif "40b" in model_name or "180B" in model_name:

        weight_map.update(

            {

                "transformer.h.{}.ln_attn.bias": "transformer.h.{}.norm_1.bias",

                "transformer.h.{}.ln_attn.weight": "transformer.h.{}.norm_1.weight",

                "transformer.h.{}.ln_mlp.bias": "transformer.h.{}.norm_2.bias",

                "transformer.h.{}.ln_mlp.weight": "transformer.h.{}.norm_2.weight",

            }

        )

    else:

        raise NotImplementedError

    for name, param in hf_weights.items():

        if "transformer.h" in name:

            from_name, number = layer_template(name, 2)

            to_name = weight_map[from_name].format(number)

        else:

            to_name = weight_map[name]

        param = load_param(param, name, dtype)

        if saver is not None:

            param = saver.store_early(param)

        state_dict[to_name] = param

def copy_weights_hf_llama(

    config: Config,

    qkv_weights: Dict[int, List[Optional[NotYetLoadedTensor]]],

    state_dict: Dict[str, torch.Tensor],

    hf_weights: Dict[str, Union[torch.Tensor, NotYetLoadedTensor]],

    saver: Optional[incremental_save] = None,

    dtype: Optional[torch.dtype] = None,

) -> None:

    weight_map = {

        "model.embed_tokens.weight": "transformer.wte.weight",

        "model.layers.{}.input_layernorm.weight": "transformer.h.{}.norm_1.weight",

        "model.layers.{}.self_attn.q_proj.weight": None,

        "model.layers.{}.self_attn.k_proj.weight": None,

        "model.layers.{}.self_attn.v_proj.weight": None,

        "model.layers.{}.self_attn.o_proj.weight": "transformer.h.{}.attn.proj.weight",

        "model.layers.{}.self_attn.rotary_emb.inv_freq": None,

        "model.layers.{}.post_attention_layernorm.weight": "transformer.h.{}.norm_2.weight",

        "model.layers.{}.mlp.gate_proj.weight": "transformer.h.{}.mlp.fc_1.weight",

        "model.layers.{}.mlp.up_proj.weight": "transformer.h.{}.mlp.fc_2.weight",

        "model.layers.{}.mlp.down_proj.weight": "transformer.h.{}.mlp.proj.weight",

        "model.norm.weight": "transformer.ln_f.weight",

        "lm_head.weight": "lm_head.weight",

    }

    for name, param in hf_weights.items():

        if "model.layers" in name:

            from_name, number = layer_template(name, 2)

            qkv = qkv_weights.setdefault(number, [None, None, None])

            if "q_proj" in name:

                qkv[0] = param

            elif "k_proj" in name:

                qkv[1] = param

            elif "v_proj" in name:

                qkv[2] = param

            to_name = weight_map[from_name]

            if to_name is None:

                continue

            to_name = to_name.format(number)

        else:

            to_name = weight_map[name]

        param = load_param(param, name, dtype)

        if saver is not None:

            param = saver.store_early(param)

        state_dict[to_name] = param

    for i, (q, k, v) in list(qkv_weights.items()):

        if q is None or k is None or v is None:

            # split across different .bin files

            continue

        q = load_param(q, f"layer {i} q", dtype)

        k = load_param(k, f"layer {i} k", dtype)

        v = load_param(v, f"layer {i} v", dtype)

        q_per_kv = config.n_head // config.n_query_groups

        qs = torch.split(q, config.head_size * q_per_kv)

        ks = torch.split(k, config.head_size)

        vs = torch.split(v, config.head_size)

        cycled = [t for group in zip(qs, ks, vs) for t in group]

        qkv = torch.cat(cycled)

        state_dict[f"transformer.h.{i}.attn.attn.weight"] = qkv

        del qkv_weights[i]

def copy_weights_phi(

    config: Config,

    state_dict: Dict[str, torch.Tensor],

    hf_weights: Dict[str, Union[torch.Tensor, NotYetLoadedTensor]],

    saver: Optional[incremental_save] = None,

    dtype: Optional[torch.dtype] = None,

) -> None:

    if any(layer_name.startswith("layers.") for layer_name in hf_weights):

        raise ValueError(

            "You are using an outdated Phi1.5 checkpoint. "

            "Please reload it as described in 'tutorials/download_phi15.md'"

        )

    weight_map = {

        "transformer.embd.wte.weight": "transformer.wte.weight",

        "transformer.h.{}.ln.bias": "transformer.h.{}.norm_1.bias",

        "transformer.h.{}.ln.weight": "transformer.h.{}.norm_1.weight",

        "transformer.h.{}.mixer.Wqkv.bias": "transformer.h.{}.attn.attn.bias",

        "transformer.h.{}.mixer.Wqkv.weight": "transformer.h.{}.attn.attn.weight",

        "transformer.h.{}.mixer.out_proj.bias": "transformer.h.{}.attn.proj.bias",

        "transformer.h.{}.mixer.out_proj.weight": "transformer.h.{}.attn.proj.weight",

        "transformer.h.{}.mixer.rotary_emb.inv_freq": None,

        "transformer.h.{}.mlp.fc1.bias": "transformer.h.{}.mlp.fc.bias",

        "transformer.h.{}.mlp.fc1.weight": "transformer.h.{}.mlp.fc.weight",

        "transformer.h.{}.mlp.fc2.bias": "transformer.h.{}.mlp.proj.bias",

        "transformer.h.{}.mlp.fc2.weight": "transformer.h.{}.mlp.proj.weight",

        "lm_head.ln.weight": "transformer.ln_f.weight",

        "lm_head.ln.bias": "transformer.ln_f.bias",

        "lm_head.linear.weight": "lm_head.weight",

        "lm_head.linear.bias": "lm_head.bias",

    }

    for name, param in hf_weights.items():

        if name.startswith("transformer.h."):

            from_name, number = layer_template(name, 2)

            to_name = weight_map[from_name].format(number)

        else:

            to_name = weight_map[name]

        param = load_param(param, name, dtype)

        if "Wqkv" in name:

            q_per_kv = config.n_head // config.n_query_groups

            total_qkv = q_per_kv + 2  # each group has 1+ queries, 1 key, and 1 value

            param = param.view(total_qkv, config.n_query_groups, -1).transpose(0, 1)

            param = param.reshape(config.n_embd * 3, -1)

            if "bias" in name:

                param = param.squeeze()

        if saver is not None:

            param = saver.store_early(param)

        state_dict[to_name] = param

def layer_template(layer_name: str, idx: int) -> Tuple[str, int]:

    split = layer_name.split(".")

    number = int(split[idx])

    split[idx] = "{}"

    from_name = ".".join(split)

    return from_name, number

def load_param(param: Union[torch.Tensor, NotYetLoadedTensor], name: str, dtype: Optional[torch.dtype]) -> torch.Tensor:

    if hasattr(param, "_load_tensor"):

        # support tensors loaded via `lazy_load()`

        print(f"Loading {name!r} into RAM")

        param = param._load_tensor()

    if dtype is not None and type(dtype) is not NotYetLoadedTensor and dtype != param.dtype:

        print(f"Converting {name!r} from {param.dtype} to {dtype}")

        param = param.to(dtype)

    return param

@torch.inference_mode()

def convert_hf_checkpoint(

    *,

    checkpoint_dir: Path = Path("checkpoints/stabilityai/stablelm-base-alpha-3b"),

    model_name: Optional[str] = None,

    dtype: Optional[str] = None,

) -> None:

    if model_name is None:

        model_name = checkpoint_dir.name

    if dtype is not None:

        dtype = getattr(torch, dtype)

    config = Config.from_name(model_name)

    config_dict = asdict(config)

    print(f"Model config {config_dict}")

    with open(checkpoint_dir / "lit_config.json", "w") as json_config:

        json.dump(config_dict, json_config)

    if "falcon" in model_name:

        copy_fn = partial(copy_weights_falcon, model_name)

    elif config._mlp_class == "LLaMAMLP":

        # holder to reconstitute the split q, k, v

        qkv_weights = {}

        copy_fn = partial(copy_weights_hf_llama, config, qkv_weights)

    elif "phi" in model_name:

        copy_fn = partial(copy_weights_phi, config)

    else:

        copy_fn = copy_weights_gpt_neox

    # initialize a new empty state dict to hold our new weights

    sd = {}

    # Load the json file containing weight mapping

    pytorch_bin_map_json_path = checkpoint_dir / "pytorch_model.bin.index.json"

    if pytorch_bin_map_json_path.is_file():  # not all checkpoints have this file

        with open(pytorch_bin_map_json_path) as json_map:

            bin_index = json.load(json_map)

        bin_files = {checkpoint_dir / bin for bin in bin_index["weight_map"].values()}

    else:

        bin_files = set(checkpoint_dir.glob("*.bin"))

        # some checkpoints serialize the training arguments

        bin_files = {f for f in bin_files if f.name != "training_args.bin"}

    if not bin_files:

        raise ValueError(f"Expected {str(checkpoint_dir)!r} to contain .bin files")

    with incremental_save(checkpoint_dir / "lit_model.pth") as saver:

        # for checkpoints that split the QKV across several files, we need to keep all the bin files

        # open, so we use `ExitStack` to close them all together at the end

        for bin_file in sorted(bin_files):

            print("Processing", bin_file)

            hf_weights = lazy_load(bin_file)

            copy_fn(sd, hf_weights, saver=saver, dtype=dtype)

        gc.collect()

        print("Saving converted checkpoint")

        saver.save(sd)

if __name__ == "__main__":

    from jsonargparse import CLI

    CLI(convert_hf_checkpoint)