convert_checkpoint.py

# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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 json
import os
import time
import traceback
from concurrent.futures import ThreadPoolExecutor, as_completed

import numpy as np
import safetensors
import torch
from transformers import AutoModelForCausalLM

import tensorrt_llm
from tensorrt_llm._utils import pad_vocab_size, str_dtype_to_torch


def torch_split(v, tensor_parallel, idx, dim=0):
    if tensor_parallel == 1:
        return v
    else:
        return (torch.split(v, v.shape[dim] // tensor_parallel,
                            dim=dim)[idx]).contiguous()


def convert_hf_phi(hf_model,
                   rank=0,
                   tensor_parallel=1,
                   dtype='float32',
                   use_parallel_embedding=False,
                   sharding_dim=0):

    hf_model_phi_block_names = [
        "input_layernorm.weight",
        "input_layernorm.bias",
        "self_attn.dense.weight",
        "self_attn.dense.bias",
        "mlp.fc1.weight",
        "mlp.fc1.bias",
        "mlp.fc2.weight",
        "mlp.fc2.bias",
    ]

    tensorrt_llm_model_phi_block_names = [
        "input_layernorm.weight",
        "input_layernorm.bias",
        "attention.dense.weight",
        "attention.dense.bias",
        "mlp.fc.weight",
        "mlp.fc.bias",
        "mlp.proj.weight",
        "mlp.proj.bias",
    ]

    weights = {}
    torch_dtype = str_dtype_to_torch(dtype)
    hf_phi_state_dict = hf_model.state_dict()

    # Embedding
    # [vocab_size, hidden_size]
    v = hf_phi_state_dict.get('model.embed_tokens.weight').to(torch_dtype).cpu()
    if use_parallel_embedding:
        v = torch_split(v, tensor_parallel, rank, sharding_dim)
    weights['transformer.vocab_embedding.weight'] = v

    # Decoder Layers
    n_layer = hf_model.config.num_hidden_layers
    for layer_idx in range(n_layer):
        hf_prefix = f"model.layers.{layer_idx}."
        tllm_prex = f'transformer.layers.{layer_idx}.'

        # MLPs
        for idx, hf_attr in enumerate(hf_model_phi_block_names):
            v = hf_phi_state_dict.get(hf_prefix + hf_attr).to(torch_dtype).cpu()

            if tensor_parallel > 1:
                if 'self_attn.dense.weight' in hf_attr:
                    # [n=hidden_size, k=hidden_size] ->
                    # [n=hidden_size, k=hidden_size // tensor_parallel]
                    v = torch_split(v, tensor_parallel, rank, dim=1)
                elif 'mlp.fc1.weight' in hf_attr:
                    # [hidden_size * 4, hidden_size] ->
                    # [hidden_size * 4 // tensor_parallel, hidden_size]
                    v = torch_split(v, tensor_parallel, rank, dim=0)
                elif 'mlp.fc1.bias' in hf_attr:
                    # [hidden_size * 4] -> [hidden_size * 4 // tensor_parallel]
                    v = torch_split(v, tensor_parallel, rank, dim=0)
                elif 'mlp.fc2.weight' in hf_attr:
                    # [hidden_size, hidden_size * 4] ->
                    # [hidden_size, hidden_size * 4 // tensor_parallel]
                    v = torch_split(v, tensor_parallel, rank, dim=1)

            tllm_attr = tensorrt_llm_model_phi_block_names[idx]
            weights[f'{tllm_prex}{tllm_attr}'] = v

        # Attention QKV Linear
        num_heads = hf_model.config.num_attention_heads
        hidden_size = hf_model.config.hidden_size
        hidden_size // num_heads

        # [(num_heads x q)|(num_heads x k)|(num_heads x v), hidden_size]
        q_weights = hf_phi_state_dict.get(hf_prefix + "self_attn.q_proj.weight")
        k_weights = hf_phi_state_dict.get(hf_prefix + "self_attn.k_proj.weight")
        v_weights = hf_phi_state_dict.get(hf_prefix + "self_attn.v_proj.weight")
        q_bias = hf_phi_state_dict.get(hf_prefix + "self_attn.q_proj.bias")
        k_bias = hf_phi_state_dict.get(hf_prefix + "self_attn.k_proj.bias")
        v_bias = hf_phi_state_dict.get(hf_prefix + "self_attn.v_proj.bias")
        qkv_weights = torch.cat((q_weights, k_weights, v_weights), dim=0)
        qkv_bias = torch.cat((q_bias, k_bias, v_bias), dim=0)

        qkv_weights = qkv_weights.reshape([hidden_size * 3, hidden_size])
        qkv_bias = qkv_bias.reshape([hidden_size * 3])

        if tensor_parallel > 1:
            qkv_weights = qkv_weights.reshape(
                3, hidden_size, hidden_size).to(torch_dtype).cpu()
            qkv_weights = torch_split(qkv_weights, tensor_parallel, rank,
                                      dim=1).reshape(
                                          3 * (hidden_size // tensor_parallel),
                                          hidden_size)

            qkv_bias = qkv_bias.reshape(3, hidden_size).to(torch_dtype).cpu()
            qkv_bias = torch_split(qkv_bias, tensor_parallel, rank,
                                   dim=1).reshape(
                                       3 * (hidden_size // tensor_parallel))

            weights[
                f"{tllm_prex}attention.qkv.weight"] = qkv_weights.contiguous()
            weights[f"{tllm_prex}attention.qkv.bias"] = qkv_bias.contiguous()
        else:
            weights[f"{tllm_prex}attention.qkv.weight"] = qkv_weights.to(
                torch_dtype).cpu()
            weights[f"{tllm_prex}attention.qkv.bias"] = qkv_bias.to(
                torch_dtype).cpu()

    # Final Layer Norm
    v = hf_phi_state_dict.get('model.final_layernorm.weight')
    weights["transformer.ln_f.weight"] = v.to(torch_dtype).cpu()

    v = hf_phi_state_dict.get('model.final_layernorm.bias')
    weights["transformer.ln_f.bias"] = v.to(torch_dtype).cpu()

    # LM Head
    v = hf_phi_state_dict.get('lm_head.weight').to(torch_dtype).cpu()
    if tensor_parallel > 1:
        # [vocab_size, hidden_size] ->
        # [vocab_size // tensor_parallel, hidden_size]
        if v.shape[0] % tensor_parallel != 0:
            # padding
            vocab_size_padded = pad_vocab_size(v.shape[0], tensor_parallel)
            pad_width = vocab_size_padded - v.shape[0]
            v = np.pad(v, ((0, pad_width), (0, 0)),
                       'constant',
                       constant_values=0)

        v = torch_split(v, tensor_parallel, rank, dim=0)
    weights["lm_head.weight"] = v

    v = hf_phi_state_dict.get('lm_head.bias').to(torch_dtype).cpu()
    if tensor_parallel > 1:
        v = torch_split(v, tensor_parallel, rank, dim=0)
    weights["lm_head.bias"] = v

    return weights


def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_dir', type=str, default=None)
    parser.add_argument('--tp_size',
                        type=int,
                        default=1,
                        help='N-way tensor parallelism size')
    parser.add_argument('--pp_size',
                        type=int,
                        default=1,
                        help='N-way pipeline parallelism size')
    parser.add_argument('--dtype',
                        type=str,
                        default='float16',
                        choices=['float32', 'bfloat16', 'float16'])
    parser.add_argument(
        '--use_parallel_embedding',
        action="store_true",
        default=False,
        help=
        'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled'
    )
    parser.add_argument(
        '--embedding_sharding_dim',
        type=int,
        default=0,
        choices=[0, 1],
        help=
        'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). '
        'To shard it along hidden dimension, set embedding_sharding_dim=1'
        'Note: embedding sharing is only enabled when embedding_sharding_dim = 0'
    )
    parser.add_argument('--output_dir',
                        type=str,
                        default='tllm_checkpoint',
                        help='The path to save the TensorRT-LLM checkpoint')
    parser.add_argument(
        '--workers',
        type=int,
        default=1,
        help='The number of workers for converting checkpoint in parallel')
    args = parser.parse_args()

    return args


if __name__ == '__main__':
    # TODO(qijun): Currently, the convert script depends on a torch op:
    # torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix,
    # which is included in tensorrt_llm Python package. Otherwise, the convert
    # script does not need to import tensorrt_llm. Will remove it after reimplementing
    # the op with PyTorch.
    print(tensorrt_llm.__version__)
    args = parse_arguments()
    world_size = args.tp_size * args.pp_size
    assert args.pp_size == 1, "Pipeline parallelism is not supported."

    tik = time.time()

    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)

    hf_model = AutoModelForCausalLM.from_pretrained(args.model_dir,
                                                    torch_dtype="auto",
                                                    trust_remote_code=True)
    hf_config = hf_model.config
    config = {
        'architecture': hf_config.architectures[0],
        'dtype': args.dtype,
        'num_hidden_layers': hf_config.num_hidden_layers,
        'num_attention_heads': hf_config.num_key_value_heads,
        'partial_rotary_factor': hf_config.partial_rotary_factor,
        'rope_theta': hf_config.rope_theta,
        'hidden_size': hf_config.hidden_size,
        'intermediate_size': hf_config.intermediate_size,
        'vocab_size': hf_config.vocab_size,
        'max_position_embeddings': hf_config.max_position_embeddings,
        'hidden_act': hf_config.hidden_act,
        'mapping': {
            'world_size': world_size,
            'tp_size': args.tp_size,
            'pp_size': args.pp_size,
        },
        'use_parallel_embedding': False,
        'embedding_sharding_dim': args.embedding_sharding_dim,
        'share_embedding_table': False,
    }

    with open(os.path.join(args.output_dir, 'config.json'), 'w') as f:
        json.dump(config, f, indent=4)

    def covert_and_save(rank):
        weights = convert_hf_phi(
            hf_model,
            rank,
            world_size,
            dtype=args.dtype,
            use_parallel_embedding=args.use_parallel_embedding,
            sharding_dim=args.embedding_sharding_dim)
        safetensors.torch.save_file(
            weights, os.path.join(args.output_dir, f'rank{rank}.safetensors'))

    if args.workers == 1:
        for rank in range(world_size):
            covert_and_save(rank)
    else:
        with ThreadPoolExecutor(max_workers=args.workers) as p:
            futures = [
                p.submit(covert_and_save, rank) for rank in range(world_size)
            ]
            exceptions = []
            for future in as_completed(futures):
                try:
                    future.result()
                except Exception as e:
                    traceback.print_exc()
                    exceptions.append(e)
            assert len(
                exceptions
            ) == 0, "Checkpoint conversion failed, please check error log."

    tok = time.time()
    t = time.strftime('%H:%M:%S', time.gmtime(tok - tik))
    print(f'Total time of converting checkpoints: {t}')