llama_ckpt_convert.py

# Copyright (c) OpenMMLab. All rights reserved.

import torch
import fire
import os.path as osp
from os import makedirs
from pathlib import Path
import safetensors
from typing import List
from tqdm import tqdm


def import_fb(ckpt_dir: str):
    checkpoints = []
    for pattern in ['*.pth', '*.pt']:
        checkpoints += sorted(Path(ckpt_dir).glob(pattern))
    print(checkpoints)
    n_ckpt = len(checkpoints)
    model_params = {}

    def get_param(name, size):
        print(name, size)
        if name not in model_params:
            model_params[name] = torch.zeros(
                size, dtype=torch.float16, device='cpu')
        return model_params[name]
    for i, ckpt_path in enumerate(checkpoints):
        ckpt = torch.load(ckpt_path, map_location='cpu')
        for param_name, param_data in ckpt.items():
            key = param_name.split('.')[-2]
            if key in ['w1', 'w3', 'wq', 'wk', 'wv', 'output']:  # column-parallel
                size = param_data.size(0)
                param = get_param(
                    param_name, [size * n_ckpt, param_data.size(1)])
                param.data[size * i: size * (i + 1), :] = param_data
            elif key in ['w2', 'wo', 'tok_embeddings']:          # row-parallel
                size = param_data.size(-1)
                param = get_param(
                    param_name, [param_data.size(0), size * n_ckpt])
                param.data[:, size * i: size * (i + 1)] = param_data
            elif i == 0:
                param = get_param(param_name, param_data.size())
                param.data = param_data
        del ckpt

    for name, param in model_params.items():
        # transpose all weights as FasterTransformer is expecting column-major weights
        # (output_dims, input_dims) -> (input_dims, output_dims)
        key = name.split('.')[-2]
        if key in ['w1', 'w3', 'wq', 'wk', 'wv', 'w2', 'wo']:
            param.data = param.data.t()

    # concat qkv projection
    for i in range(1000):
        _qkv = [f'layers.{i}.attention.{k}.weight' for k in ['wq', 'wk', 'wv']]
        try:
            qkv = tuple(map(model_params.pop, _qkv))
        except KeyError:
            break
        qkv = torch.stack(qkv, dim=1)
        model_params[f'layers.{i}.attention.w_qkv.weight'] = qkv
        print(qkv.shape, qkv.dtype)

    return model_params


def permute(x: torch.Tensor):
    SIZE_PER_HEAD = 128
    if x.shape[-1] > 1:  # qweights
        dim = x.shape[-1]
        n_heads = dim // SIZE_PER_HEAD
        return x.view(-1, n_heads, 2, dim // n_heads // 2).transpose(2, 3).reshape(-1, dim)
    else:  # scales, zeros
        dim = x.shape[0]
        n_heads = dim // SIZE_PER_HEAD
        return x.view(n_heads, 2, dim // n_heads // 2, 1).transpose(1, 2).reshape(dim, 1)


def check_zero(x: torch.Tensor):
    sum = x.flatten().sum().item()
    assert sum == 0, str(sum)


def import_gptq(path: str):
    model_params = {}

    _qweight = 'weight'
    _suffixes = [_qweight]
    n_split = 3
    if True:
        _params = {}
        for i in tqdm(range(0, n_split)):
            filename = "pytorch_model-{:05d}-of-{:05d}.bin".format(i + 1, n_split)
            _tmp = torch.load(osp.join(path, filename), map_location='cpu')
            _params.update(_tmp)
        # print('\n'.join(_params.keys()))
        def get_tensor(name):
            return _params[name]
        def get_tensor_transposed(name):
            return _params[name].t()

    # _qweight = 'qweight'
    # _suffixes = [_qweight, 'bias', 'scales', 'zeros']
    # with safetensors.safe_open(path, framework='pt') as f:
    #     get_tensor = f.get_tensor
    #     # quantized weights are already in column major, no need to transpose
    #     get_tensor_transposed = get_tensor
        for i in range(1000):
            try:
                # attention weights
                _qkvo = [f'model.layers.{i}.self_attn.{t}_proj' for t in 'qkvo']
                for suffix in _suffixes:
                    q, k, v, o = map(get_tensor_transposed, map(('{}.' + suffix).format, _qkvo))
                    if suffix == 'bias':
                        check_zero(q), check_zero(k), check_zero(v), check_zero(o)
                    else:
                        # q, k has different layout for fb & hf, convert to fb's layout
                        q = permute(q)
                        k = permute(k)
                        if suffix == _qweight:  # weight, qweight
                            # insert a dimension for splitting heads later
                            # qkv = torch.cat([q[:, None, :], k[:, None, :], v[:, None, :]], dim=1)
                            qkv = torch.stack((q, k, v), dim=1)
                        else:  # scales, zeros
                            # qkv = torch.cat([q[None, :], k[None, :], v[None, :]], dim=0).squeeze(dim=-1)
                            qkv = torch.stack((q, k, v), dim=0).squeeze(dim=-1)
                        for k, v in [('w_qkv', qkv), ('wo', o)]:
                            model_params[f'layers.{i}.attention.{k}.{suffix}'] = v
                # ffn weights
                _w123 = [f'model.layers.{i}.mlp.{t}_proj' for t in ['gate', 'down', 'up']]
                for suffix in _suffixes:
                    w1, w2, w3 = map(get_tensor_transposed, map(('{}.' + suffix).format, _w123))
                    if suffix == 'bias':
                        check_zero(w1), check_zero(w2), check_zero(w3)
                    else:
                        if suffix in ['scales', 'zeros']:
                            w1, w2, w3 = map(lambda x: x.squeeze(dim=-1), [w1, w2, w3])
                        for k, v in [('w1', w1), ('w2', w2), ('w3', w3)]:
                            model_params[f'layers.{i}.feed_forward.{k}.{suffix}'] = v
                other = [('attention_norm.weight', 'input_layernorm.weight'),
                         ('ffn_norm.weight', 'post_attention_layernorm.weight')]
                for ours, theirs in other:
                    model_params[f'layers.{i}.' + ours] = get_tensor(f'model.layers.{i}.' + theirs)
            except safetensors.SafetensorError:
                break
            except KeyError:
                break
            print(i)

        other = [('tok_embeddings.weight', 'model.embed_tokens.weight'),
                 ('norm.weight', 'model.norm.weight'),
                 ('output.weight', 'lm_head.weight')]
        for ours, theirs in other:
            model_params[ours] = get_tensor(theirs)

        return model_params


def export(model_params: dict, out_dir: str, n_inference: int):
    makedirs(out_dir, exist_ok=True)

    def save_bin(param: torch.Tensor, name):
        print(name, param.shape)
        if param.dtype in [torch.float, torch.bfloat16]:
            param = param.half()
        param.contiguous().numpy().tofile(osp.join(out_dir, name))

    # reverse the spliting axes since the weights are transposed above
    for param_name, param_data in model_params.items():
        split_dim = None
        key, ext = param_name.split('.')[-2:]
        copy = False
        if key in ['w1', 'w3', 'w_qkv']:
            split_dim = -1
        elif key in ['w2', 'wo']:
            if ext in ['scales', 'zeros']:
                copy = True
            else:
                split_dim = 0
        if split_dim is not None:
            print(f'*** spliting {param_name}, shape={param_data.shape}, split_dim={split_dim}')
            assert param_data.shape[split_dim] % n_inference == 0
            split_size = param_data.shape[split_dim] // n_inference
            splits = torch.split(param_data, split_size, dim=split_dim)
            for i, split in enumerate(splits):
                prefix, ext = osp.splitext(param_name)
                save_bin(split, f'{prefix}.{i}{ext}')
        elif copy:
            print(f'### copying {param_name}, shape={param_data.shape}')
            copies = [param_data] * n_inference
            for i, copy in enumerate(copies):
                prefix, ext = osp.splitext(param_name)
                save_bin(copy, f'{prefix}.{i}{ext}')
        else:
            save_bin(param_data, param_name)


def main(kind: str, input_path: str, out_dir: str, n_inference: int = 1):
    if kind == 'fb':
        model_params = import_fb(input_path)
    elif kind == 'gptq':
        model_params = import_gptq(input_path)
    else:
        raise RuntimeError(f'Unsupported kind: {kind}')

    export(model_params, out_dir, n_inference)


if __name__ == '__main__':
    fire.Fire(main)