ConvertWeightToOpmx.py

# Copyright 2022 EleutherAI and The HuggingFace Inc. team. All rights reserved.
#
# 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 argparse
import gc
import json
import os
import shutil
import warnings

import torch

from pathlib import Path

"""
Sample usage:

```
python convert_hf_weights_to_pmx.py \
    --input_dir /path/to/downloaded/hf/weights/7B --output_dir /output/path
```

Thereafter, models can be loaded via:

```
Important note: you need to be able to host the whole model in RAM to execute this script (even if the biggest versions
come in several checkpoints they each contain a part of each weight of the model, so we need to load them all in RAM).
"""

def read_json(path):
    with open(path, "r") as f:
        return json.load(f)


def write_json(text, path):
    with open(path, "w") as f:
        json.dump(text, f)


def write_pmx_model(model_path, input_base_path):
    os.makedirs(model_path, exist_ok=True)
    print("Loading the checkpoint in a HF model")

    # convert opmx params
    pmx_params_dict = {}
    params = read_json((os.path.join(input_base_path, "config.json")))
    pmx_params_dict['hidden_dim'] = params['hidden_size']
    pmx_params_dict['num_heads'] = params['num_attention_heads']
    pmx_params_dict['num_layers'] = params['num_layers']
    pmx_params_dict['norm_eps'] = params['layernorm_epsilon']
    pmx_params_dict['vocab_size'] = params['padded_vocab_size']
    pmx_params_dict['num_kv_heads'] = params['multi_query_group_num']
    pmx_params_dict['intermediate_dim'] = params["ffn_hidden_size"]
    rope_ratio = params.get('rope_ratio', 1.0)
    pmx_params_dict['rope_theta'] = 10000.0 * rope_ratio
    
    write_json(pmx_params_dict, os.path.join(model_path, "opmx_params.json"))
    print(pmx_params_dict)
    
    hidden_dim = pmx_params_dict['hidden_dim']
    num_heads = pmx_params_dict['num_heads']
    num_kv_heads = pmx_params_dict['num_kv_heads']
    dims_per_head = hidden_dim // num_heads
    key_value_dim = pmx_params_dict['num_kv_heads'] * dims_per_head
    
    # load weights
    def unpermute_weight(wqkv, n_heads, num_kv_heads, hidden_dim, key_value_dim):
        # wqkv: [hidden_dim + kv_dim + kv_dim, hidden_dim]
        wq, wk, wv = wqkv.split([hidden_dim, key_value_dim, key_value_dim], dim=0)
        
        # 由于chatglm2使用的是复数形式的rope，因此此处不需要转置。下方注释为任意rotary dim时的实数转复数转换代码，由于较为复杂因此将其保留，可之后复用于其他模型
        # head_dim = hidden_dim // n_heads
        # rotary_dim = head_dim // 2
        
        # wq_rot = wq.reshape(n_heads, head_dim, hidden_dim)[:, : rotary_dim, :]
        # wq_pass = wq.reshape(n_heads, head_dim, hidden_dim)[:, rotary_dim: , :]
        # wq_rot = wq_rot.reshape(n_heads, 2, rotary_dim // 2, hidden_dim).transpose(1, 2).reshape(n_heads, rotary_dim, hidden_dim)
        # wq = torch.cat([wq_rot, wq_pass], dim=1).reshape(n_heads * head_dim, hidden_dim)
        
        # wk_rot = wk.reshape(num_kv_heads, head_dim, hidden_dim)[:, : rotary_dim, :]
        # wk_pass = wk.reshape(num_kv_heads, head_dim, hidden_dim)[:, rotary_dim: , :]
        # wk_rot = wk_rot.reshape(num_kv_heads, 2, rotary_dim // 2, hidden_dim).transpose(1, 2).reshape(num_kv_heads, rotary_dim, hidden_dim)
        # wk = torch.cat([wk_rot, wk_pass], dim=1).reshape(num_kv_heads * head_dim, hidden_dim)

        return wq, wk, wv

    def unpermute_bias(wqkv_bias, n_heads, num_kv_heads, hidden_dim, key_value_dim):
        wq_bias, wk_bias, wv_bias = wqkv_bias.split([hidden_dim, key_value_dim, key_value_dim])
        
        # # 同unpermute_weight
        # head_dim = hidden_dim // n_heads
        # rotary_dim = head_dim // 2

        # wq_bias0 = wq_bias.reshape(n_heads, head_dim)[:, : rotary_dim]
        # wq_bias1 = wq_bias.reshape(n_heads, head_dim)[:, rotary_dim:]
        # wq_bias0 = wq_bias0.reshape(n_heads, 2, rotary_dim // 2).transpose(1, 2).reshape(n_heads, rotary_dim)
        # wq_bias = torch.cat([wq_bias0, wq_bias1], dim=1).reshape(n_heads * head_dim)
        
        # wk_bias0 = wk_bias.reshape(num_kv_heads, head_dim)[:, : rotary_dim]
        # wk_bias1 = wk_bias.reshape(num_kv_heads, head_dim)[:, rotary_dim:]
        # wk_bias0 = wk_bias0.reshape(num_kv_heads, 2, rotary_dim // 2).transpose(1, 2).reshape(num_kv_heads, rotary_dim)
        # wk_bias = torch.cat([wk_bias0, wk_bias1], dim=1).reshape(num_kv_heads * head_dim)
        
        return wq_bias, wk_bias, wv_bias
        
    hf_model_state_dict, state_dict = {}, {}
    for ckpt_path in sorted(Path(input_base_path).glob("*.bin")):
        hf_model_state_dict.update(torch.load(ckpt_path, map_location="cpu"))
        
    for layer_i in range(pmx_params_dict['num_layers']):
        wq, wk, wv = unpermute_weight(hf_model_state_dict[f"transformer.encoder.layers.{layer_i}.self_attention.query_key_value.weight"], num_heads, num_kv_heads, hidden_dim, key_value_dim)
        
        wq_bias, wk_bias, wv_bias = unpermute_bias(hf_model_state_dict[f"transformer.encoder.layers.{layer_i}.self_attention.query_key_value.bias"], num_heads, num_kv_heads, hidden_dim, key_value_dim)

        w1, w3 = hf_model_state_dict[f"transformer.encoder.layers.{layer_i}.mlp.dense_h_to_4h.weight"].chunk(2, dim=0)
        
        state_dict.update({
            f"layers.{layer_i}.attention.wq.weight": wq,
            f"layers.{layer_i}.attention.wk.weight": wk,
            f"layers.{layer_i}.attention.wv.weight": wv,
            f"layers.{layer_i}.attention.wq.bias": wq_bias,
            f"layers.{layer_i}.attention.wk.bias": wk_bias,
            f"layers.{layer_i}.attention.wv.bias": wv_bias,
            f"layers.{layer_i}.attention.wo.weight": hf_model_state_dict[f"transformer.encoder.layers.{layer_i}.self_attention.dense.weight"],

            f"layers.{layer_i}.feed_forward.w1.weight": w1,
            f"layers.{layer_i}.feed_forward.w3.weight": w3,
            f"layers.{layer_i}.feed_forward.w2.weight": hf_model_state_dict[f"transformer.encoder.layers.{layer_i}.mlp.dense_4h_to_h.weight"],
            
            f"layers.{layer_i}.attention_norm.weight": hf_model_state_dict[f"transformer.encoder.layers.{layer_i}.input_layernorm.weight"],
            f"layers.{layer_i}.ffn_norm.weight": hf_model_state_dict[f"transformer.encoder.layers.{layer_i}.post_attention_layernorm.weight"],
        })

    state_dict.update({
        "tok_embeddings.weight": hf_model_state_dict["transformer.embedding.word_embeddings.weight"],
        "norm.weight": hf_model_state_dict["transformer.encoder.final_layernorm.weight"],
        "output.weight": hf_model_state_dict["transformer.output_layer.weight"]
    })

    torch.save(state_dict, os.path.join(model_path, "model.pth"))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input_dir",
        help="Location of HF weights, which contains tokenizer.model and model folders",
    )
    parser.add_argument(
        "--output_dir",
        help="Location to write OPMX model",
    )
    args = parser.parse_args()
    write_pmx_model(
        model_path=args.output_dir,
        input_base_path=args.input_dir
    )

if __name__ == "__main__":
    main()