TensorRT-LLMs/examples/bloom/weight.py

# SPDX-FileCopyrightText: Copyright (c) 2022-2023 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 time

import numpy as np
import torch

import tensorrt_llm


def split(v, tp_size, idx, dim=0):
    if tp_size == 1:
        return v
    if len(v.shape) == 1:
        return np.ascontiguousarray(np.split(v, tp_size)[idx])
    else:
        return np.ascontiguousarray(np.split(v, tp_size, axis=dim)[idx])


def reorder_qkv_weight_or_bias(v, n_head, n_hidden, is_bias=False):
    """ Reorder the qkv weight.

    Note that the shape of the fused QKV weights in HF is different from the
    shape that TRT-LLM requires.
       HF: (num_heads x 3 x head_dim, hidden_size)
       TRT-LLM: (3 x num_heads x head_dim, hidden_size)
    This is unlike to the other models in HF e.g. GPT where they have the
    same shape with TRT-LLM, i.e., (3 x num_heads x head_dim, hidden_size). Also,
    to split across attention heads in tensor parallel, we reshape the qkv
        weight: (3, num_heads x head_dim, hidden).
        bias  : (3, num_heads x head_dim).
    """

    head_dim = n_hidden // n_head

    # (3 x hidden, ...) view as (num_heads, 3, head_dim, ...)
    v = v.reshape(n_head, 3, head_dim, -1)
    # permute to (3, num_heads, head_dim, ...)
    v = v.transpose((1, 0, 2, 3))
    # final shape: weight=(3, hidden, hidden) or bias=(3, hidden)
    if is_bias:
        return v.reshape(3, n_hidden)
    return v.reshape(3, n_hidden, n_hidden)


def split_qkv_tp(tensorrt_llm_bloom, v, tensor_parallel, rank):
    """
    Splits the QKV matrix according to tensor parallelism
    """
    n_heads = tensorrt_llm_bloom._num_heads
    hidden_size = tensorrt_llm_bloom._hidden_size
    v = reorder_qkv_weight_or_bias(v, n_heads, hidden_size, is_bias=False)
    split_v = split(v, tensor_parallel, rank, dim=1)
    split_v = split_v.reshape(3 * (hidden_size // tensor_parallel), hidden_size)

    return np.ascontiguousarray(split_v)


def split_qkv_bias_tp(tensorrt_llm_bloom, v, tensor_parallel, rank):
    """
    Splits the QKV bias according to tensor parallelism
    """
    layer = tensorrt_llm_bloom.layers[0]
    n_heads = layer.num_attention_heads
    hidden_size = layer.hidden_size
    v = reorder_qkv_weight_or_bias(v, n_heads, hidden_size, is_bias=True)
    split_v = split(v, tensor_parallel, rank, dim=1)
    split_v = split_v.reshape(3 * (hidden_size // tensor_parallel))
    return np.ascontiguousarray(split_v)


def split_matrix_tp(v, tensor_parallel, rank, dim):
    return np.ascontiguousarray(split(v, tensor_parallel, rank, dim=dim))


def get_weight(config, prefix, dtype):
    return config[prefix + '.weight'].to(dtype).detach().cpu().numpy()


def get_bias(config, prefix, dtype):
    return config[prefix + '.bias'].to(dtype).detach().cpu().numpy()


def get_weight_and_bias(config, prefix, dtype):
    return get_weight(config, prefix, dtype), get_bias(config, prefix, dtype)


def load_from_hf_bloom(tensorrt_llm_bloom,
                       hf_bloom,
                       rank=0,
                       tensor_parallel=1,
                       fp16=False,
                       use_parallel_embedding=False,
                       sharding_dim=0):
    tensorrt_llm.logger.info('Loading weights from HF BLOOM...')
    tik = time.time()

    model_params = dict(hf_bloom.named_parameters())
    dtype = torch.float16 if fp16 else torch.float32
    for l in range(hf_bloom.config.num_hidden_layers):
        prefix = f'transformer.h.{l}.'

        qkv_weight, qkv_bias = get_weight_and_bias(
            model_params, prefix + 'self_attention.query_key_value', dtype)
        tensorrt_llm_bloom.layers[l].attention.qkv.weight.value = split_qkv_tp(
            tensorrt_llm_bloom, qkv_weight, tensor_parallel, rank)
        tensorrt_llm_bloom.layers[
            l].attention.qkv.bias.value = split_qkv_bias_tp(
                tensorrt_llm_bloom, qkv_bias, tensor_parallel, rank)

        attn_dense_weight, attn_dense_bias = get_weight_and_bias(
            model_params, prefix + 'self_attention.dense', dtype)
        tensorrt_llm_bloom.layers[
            l].attention.dense.weight.value = split_matrix_tp(attn_dense_weight,
                                                              tensor_parallel,
                                                              rank,
                                                              dim=1)
        tensorrt_llm_bloom.layers[
            l].attention.dense.bias.value = attn_dense_bias

        mlp_fc_weight, mlp_fc_bias = get_weight_and_bias(
            model_params, prefix + 'mlp.dense_h_to_4h', dtype)
        tensorrt_llm_bloom.layers[l].mlp.fc.weight.value = split_matrix_tp(
            mlp_fc_weight, tensor_parallel, rank, dim=0)
        tensorrt_llm_bloom.layers[l].mlp.fc.bias.value = split_matrix_tp(
            mlp_fc_bias, tensor_parallel, rank, dim=0)

        mlp_proj_weight, mlp_proj_bias = get_weight_and_bias(
            model_params, prefix + 'mlp.dense_4h_to_h', dtype)
        tensorrt_llm_bloom.layers[l].mlp.proj.weight.value = split_matrix_tp(
            mlp_proj_weight, tensor_parallel, rank, dim=1)
        tensorrt_llm_bloom.layers[l].mlp.proj.bias.value = mlp_proj_bias

        # Layer norms do not use tensor parallelism
        input_ln_weight, input_ln_bias = get_weight_and_bias(
            model_params, prefix + 'input_layernorm', dtype)
        tensorrt_llm_bloom.layers[
            l].input_layernorm.weight.value = input_ln_weight
        tensorrt_llm_bloom.layers[l].input_layernorm.bias.value = input_ln_bias

        post_ln_weight, post_ln_bias = get_weight_and_bias(
            model_params, prefix + 'post_attention_layernorm', dtype)
        tensorrt_llm_bloom.layers[
            l].post_layernorm.weight.value = post_ln_weight
        tensorrt_llm_bloom.layers[l].post_layernorm.bias.value = post_ln_bias

    embed_w = get_weight(model_params, 'transformer.word_embeddings', dtype)
    if not use_parallel_embedding:
        tensorrt_llm_bloom.embedding.weight.value = embed_w.copy()
    else:
        assert hf_bloom.config.vocab_size % tensor_parallel == 0
        tensorrt_llm_bloom.embedding.weight.value = np.ascontiguousarray(
            split(embed_w.copy(), tensor_parallel, rank, dim=sharding_dim))
    tensorrt_llm_bloom.lm_head.weight.value = split_matrix_tp(embed_w,
                                                              tensor_parallel,
                                                              rank,
                                                              dim=0)

    embed_f_w, embed_f_b = get_weight_and_bias(
        model_params, 'transformer.word_embeddings_layernorm', dtype)
    tensorrt_llm_bloom.ln_embed.weight.value = embed_f_w
    tensorrt_llm_bloom.ln_embed.bias.value = embed_f_b

    ln_f_w, ln_f_b = get_weight_and_bias(model_params, 'transformer.ln_f',
                                         dtype)
    tensorrt_llm_bloom.ln_f.weight.value = ln_f_w
    tensorrt_llm_bloom.ln_f.bias.value = ln_f_b

    tok = time.time()
    t = time.strftime('%H:%M:%S', time.gmtime(tok - tik))
    tensorrt_llm.logger.info(f'Weights loaded. Total time: {t}')