tools/ckpts/convert_neox_to_hf.py

# Copyright (c) 2023, EleutherAI
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import sys

import yaml
import argparse
from tqdm import tqdm

import torch
from transformers import (
    MistralConfig,
    LlamaConfig,
    GPTNeoXConfig,
    AutoModelForCausalLM,
    AutoConfig,
    AutoModelForSequenceClassification,
)

from typing import List, Literal

sys.path.append(
    os.path.abspath(
        os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir)
    )
)
from megatron.tokenizer import build_tokenizer


"""
A script for converting saved NeoX Checkpoints to Huggingface (HF) compatible GPT-NeoX type models.

Note that this script does not support all NeoX features.
Please investigate carefully whether your model is compatible with all architectures supported by the GPTNeoXForCausalLM class in HF.

(e.g. position embeddings such as AliBi may not be supported by Huggingface's GPT-NeoX architecture).
"""


# Model definitions: a list of keys, and where they fall in terms of handling them in the presence of TP.
# in format : {model arch: {param type: {param in neox: param in HF}}}
MODEL_KEYS = {
    "neox": {
        "new": {
            "COLUMN_PARALLEL_LINEAR_KEYS": {
                "mlp.linear1.weight": "mlp.dense_h_to_4h.weight",
                "mlp.linear1.bias": "mlp.dense_h_to_4h.bias",
                "attention.query_key_value.weight": "attention.query_key_value.weight",
                "attention.query_key_value.bias": "attention.query_key_value.bias",  # TODO: handle GQA separately?
            },
            "ROW_PARALLEL_LINEAR_KEYS": {
                "attention.dense.weight": "attention.dense.weight",
                "mlp.linear2.weight": "mlp.dense_4h_to_h.weight",
            },
            "ROW_PARALLEL_BIAS_KEYS": {
                "mlp.linear2.bias": "mlp.dense_4h_to_h.bias",
                "attention.dense.bias": "attention.dense.bias",
            },
            "NORM_KEYS": {
                "input_layernorm.weight": "input_layernorm.weight",
                "input_layernorm.bias": "input_layernorm.bias",
                "post_attention_layernorm.weight": "post_attention_layernorm.weight",
                "post_attention_layernorm.bias": "post_attention_layernorm.bias",
            },
            "FINAL_NORM_KEYS": {
                "norm.weight": "weight",
                "norm.bias": "bias",
            },
        },
        "legacy": {
            "COLUMN_PARALLEL_LINEAR_KEYS": {
                "mlp.dense_h_to_4h.weight": "mlp.dense_h_to_4h.weight",
                "mlp.dense_h_to_4h.bias": "mlp.dense_h_to_4h.bias",
                "attention.query_key_value.weight": "attention.query_key_value.weight",
                "attention.query_key_value.bias": "attention.query_key_value.bias",  # TODO: handle GQA separately?
            },
            "ROW_PARALLEL_LINEAR_KEYS": {
                "attention.dense.weight": "attention.dense.weight",
                "mlp.dense_4h_to_h.weight": "mlp.dense_4h_to_h.weight",
            },
            "ROW_PARALLEL_BIAS_KEYS": {
                "mlp.dense_4h_to_h.bias": "mlp.dense_4h_to_h.bias",
                "attention.dense.bias": "attention.dense.bias",
            },
            "NORM_KEYS": {
                "input_layernorm.weight": "input_layernorm.weight",
                "input_layernorm.bias": "input_layernorm.bias",
                "post_attention_layernorm.weight": "post_attention_layernorm.weight",
                "post_attention_layernorm.bias": "post_attention_layernorm.bias",
            },
            "FINAL_NORM_KEYS": {
                "norm.weight": "weight",
                "norm.bias": "bias",
            },
        },
    },
    "llama": {
        "new": {
            "COLUMN_PARALLEL_LINEAR_KEYS": {
                "mlp.linear1.weight": ["mlp.up_proj.weight", "mlp.gate_proj.weight"]
            },
            "ROW_PARALLEL_LINEAR_KEYS": {
                "attention.dense.weight": "self_attn.o_proj.weight",
                "mlp.linear2.weight": "mlp.down_proj.weight",
            },
            "ROW_PARALLEL_BIAS_KEYS": {},  # No biases in RowParallelLinear layers
            "NORM_KEYS": {
                "input_layernorm.scale": "input_layernorm.weight",
                "post_attention_layernorm.scale": "post_attention_layernorm.weight",
            },
            "FINAL_NORM_KEYS": {
                "norm.scale": "weight",
            },
            "GQA_QKV_KEYS": {  # because Llama can have Grouped Query Attention and has separate Q, K, and V linear proj params, handle them separately.
                "attention.query_key_value.weight": [
                    "self_attn.q_proj.weight",
                    "self_attn.k_proj.weight",
                    "self_attn.v_proj.weight",
                ],
            },
        },
        "legacy": {
            "COLUMN_PARALLEL_LINEAR_KEYS": {
                "mlp.w1.weight": "mlp.gate_proj.weight",
                "mlp.w3.weight": "mlp.up_proj.weight",
            },
            "ROW_PARALLEL_LINEAR_KEYS": {
                "attention.dense.weight": "self_attn.o_proj.weight",
                "mlp.w2.weight": "mlp.down_proj.weight",
            },
            "ROW_PARALLEL_BIAS_KEYS": {},  # No biases in RowParallelLinear layers
            "NORM_KEYS": {
                "input_layernorm.scale": "input_layernorm.weight",
                "post_attention_layernorm.scale": "post_attention_layernorm.weight",
            },
            "FINAL_NORM_KEYS": {
                "norm.scale": "weight",
            },
            "GQA_QKV_KEYS": {  # because Llama can have Grouped Query Attention and has separate Q, K, and V linear proj params, handle them separately.
                "attention.query_key_value.weight": [
                    "self_attn.q_proj.weight",
                    "self_attn.k_proj.weight",
                    "self_attn.v_proj.weight",
                ],
            },
        },
    },
}

MODEL_KEYS["mistral"] = MODEL_KEYS["llama"]


def load_partitions(
    input_checkpoint_path: str, mp_partitions: int, layer_idx: int, sequential: bool
) -> List[torch.Tensor]:
    """Returns a list containing all states from a model (across MP partitions)"""

    if sequential:
        filename_format = f"mp_rank_{{i:02}}_model_states.pt"
    else:
        filename_format = f"layer_{layer_idx:02}-model_{{i:02}}-model_states.pt"

    loaded_tp_ranks = [
        torch.load(
            os.path.join(
                input_checkpoint_path,
                filename_format.format(i=i),
            ),
            map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
        )
        for i in range(mp_partitions)
    ]

    return loaded_tp_ranks


def get_state(
    state_dicts: List[torch.Tensor], key: str, layer_idx: int, sequential: bool
) -> torch.Tensor:
    """Helper that returns a list containing a given weight's state from each MP partition, for a given layer in the model."""

    if sequential:
        # use the correct key into the sequential dict for given weight/provided key
        key = f"sequential.{layer_idx}.{key}"

        return [state_dict["module"][key] for state_dict in state_dicts]
    else:
        # For the PipelineModule case, we don't need any key / module prefix. just grab this weight value.
        # layer_idx is also ignored because we've loaded only this layer's weights, ahead of time.
        key = key

        return [state_dict[key] for state_dict in state_dicts]


def get_key(loaded_config, key, default=None):
    """
    Search for a given key in a NeoX yaml. normalizes underscores -> hyphens
    """
    key = key.replace("_", "-")
    try:
        return loaded_config[key]
    except KeyError:
        key = key.replace("-", "_")
        try:
            return loaded_config[key]
        except KeyError:
            return default


def create_config(neox_config, architecture="neox", is_rm=False, pad_token_id=-1):
    """take in a loaded yaml from NeoX and assign relevant values to HF config.
    Returns: GPTNeoXConfig() object
    """

    def gated_size(hidden_dim):
        # takes in a hidden dim and calculates intermediate dim of a LLaMAParallelMLP.
        # (only used if intermediate_size not specified in config)
        # hidden-size * 8 / 3 , rounded up to nearest multiple of 256
        ff_dim = int(2 * hidden_dim * 4 / 3)
        ff_dim = 256 * ((ff_dim + 256 - 1) // 256)
        return ff_dim

    class TokenizerArgs:
        # kinda hacky.
        # this is to get something with the same interface as is used in build_tokenizer()
        # without diving into loading a neox_args object or using argparse etc.
        def __init__(self, neox_config):
            self.make_vocab_size_divisible_by = get_key(
                neox_config, "make-vocab-size-divisible-by", default=128
            )
            self.model_parallel_size = get_key(neox_config, "model-parallel-size")
            self.vocab_file = get_key(neox_config, "vocab-file")
            self.merge_file = get_key(neox_config, "merge-file")
            self.tokenizer_type = get_key(neox_config, "tokenizer-type")

            self.rank = 0

    args = TokenizerArgs(neox_config)
    tokenizer = build_tokenizer(args)
    try:  # GPT2TokenizerFast raises NotImplementedError
        pad_token = tokenizer.pad
    except:
        pad_token = (
            1  # pad defaulting to 1. follows convention from GPT-NeoX-20b tokenizer
        )

    # TODO: change the default value here based on discussion regarding `gpt_j_tied` config parameter's default
    use_tied_lns = get_key(neox_config, "gpt-j-tied", False)

    if use_tied_lns:
        raise NotImplementedError(
            """ERROR: Huggingface Transformers does not yet support a single shared layernorm
                per transformer block for GPT-NeoX models trained  w/ GPT-J parallel residuals.
                See https://github.com/EleutherAI/gpt-neox/pull/481 for further details."""
        )

    # set all config values.

    # shared config parameters.
    args = {
        "vocab_size": args.padded_vocab_size,
        "hidden_size": get_key(neox_config, "hidden-size"),
        "num_hidden_layers": get_key(neox_config, "num-layers"),
        "num_attention_heads": get_key(neox_config, "num-attention-heads"),
        "max_position_embeddings": get_key(neox_config, "max-position-embeddings"),
        "initializer_range": get_key(neox_config, "init-method-std", 0.02),
        "tie_word_embeddings": (not get_key(neox_config, "no-weight-tying", False)),
        "use_cache": True,
    }
    if architecture == "mistral" or architecture == "llama":
        args.update(
            {
                "intermediate_size": get_key(
                    neox_config,
                    "intermediate-size",
                    gated_size(get_key(neox_config, "hidden-size")),
                ),
                "num_key_value_heads": get_key(
                    neox_config,
                    "num-kv-heads",
                    get_key(neox_config, "num-attention-heads"),
                ),
                "hidden_act": get_key(
                    neox_config, "activation", default="silu"
                ).replace("swiglu", "silu"),
                "rms_norm_eps": get_key(neox_config, "rms-norm-epsilon", 1.0e-6),
                "bos_token_id": tokenizer.eod,
                "eos_token_id": tokenizer.eod,
                "rope_theta": get_key(neox_config, "rotary-emb-base", 10000.0),
            }
        )

        if architecture == "mistral":
            # mistral-specific options
            args.update(
                {
                    "sliding_window": get_key(
                        neox_config, "sliding-window-width", 4096
                    ),
                }
            )
            hf_config = MistralConfig(**args)
        elif architecture == "llama":
            # llama-specific options
            args.update(
                {
                    # NeoX library defaults to using bias in attention
                    "attention_bias": get_key(
                        neox_config, "use_bias_in_attn_linear", True
                    ),
                }
            )
            hf_config = LlamaConfig(**args)
    else:
        # GPT-NeoX HF model class-specific options
        args.update(
            {
                "rotary_pct": get_key(neox_config, "rotary-pct", default=1.0),
                "rotary_emb_base": get_key(
                    neox_config, "rotary-emb-base", default=1000.0
                ),
                "use_parallel_residual": get_key(neox_config, "gpt-j-residual", False),
                "layer_norm_eps": get_key(neox_config, "layernorm-epsilon", 1e-5),
                "intermediate_size": get_key(
                    neox_config,
                    "intermediate-size",
                    4 * get_key(neox_config, "hidden-size"),
                ),
            }
        )
        hf_config = GPTNeoXConfig(**args)
    if is_rm:
        hf_config.num_labels = 1
        hf_config.pad_token_id = pad_token_id

    return hf_config


def reshard_and_split_qkv(
    param_mapping: dict,  # a dictionary mapping the QKV weight keys in GPT-NeoX -> a list of keys representing the Q, K, and V weight keys the HF model will use
    hf_config: AutoConfig,  # a HF model config for the model
    loaded_tp_ranks: List[torch.Tensor],
    layer_idx: int,
    sequential: bool,
):
    """
    A helper function which performs reshaping and sharding to make the QKV projection from NeoX compatible with HF Llama models,
    even when grouped-query attention is required.
    """
    for key, hf_keys in param_mapping.items():
        assert (
            isinstance(hf_keys, list) and len(hf_keys) == 3
        ), "Must map QKV to precisely 3 resulting weight matrices."

    for key, hf_keys in param_mapping.items():
        # we first merge the QKV proj. across TP ranks
        sharded_qkv = torch.stack(
            get_state(loaded_tp_ranks, key, layer_idx, sequential), dim=0
        )
        # should now have shape [TP_SIZE, (hidden_size + 2 * kv_hidden_size) / TP_SIZE, hidden_size].

        sharded_qkv = sharded_qkv.view(
            len(loaded_tp_ranks),
            hf_config.num_attention_heads // len(loaded_tp_ranks),
            int(
                hf_config.hidden_size
                // hf_config.num_attention_heads
                * (
                    1
                    + 2 * hf_config.num_key_value_heads / hf_config.num_attention_heads
                )
            ),
            hf_config.hidden_size,
        )  # is meant to convert to shape [TP_SIZE, NUM_QUERY_HEADS_PER_SHARD, dims_per_head * (1 + 2 * kv-to-q head ratio), hidden_size]

        q, k, v = torch.split(
            sharded_qkv,
            [
                hf_config.hidden_size // hf_config.num_attention_heads,
                int(
                    (hf_config.num_key_value_heads / hf_config.num_attention_heads)
                    * hf_config.hidden_size
                    // hf_config.num_attention_heads
                ),
                int(
                    (hf_config.num_key_value_heads / hf_config.num_attention_heads)
                    * hf_config.hidden_size
                    // hf_config.num_attention_heads
                ),
            ],
            dim=2,
        )
        # splits along the (dims_per_head * (1 + 2 * kv-to-q head ratio)_ dim to get 3 tensors:
        # 1 x [TP_SIZE, NUM_Q_HEADS_PER_SHARD, dims_per_head, hidden_size] and 2 x [TP_SIZE, NUM_Q_HEADS_PER_SHARD, (dims_per_head / kv-to-q head ratio), hidden_size]
        # these are the Q, and K, V tensors respectively.

        # we have to do additional reshape for each individual tensor now,
        # into the expected square (or smaller than square, for K/V tensors) shape
        q, k, v = q.squeeze(dim=2), k.squeeze(dim=2), v.squeeze(dim=2)
        q = q.view(
            hf_config.num_attention_heads,
            hf_config.hidden_size // hf_config.num_attention_heads,
            hf_config.hidden_size,
        ).reshape(hf_config.hidden_size, hf_config.hidden_size)
        k = k.reshape(
            hf_config.num_key_value_heads,
            hf_config.hidden_size // hf_config.num_attention_heads,
            hf_config.hidden_size,
        ).reshape(
            hf_config.hidden_size
            // hf_config.num_attention_heads
            * hf_config.num_key_value_heads,
            hf_config.hidden_size,
        )
        v = v.reshape(
            hf_config.num_key_value_heads,
            hf_config.hidden_size // hf_config.num_attention_heads,
            hf_config.hidden_size,
        ).reshape(
            hf_config.hidden_size
            // hf_config.num_attention_heads
            * hf_config.num_key_value_heads,
            hf_config.hidden_size,
        )

        # return these
        state_dict = {}
        for hf_key, proj in zip(hf_keys, [q, k, v]):
            state_dict[hf_key] = proj.clone()
        return state_dict


def get_mlp_naming_convention(loaded_tp_ranks, layer_idx, sequential):
    """Determine whether the checkpoint uses the legacy or new MLP naming convention."""
    print(list(loaded_tp_ranks[0]["module"].keys()))
    if any(
        [
            ["mlp.linear1.weight" in key for key in list(state_dict["module"].keys())]
            for state_dict in loaded_tp_ranks
        ]
    ):
        return "new"
    elif any(
        [
            [
                "mlp.dense_h_to_4h.weight" in key
                for key in list(state_dict["module"].keys())
            ]
            for state_dict in loaded_tp_ranks
        ]
    ):
        return "legacy"
    else:
        raise ValueError("Unable to determine MLP naming convention in checkpoint")


def convert(
    input_checkpoint_path,
    loaded_config,
    output_checkpoint_path,
    sequential: bool = True,
    precision: Literal["auto", "fp16", "bf16", "fp32"] = "auto",
    architecture: Literal["neox", "llama", "mistral"] = "neox",
    is_rm: bool = False,
    pad_token_id: int = -1,
):
    """convert a NeoX checkpoint to a HF model format.
    should perform model-parallel merging correctly
    but only supports features allowed by HF GPT-NeoX implementation (e.g. rotary embeddings)
    """

    ARCH = MODEL_KEYS[architecture]

    hf_config = create_config(
        loaded_config, architecture=architecture, is_rm=is_rm, pad_token_id=pad_token_id
    )

    if not is_rm:
        hf_model = AutoModelForCausalLM.from_config(hf_config)
    else:
        hf_model = AutoModelForSequenceClassification.from_config(hf_config)

    if architecture == "neox":
        hf_transformer = hf_model.gpt_neox
    else:
        hf_transformer = hf_model.model

    if precision == "auto":
        print("Auto-detecting precision to save model into...")
        # save model in FP16 if Deepspeed fp16 was used in config, else 32 bit
        fp16 = get_key(loaded_config, "fp16")

        if fp16:
            try:
                # current behavior is to pass "fp16": {"enabled": true}, when using upstream Deepspeed
                if fp16["enabled"]:
                    hf_model.half()
                    print("Saving weights in fp16 precision...")
            except:
                try:
                    # attempt to access bf16 dict in yaml file, if fp16 not enabled
                    bf16 = get_key(loaded_config, "bf16")
                    if bf16:
                        hf_model.to(dtype=torch.bfloat16)
                        print("Saving weights in bf16 precision...")
                except:
                    hf_model.to(dtype=torch.float)
                    print(
                        "Model not trained in fp16 / bf16 mixed precision, saving weights in fp32..."
                    )
    else:
        name_to_dtype = {
            "bf16": torch.bfloat16,
            "fp16": torch.float16,
            "fp32": torch.float,
        }
        print(f"Saving model into specified {precision} precision...")
        hf_model.to(dtype=name_to_dtype[precision])

    mp_partitions = get_key(loaded_config, "model-parallel-size")

    # Sequential saves all model states from an MP rank in one file.
    # so we only load the MP ranks only once and index into them with get_state().
    # for the pipeline-parallel case (pipeline-parallel-size >= 1),
    # we must load the correct layer's states at each step.
    # (this does mean that less memory is required for PP conversion.)
    loaded_tp_ranks = load_partitions(
        input_checkpoint_path, mp_partitions, layer_idx=0, sequential=sequential
    )

    ### Embedding layer ###
    # Embedding is layer idx 0
    if architecture == "neox":
        embed_in = hf_transformer.embed_in
    else:
        embed_in = hf_transformer.embed_tokens
    embed_in.load_state_dict(  # TODO: embed_in is not always model's name for embedding
        {
            "weight": torch.cat(
                get_state(
                    loaded_tp_ranks,
                    "word_embeddings.weight",
                    layer_idx=0,
                    sequential=sequential,
                ),
                dim=0,
            )
        }
    )
    assert (
        hf_config.vocab_size == embed_in.weight.shape[0]
    ), f"ERROR: calculated vocab size {hf_config.vocab_size} != embed param size {embed_in.shape[0]}"
    ### End Embedding Layer ###

    # grab from 3rd layer to pass embeddings
    mlp_naming = get_mlp_naming_convention(
        load_partitions(
            input_checkpoint_path,
            mp_partitions,
            layer_idx=3,
            sequential=sequential,
        ),
        0,
        sequential,
    )
    print(f"Detected MLP naming convention: {mlp_naming}")
    ARCH = ARCH[mlp_naming]

    for layer_i in tqdm(range(get_key(loaded_config, "num-layers"))):

        # get layer from hf model
        hf_layer = hf_transformer.layers[layer_i]  # TODO: model module names

        if not sequential:
            # in the non-sequential case, must load from each layer individually.
            # use layer index + 2 bc of embed layer and a dummy _pre_transformer_block, which are "layers 0 and 1"
            loaded_tp_ranks = load_partitions(
                input_checkpoint_path,
                mp_partitions,
                layer_idx=layer_i + 2,
                sequential=sequential,
            )

        # + 2 bc of embed layer and a dummy _pre_transformer_block
        state_dict = {}
        for key, hf_key in ARCH["ROW_PARALLEL_LINEAR_KEYS"].items():
            state_dict[hf_key] = torch.cat(
                get_state(
                    loaded_tp_ranks, key, layer_idx=layer_i + 2, sequential=sequential
                ),
                dim=1,
            )

        # average layernorm stats over mp ranks
        for key, hf_key in ARCH["NORM_KEYS"].items():
            state_dict[hf_key] = sum(
                get_state(
                    loaded_tp_ranks, key, layer_idx=layer_i + 2, sequential=sequential
                )
            ) / len(loaded_tp_ranks)

        # LinearWithTPMerge
        for key, hf_key in ARCH["COLUMN_PARALLEL_LINEAR_KEYS"].items():
            if type(hf_key) == list:
                # Llama magic - split the weight into two parts for the gate and up proj
                states = [
                    torch.chunk(state, chunks=2, dim=0)
                    for state in get_state(
                        loaded_tp_ranks,
                        key,
                        layer_idx=layer_i + 2,
                        sequential=sequential,
                    )
                ]
                # Set up proj...
                state_dict[hf_key[0]] = torch.cat([state[0] for state in states], dim=0)
                # Set gate proj...
                state_dict[hf_key[1]] = torch.cat([state[1] for state in states], dim=0)
            else:
                state_dict[hf_key] = torch.cat(
                    get_state(
                        loaded_tp_ranks,
                        key,
                        layer_idx=layer_i + 2,
                        sequential=sequential,
                    ),
                    dim=0,
                )

        # LinearWithTPSplitBias
        for key, hf_key in ARCH["ROW_PARALLEL_BIAS_KEYS"].items():
            state_dict[hf_key] = sum(
                get_state(
                    loaded_tp_ranks, key, layer_idx=layer_i + 2, sequential=sequential
                )
            )

        # Just take one
        if "attention.bias" in hf_layer.state_dict():
            state_dict["attention.bias"] = hf_layer.state_dict()["attention.bias"]
        if "attention.masked_bias" in hf_layer.state_dict():
            state_dict["attention.masked_bias"] = hf_layer.state_dict()[
                "attention.masked_bias"
            ]

        # some architectures, like Mistral and Llama, have the following which must be handled specially:
        # - Q, K, V projections are performed separately, so we must split apart GPT-NeoX library's single QKV proj
        # - Support for Grouped-Query Attention, meaning the Q and the K, V projections may not be the same size
        if "GQA_QKV_KEYS" in ARCH:
            state_dict.update(
                reshard_and_split_qkv(
                    param_mapping=ARCH["GQA_QKV_KEYS"],
                    hf_config=hf_config,
                    loaded_tp_ranks=loaded_tp_ranks,
                    layer_idx=layer_i + 2,
                    sequential=sequential,
                )
            )
        # load state_dict into layer
        hf_layer.load_state_dict(state_dict)

    if not sequential:
        loaded_tp_ranks = load_partitions(
            input_checkpoint_path,
            mp_partitions,
            get_key(loaded_config, "num-layers") + 3,
            sequential=sequential,
        )
    # Load final layer norm
    norm_state_dict = {}
    for key, hf_key in ARCH["FINAL_NORM_KEYS"].items():
        norm_state_dict[hf_key] = sum(
            get_state(
                loaded_tp_ranks,
                key,
                layer_idx=get_key(loaded_config, "num-layers") + 3,
                sequential=sequential,
            )
        ) / len(loaded_tp_ranks)

    if architecture == "neox":
        final_layer_norm = hf_transformer.final_layer_norm
    else:
        final_layer_norm = hf_transformer.norm

    final_layer_norm.load_state_dict(norm_state_dict)

    # Load output embedding
    if not sequential:
        if get_key(loaded_config, "no-weight-tying", False):
            # if we have trained input + output embedding layers without tied weights
            loaded_tp_ranks = load_partitions(
                input_checkpoint_path,
                mp_partitions,
                get_key(loaded_config, "num-layers") + 4,
                sequential=sequential,
            )
        else:
            # in this case, output embedding layer and input embedding layer are tied.
            # load + save the input embed weights into the output embedding layer's place.
            loaded_tp_ranks = load_partitions(
                input_checkpoint_path,
                mp_partitions,
                layer_idx=0,
                sequential=sequential,
            )
    # output embedding / LM head
    if not is_rm:
        if architecture == "neox":  # name of lm head / final linear proj varies
            lm_head = hf_model.embed_out
        else:
            lm_head = hf_model.lm_head
    else:
        lm_head = hf_model.score

    if get_key(loaded_config, "no-weight-tying", False):
        # save the (untied) final linear into LM head for HF
        lm_head.load_state_dict(
            {
                "weight": torch.cat(
                    get_state(
                        loaded_tp_ranks,
                        "final_linear.weight" if not is_rm else "rm_linear.weight",
                        layer_idx=get_key(loaded_config, "num-layers") + 4,
                        sequential=sequential,
                    ),
                    dim=0 if not is_rm else 1,
                ),
            }
        )
    else:
        # don't need to worry about rm here since you can't really tie them...

        # embedding layers are tied. transpose input layer and save
        lm_head.load_state_dict(
            {
                "weight": torch.cat(
                    get_state(
                        loaded_tp_ranks,
                        "word_embeddings.weight",
                        layer_idx=0,
                        sequential=sequential,
                    ),
                    dim=0,
                ),
            }
        )

    del loaded_tp_ranks

    return hf_model


def main(input_args=None, overwrite_values=None):
    from huggingface_hub import create_repo, HfApi

    parser = argparse.ArgumentParser(
        description="Merge MP partitions and convert to HF Model."
    )
    parser.add_argument(
        "--input_dir",
        type=str,
        help="Path to NeoX checkpoint, e.g. /path/to/model/global_step143000",
    )
    parser.add_argument(
        "--config_file",
        type=str,
        help="Path to config file for the input NeoX checkpoint.",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        help="Output dir, where to save the HF Model, tokenizer, and configs",
    )
    parser.add_argument(
        "--precision",
        type=str,
        default="auto",
        help="What precision to save the model into. Defaults to auto, which auto-detects which 16-bit dtype to save into, or falls back to fp32.",
    )
    parser.add_argument(
        "--no_save_tokenizer",
        action="store_true",
        help="Whether to skip saving the tokenizer alongside a model.",
    )
    parser.add_argument(
        "--vocab-is-hf-tokenizer",
        action="store_true",
        help="Whether the vocab file is in a Huggingface tokenizer path.",
    )
    parser.add_argument(
        "--pad-token-id",
        type=int,
        default=-1,
        help="Pad token id to set in tokenizer. Required for RM style models.",
    )
    parser.add_argument(
        "--architecture",
        type=str,
        default="neox",
        help="What HF model class type to export into.",
    )
    args = parser.parse_args(input_args)

    # validate arguments
    assert args.precision in [
        "auto",
        "fp16",
        "bf16",
        "fp32",
    ], f"expected --precision to be one of 'auto', 'fp16', 'bf16', 'fp32' but got '{args.precision}' !"
    assert args.architecture in [
        "neox",
        "llama",
        "mistral",
    ], f"expected --architecture to be one of 'neox', 'mistral', 'llama', but got '{args.architecture}' !"

    with open(args.config_file) as f:
        loaded_config = yaml.full_load(f)
        if overwrite_values:
            loaded_config.update(overwrite_values)

    # Determine the checkpoint format of the model.
    # DeepSpeed saves models wrapped in a PipelineModule differently from those not.
    # PipelineModule models are saved as per-layer state dicts per TP shard,
    # while Sequential model state dicts are saved all together in one mp_rank_xx_model_states.pt
    # file per tensor/model parallel shard.
    pipeline_world_size = get_key(loaded_config, "pipe-parallel-size", 1)
    is_rm = get_key(loaded_config, "train_impl", "normal") == "rm"
    if is_rm and args.pad_token_id == -1:
        raise ValueError("RM models require a pad token id to be set.")
    if pipeline_world_size == 0:
        sequential = True
        print(
            f"Detected 'pipe-parallel-size' of {pipeline_world_size}, assuming model is saved as Sequential..."
        )
    else:
        sequential = False
        print(
            f"Detected 'pipe-parallel-size' of {pipeline_world_size}, assuming model is saved as PipelineModule..."
        )

    # convert the model to HF.
    hf_model = convert(
        args.input_dir,
        loaded_config,
        args.output_dir,
        sequential=sequential,
        architecture=args.architecture,
        is_rm=is_rm,
        pad_token_id=args.pad_token_id,
    )

    # Save to disk.
    hf_model.save_pretrained(args.output_dir)

    if not args.no_save_tokenizer:
        # save tokenizer to directory as well, for easy loading of model as a HF model.
        tokenizer_type = get_key(loaded_config, "tokenizer-type")
        if args.vocab_is_hf_tokenizer:
            from transformers import AutoTokenizer

            tokenizer = AutoTokenizer.from_pretrained(
                os.path.dirname(get_key(loaded_config, "vocab-file"))
            )
            if args.pad_token_id != -1:
                tokenizer.pad_token_id = args.pad_token_id
            print("loaded tokenizer: ", tokenizer)
            tokenizer.save_pretrained(args.output_dir)
            print("tokenizer saved!")
        elif tokenizer_type == "HFTokenizer":  # TODO: handle sentencepiece tokenizers?
            print(f"saving tokenizer from file {get_key(loaded_config, 'vocab-file')}")
            print(
                "Warning: please check that your model config and tokenizer end with the correct special tokens (EOS, BOS)."
            )
            from transformers import PreTrainedTokenizerFast

            tokenizer = PreTrainedTokenizerFast(
                tokenizer_file=get_key(loaded_config, "vocab-file")
            )
            if args.pad_token_id != -1:
                tokenizer.pad_token_id = args.pad_token_id
            print("loaded tokenizer: ", tokenizer)
            tokenizer.save_pretrained(args.output_dir)
            print("tokenizer saved!")


if __name__ == "__main__":

    # before running script:
    # `pip install --upgrade transformers`
    # `huggingface-cli login`
    #
    main()