import argparse
import functools
import json
import os
import time
from collections import defaultdict
from pathlib import Path
from typing import Any, Dict, Iterable, Optional, Union

import numpy as np
import safetensors
import torch
import torch.nn as nn
from tqdm import tqdm
from transformers import BloomConfig, BloomForCausalLM, BloomTokenizerFast
from transformers.models.bloom.modeling_bloom import BloomBlock
from transformers.pytorch_utils import Conv1D

# isort: off
import tensorrt_llm
from tensorrt_llm import logger
from tensorrt_llm.quantization import QuantAlgo, QuantMode
from tensorrt_llm.models.convert_utils import iterate_shard_files, load_state_dict, \
     load_calib_dataset, split_matrix_tp, get_weight_and_bias, split, smooth_gemm, \
     generate_int8,get_weight
# isort: on


@torch.no_grad()
def capture_activation_range(model,
                             tokenizer,
                             dataset,
                             num_samples=512,
                             seq_len=512):
    model.eval()
    device = next(model.parameters()).device
    act_scales = defaultdict(lambda: {"x": None, "y": None, "w": None})

    def stat_tensor(name, tensor, act_scales, key):
        hidden_dim = tensor.shape[-1]
        tensor = tensor.view(-1, hidden_dim).abs().detach()
        comming_max = torch.max(tensor, dim=0)[0].float()

        if act_scales[name][key] is None:
            act_scales[name][key] = comming_max
        else:
            act_scales[name][key] = torch.max(act_scales[name][key],
                                              comming_max)

    def stat_input_hook(m, x, y, name):
        if isinstance(x, tuple):
            x = x[0]
        stat_tensor(name, x, act_scales, "x")
        stat_tensor(name, y, act_scales, "y")

        if act_scales[name]["w"] is None:
            act_scales[name]["w"] = m.weight.abs().clip(1e-8,
                                                        None).max(dim=1)[0]

    hooks = []
    for name, m in model.named_modules():
        if isinstance(m, nn.Linear) or isinstance(m, Conv1D):
            hooks.append(
                m.register_forward_hook(
                    functools.partial(stat_input_hook, name=name)))

    for i in tqdm(range(num_samples), desc="calibrating model"):
        input_ids = tokenizer(dataset[i],
                              return_tensors="pt",
                              max_length=seq_len,
                              truncation=True).input_ids.to(device)
        model(input_ids)

    for h in hooks:
        h.remove()

    return act_scales


def reorder_torch_qkv_weight_or_bias(v, model, 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). We reshape the qkv
        weight: (3 x num_heads x head_dim, hidden).
        bias  : (3 x num_heads x head_dim).
    """

    n_head = model.transformer.num_heads
    hidden_size = model.transformer.embed_dim
    head_dim = hidden_size // 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.permute((1, 0, 2, 3))
    # final shape: weight=(3 x hidden, hidden) or bias=(3 x hidden)
    if is_bias:
        return v.reshape(3 * hidden_size)
    return v.reshape(3 * hidden_size, hidden_size)


def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_dir', type=Path, 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_weight_only',
        default=False,
        action="store_true",
        help='Quantize weights for the various GEMMs to INT4/INT8.'
        'See --weight_only_precision to set the precision')
    parser.add_argument(
        '--weight_only_precision',
        const='int8',
        type=str,
        nargs='?',
        default='int8',
        choices=['int8', 'int4'],
        help=
        'Define the precision for the weights when using weight-only quantization.'
        'You must also use --use_weight_only for that argument to have an impact.'
    )
    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=Path,
                        default='tllm_checkpoint',
                        help='The path to save the TensorRT-LLM checkpoint')
    parser.add_argument(
        '--calib_dataset',
        type=str,
        default='lambada',
        help=
        "The huggingface dataset name or the local directory of the dataset for calibration."
    )
    parser.add_argument(
        "--smoothquant",
        "-sq",
        type=float,
        default=None,
        help="Set the α parameter (see https://arxiv.org/pdf/2211.10438.pdf)"
        " to Smoothquant the model, and output int8 weights."
        " A good first try is 0.5. Must be in [0, 1]")
    parser.add_argument(
        '--per_channel',
        action="store_true",
        default=False,
        help=
        'By default, we use a single static scaling factor for the GEMM\'s result. '
        'per_channel instead uses a different static scaling factor for each channel. '
        'The latter is usually more accurate, but a little slower.')
    parser.add_argument(
        '--per_token',
        action="store_true",
        default=False,
        help=
        'By default, we use a single static scaling factor to scale activations in the int8 range. '
        'per_token chooses at run time, and for each token, a custom scaling factor. '
        'The latter is usually more accurate, but a little slower.')
    parser.add_argument(
        '--int8_kv_cache',
        default=False,
        action="store_true",
        help=
        'By default, we use dtype for KV cache. int8_kv_cache chooses int8 quantization for KV'
    )
    parser.add_argument(
        '--workers',
        type=int,
        default=1,
        help='The number of workers to convert checkpoint in parallel')
    parser.add_argument('--log_level', type=str, default='info')
    args = parser.parse_args()

    return args


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(0, 1)
    # 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(v, n_head, n_hidden, tensor_parallel, rank):
    """
    Splits the QKV matrix according to tensor parallelism
    """
    v = reorder_qkv_weight_or_bias(v, n_head, n_hidden, is_bias=False)
    split_v = split(v, tensor_parallel, rank, dim=1)
    split_v = split_v.reshape(3 * (n_hidden // tensor_parallel), n_hidden)
    return split_v.contiguous()


def split_qkv_bias_tp(v, n_head, n_hidden, tensor_parallel, rank):
    """
    Splits the QKV bias according to tensor parallelism
    """
    v = reorder_qkv_weight_or_bias(v, n_head, n_hidden, is_bias=True)
    split_v = split(v, tensor_parallel, rank, dim=1)
    split_v = split_v.reshape(3 * (n_hidden // tensor_parallel))
    return split_v.contiguous()


def get_tllm_linear_weight(weight,
                           prefix,
                           bias=None,
                           use_weight_only=False,
                           plugin_weight_only_quant_type=torch.int8):
    results = {}
    if use_weight_only:
        v = weight.cpu().t().contiguous()
        processed_torch_weights, torch_weight_scales = \
            torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix(
                v, plugin_weight_only_quant_type)
        results[prefix + 'weight'] = processed_torch_weights
        results[prefix + 'per_channel_scale'] = torch_weight_scales
    else:
        results[prefix + 'weight'] = weight.contiguous()

    if bias is not None:
        results[prefix + 'bias'] = bias

    return results


def add_tllm_weight(
        weights: Dict[str, torch.Tensor],
        name: str,
        param: torch.Tensor,
        quant_mode: QuantMode = QuantMode(0),
):
    assert name not in weights, f'{name} is already added.'

    if name.endswith('.weight') and quant_mode.is_weight_only():
        if quant_mode.is_int8_weight_only():
            quant_dtype = torch.int8
        elif quant_mode.is_int4_weight_only():
            quant_dtype = torch.quint4x2
        else:
            raise ValueError(
                f'Invalid configuration, got quant_mode={quant_mode}')
        processed_torch_weights, torch_weight_scales = \
            torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix(
                param.t().contiguous(), quant_dtype)
        weights[name] = processed_torch_weights
        scale_name = name.replace('.weight', '.per_channel_scale')
        weights[scale_name] = torch_weight_scales
    else:
        weights[name] = param.contiguous()


@torch.no_grad()
def smooth_bloom_model(model, scales, alpha, bloom_qkv_param, bloom_smoother):
    # Smooth the activation and weights with smoother = $\diag{s}$
    for name, module in model.named_modules():
        if not isinstance(module, BloomBlock):
            continue

        # reorder qkv weight/bias and scales
        param = module.self_attention.query_key_value.weight
        param = reorder_torch_qkv_weight_or_bias(param, model, is_bias=False)

        layer_name = name + ".self_attention.query_key_value"
        act_range_qkv = scales.get(layer_name)
        # (n_head x 3 x head_dim) -> (3 x n_head x head_dim)
        act_range_qkv['w'] = reorder_torch_qkv_weight_or_bias(
            act_range_qkv['w'], model, is_bias=True)
        act_range_qkv['y'] = reorder_torch_qkv_weight_or_bias(
            act_range_qkv['y'], model, is_bias=True)
        scales[layer_name] = act_range_qkv

        # qkv_proj
        smoother = smooth_gemm(param, scales[layer_name]["x"],
                               module.input_layernorm.weight,
                               module.input_layernorm.bias, alpha)
        scales[layer_name]["x"] = scales[layer_name]["x"] / smoother
        scales[layer_name]["w"] = param.abs().max(dim=1)[0]
        bloom_qkv_param[layer_name] = param

        # dense
        # enabled for better accuracy with perf overhead of quantization
        layer_name = name + ".self_attention.dense"
        smoother = smooth_gemm(module.self_attention.dense.weight,
                               scales[layer_name]["x"], None, None, alpha)
        bloom_smoother[layer_name] = smoother

        scales[layer_name]["x"] = scales[layer_name]["x"] / smoother
        scales[layer_name]["w"] = module.self_attention.dense.weight.abs().max(
            dim=1)[0]

        # fc1
        layer_name = name + ".mlp.dense_h_to_4h"
        smoother = smooth_gemm(module.mlp.dense_h_to_4h.weight,
                               scales[layer_name]["x"],
                               module.post_attention_layernorm.weight,
                               module.post_attention_layernorm.bias, alpha)
        scales[layer_name]["x"] = scales[layer_name]["x"] / smoother
        scales[layer_name]["w"] = module.mlp.dense_h_to_4h.weight.abs().max(
            dim=1)[0]

        # fc2
        # enabled for better accuracy with perf overhead of quantization
        layer_name = name + ".mlp.dense_4h_to_h"
        smoother = smooth_gemm(module.mlp.dense_4h_to_h.weight,
                               scales[layer_name]["x"], None, None, alpha)
        bloom_smoother[layer_name] = smoother
        scales[layer_name]["x"] = scales[layer_name]["x"] / smoother
        scales[layer_name]["w"] = module.mlp.dense_4h_to_h.weight.abs().max(
            dim=1)[0]


def get_tllm_linear_sq_weight(
    vals,
    prefix,
    shape,
    tensor_parallel,
    is_qkv=False,
    per_token=False,
    per_channel=False,
    last_prefix=None,
    bias=None,
    smoother_value=None,
    smoother_shape=None,
    rank=0,
    cat_dim=0,
):
    results = {}

    col_shape = shape if (is_qkv or per_channel) else [1, 1]

    if per_token:

        original_weights = vals["weight.int8.col"]
        cur_weights = np.split(original_weights, tensor_parallel,
                               axis=cat_dim)[rank]
        if is_qkv:
            hidden_dim = cur_weights.shape[0]
            cur_weights = cur_weights.reshape(hidden_dim, -1)
        results[prefix + 'weight'] = cur_weights.t().contiguous()
        if smoother_value is None:
            results[last_prefix] = torch.from_numpy(
                np.array([1.0], dtype=np.float32))

        if smoother_value is None:
            cur_per_channel_value = np.split(vals["scale_w_quant_orig.col"],
                                             tensor_parallel,
                                             axis=cat_dim)[rank]
        else:
            cur_per_channel_value = vals["scale_w_quant_orig.col"]
        results[prefix + 'per_channel_scale'] = cur_per_channel_value.reshape(
            col_shape).contiguous()
    else:
        original_weights = vals["weight.int8"]
        cur_weights = np.split(original_weights, tensor_parallel,
                               axis=cat_dim)[rank]

        if is_qkv:
            hidden_dim = cur_weights.shape[0]
            cur_weights = cur_weights.reshape(hidden_dim, -1)
        results[prefix + 'weight'] = cur_weights.t().contiguous()

        cur_per_channel_value = vals["scale_y_accum_quant"]

        results[prefix + 'per_channel_scale'] = cur_per_channel_value.reshape(
            col_shape).contiguous()

        results[last_prefix] = vals['scale_x_orig_quant'].contiguous()

        results[prefix + 'act_scale'] = vals["scale_y_quant_orig"].contiguous()

    if smoother_value is not None:
        cur_smoother_value = np.split(smoother_value,
                                      tensor_parallel,
                                      axis=cat_dim)[rank]
        results[prefix + 'smoother'] = cur_smoother_value.reshape(
            smoother_shape).contiguous().to(torch.float32)

    if bias is not None:
        results[prefix + 'bias'] = bias

    return results


def convert_hf_bloom(hf_bloom,
                     rank=0,
                     tensor_parallel=1,
                     dtype='float32',
                     use_parallel_embedding=False,
                     sharding_dim=0,
                     use_weight_only=False,
                     plugin_weight_only_quant_type=torch.int8,
                     use_smooth_quant=False,
                     bloom_qkv_param={},
                     act_range=None,
                     smoother=None,
                     per_channel=False,
                     per_token=False,
                     int8_kv_cache=False):

    weights = {}
    tik = time.time()

    model_params = dict(hf_bloom.named_parameters())
    dtype = getattr(torch, dtype)
    num_attention_heads = hf_bloom.config.n_head
    hidden_size = hf_bloom.config.hidden_size

    for l in range(hf_bloom.config.num_hidden_layers):
        prefix = f'transformer.h.{l}.'
        tllm_prex = f'transformer.layers.{l}.'
        qkv_weight, qkv_bias = get_weight_and_bias(
            model_params, prefix + 'self_attention.query_key_value', dtype)
        split_v = split_qkv_tp(qkv_weight, num_attention_heads, hidden_size,
                               tensor_parallel, rank)
        bias = split_qkv_bias_tp(qkv_bias, num_attention_heads, hidden_size,
                                 tensor_parallel, rank)

        if use_smooth_quant:
            qkv_weight = bloom_qkv_param[prefix +
                                         'self_attention.query_key_value'].t()

            qkv_weight = qkv_weight.reshape(hidden_size, 3, hidden_size)
            int8_weights = generate_int8(
                qkv_weight,
                act_range.get(
                    (tllm_prex + 'self_attention.query_key_value').replace(
                        ".layers.", ".h.")),
                is_qkv=True)

            weights.update(
                get_tllm_linear_sq_weight(
                    int8_weights,
                    tllm_prex + 'attention.qkv.',
                    [1, 3 * hidden_size // tensor_parallel],
                    tensor_parallel,
                    is_qkv=True,
                    per_token=per_token,
                    per_channel=per_channel,
                    last_prefix=tllm_prex + 'input_layernorm.scale_to_int',
                    bias=bias,
                    smoother_value=None,
                    smoother_shape=None,
                    rank=rank,
                    cat_dim=-1))
        else:
            split_v = split_qkv_tp(qkv_weight, num_attention_heads, hidden_size,
                                   tensor_parallel, rank)
            bias = split_qkv_bias_tp(qkv_bias, num_attention_heads, hidden_size,
                                     tensor_parallel, rank)
            weights.update(
                get_tllm_linear_weight(split_v, tllm_prex + 'attention.qkv.',
                                       bias, use_weight_only,
                                       plugin_weight_only_quant_type))
        if int8_kv_cache:
            qkv_weight = qkv_weight.reshape(hidden_size, 3, hidden_size)

            int8_weights = generate_int8(
                qkv_weight,
                act_range.get(
                    (tllm_prex + 'self_attention.query_key_value').replace(
                        ".layers.", ".h.")),
                is_qkv=True)

            kv_cache_weights = {}

            kv_cache_weights[
                tllm_prex +
                'attention.kv_cache_scaling_factor'] = torch.from_numpy(
                    np.array([int8_weights['scale_y_quant_orig']],
                             dtype=np.float32)).contiguous()
            weights.update(kv_cache_weights)
        attn_dense_weight, attn_dense_bias = get_weight_and_bias(
            model_params, prefix + 'self_attention.dense', dtype)
        if use_smooth_quant:
            attn_dense_weight = attn_dense_weight.t()
            int8_weights = generate_int8(
                attn_dense_weight,
                act_range.get((tllm_prex + 'self_attention.dense').replace(
                    ".layers.", ".h.")))

            weights.update(
                get_tllm_linear_sq_weight(
                    int8_weights,
                    tllm_prex + 'attention.dense.', [1, hidden_size],
                    tensor_parallel,
                    is_qkv=False,
                    per_token=per_token,
                    per_channel=per_channel,
                    last_prefix=tllm_prex +
                    'attention.quantization_scaling_factor',
                    bias=attn_dense_bias,
                    smoother_value=smoother[(tllm_prex +
                                             'self_attention.dense').replace(
                                                 ".layers.", ".h.")],
                    smoother_shape=[1, hidden_size // tensor_parallel],
                    rank=rank,
                    cat_dim=0))
        else:

            split_v = split_matrix_tp(attn_dense_weight,
                                      tensor_parallel,
                                      rank,
                                      dim=1)
            weights.update(
                get_tllm_linear_weight(split_v, tllm_prex + 'attention.dense.',
                                       attn_dense_bias, use_weight_only,
                                       plugin_weight_only_quant_type))
        mlp_fc_weight, mlp_fc_bias = get_weight_and_bias(
            model_params, prefix + 'mlp.dense_h_to_4h', dtype)
        bias = split_matrix_tp(mlp_fc_bias, tensor_parallel, rank, dim=0)

        if use_smooth_quant:
            mlp_fc_weight = mlp_fc_weight.t()
            int8_weights = generate_int8(
                mlp_fc_weight,
                act_range.get((tllm_prex + 'mlp.dense_h_to_4h').replace(
                    ".layers.", ".h.")))

            weights.update(
                get_tllm_linear_sq_weight(
                    int8_weights,
                    tllm_prex + 'mlp.fc.',
                    [1, 4 * hidden_size // tensor_parallel],
                    tensor_parallel,
                    is_qkv=False,
                    per_token=per_token,
                    per_channel=per_channel,
                    last_prefix=tllm_prex + 'post_layernorm.scale_to_int',
                    bias=bias,
                    smoother_value=None,
                    smoother_shape=None,
                    rank=rank,
                    cat_dim=-1))
        else:
            split_v = split_matrix_tp(mlp_fc_weight,
                                      tensor_parallel,
                                      rank,
                                      dim=0)
            bias = split_matrix_tp(mlp_fc_bias, tensor_parallel, rank, dim=0)
            weights.update(
                get_tllm_linear_weight(split_v, tllm_prex + 'mlp.fc.', bias,
                                       use_weight_only,
                                       plugin_weight_only_quant_type))
        mlp_proj_weight, mlp_proj_bias = get_weight_and_bias(
            model_params, prefix + 'mlp.dense_4h_to_h', dtype)
        if use_smooth_quant:
            mlp_proj_weight = mlp_proj_weight.t()
            int8_weights = generate_int8(
                mlp_proj_weight,
                act_range.get((tllm_prex + 'mlp.dense_4h_to_h').replace(
                    ".layers.", ".h.")))

            weights.update(
                get_tllm_linear_sq_weight(
                    int8_weights,
                    tllm_prex + 'mlp.proj.', [1, hidden_size],
                    tensor_parallel,
                    is_qkv=False,
                    per_token=per_token,
                    per_channel=per_channel,
                    last_prefix=tllm_prex + 'mlp.quantization_scaling_factor',
                    bias=mlp_proj_bias,
                    smoother_value=smoother[(tllm_prex +
                                             'mlp.dense_4h_to_h').replace(
                                                 ".layers.", ".h.")],
                    smoother_shape=[1, 4 * hidden_size // tensor_parallel],
                    rank=rank,
                    cat_dim=0))
        else:
            split_v = split_matrix_tp(mlp_proj_weight,
                                      tensor_parallel,
                                      rank,
                                      dim=1)
            weights.update(
                get_tllm_linear_weight(split_v, tllm_prex + 'mlp.proj.',
                                       mlp_proj_bias, use_weight_only,
                                       plugin_weight_only_quant_type))

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

        post_ln_weight, post_ln_bias = get_weight_and_bias(
            model_params, prefix + 'post_attention_layernorm', dtype)
        weights[tllm_prex + 'post_layernorm.weight'] = post_ln_weight
        weights[tllm_prex + 'post_layernorm.bias'] = post_ln_bias

    embed_w = get_weight(model_params, 'transformer.word_embeddings', dtype)
    weights['lm_head.weight'] = split_matrix_tp(embed_w.clone(),
                                                tensor_parallel,
                                                rank,
                                                dim=0)

    if not use_parallel_embedding:
        weights['transformer.vocab_embedding.weight'] = embed_w
    else:
        assert hf_bloom.config.vocab_size % tensor_parallel == 0
        weights['transformer.vocab_embedding.weight'] = split_matrix_tp(
            embed_w, tensor_parallel, rank, dim=sharding_dim)

    embed_f_w, embed_f_b = get_weight_and_bias(
        model_params, 'transformer.word_embeddings_layernorm', dtype)
    weights['transformer.ln_embed.weight'] = embed_f_w
    weights['transformer.ln_embed.bias'] = embed_f_b

    ln_f_w, ln_f_b = get_weight_and_bias(model_params, 'transformer.ln_f',
                                         dtype)
    weights['transformer.ln_f.weight'] = ln_f_w
    weights['transformer.ln_f.bias'] = ln_f_b

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


def rename_hf_to_tllm(name: str):
    """ Rename a HF parameter name by the corresponding TRT-LLM style name. """
    if 'word_embeddings_layernorm.' in name:
        name = name.replace('word_embeddings_layernorm', 'ln_embed')
        if not name.startswith('transformer.'):
            name = f'transformer.{name}'
    elif 'word_embeddings.' in name:
        name = name.replace('word_embeddings', 'vocab_embedding')
    if name.startswith(('ln_embed.', 'vocab_embedding.', 'ln_f.')):
        name = f'transformer.{name}'

    # Parameter names in layers
    if name.startswith(('transformer.h.', 'h.')):
        import re
        name = re.sub(r'^(transformer.h.|h.)', 'transformer.layers.', name, 1)
    if 'post_attention_layernorm' in name:
        name = name.replace('post_attention_layernorm', 'post_layernorm')
    elif 'self_attention.query_key_value' in name:
        name = name.replace('self_attention.query_key_value', 'attention.qkv')
    elif 'self_attention.dense' in name:
        name = name.replace('self_attention.dense', 'attention.dense')
    elif 'mlp.dense_h_to_4h' in name:
        name = name.replace('mlp.dense_h_to_4h', 'mlp.fc')
    elif 'mlp.dense_4h_to_h' in name:
        name = name.replace('mlp.dense_4h_to_h', 'mlp.proj')
    return name


def contain_any(name: str, words: Iterable[str]):
    for word in words:
        if word in name:
            return True
    return False


def convert_from_hf_checkpoint(
    model_dir: Union[str, Path],
    rank=0,
    tensor_parallel=1,
    dtype: Union[str, torch.dtype] = torch.float32,
    use_parallel_embedding: bool = False,
    sharding_dim: int = 0,
    use_weight_only: bool = False,
    plugin_weight_only_quant_type: torch.dtype = torch.int8,
    use_smooth_quant: bool = False,
    bloom_qkv_param: Optional[Dict] = None,
    act_range: Optional[Any] = None,
    smoother: Optional[Any] = None,
    per_channel: bool = False,
    per_token: bool = False,
    int8_kv_cache: bool = False,
):
    logger.info('Loading weights from HF BLOOM...')
    tik = time.time()

    weights = {}
    hf_config = BloomConfig.from_pretrained(model_dir)
    num_heads = hf_config.n_head
    hidden_size = hf_config.hidden_size
    if isinstance(dtype, str):
        dtype = tensorrt_llm.str_dtype_to_torch(dtype)
    tp_rank = rank
    tp_size = tensor_parallel

    if use_smooth_quant:
        quant_mode = QuantMode.use_smooth_quant(per_token, per_channel)
    elif use_weight_only:
        quant_mode = QuantMode.from_description(
            quantize_weights=True,
            quantize_activations=False,
            per_token=False,
            per_channel=False,
            use_int8_kv_cache=int8_kv_cache,
            use_int4_weights=plugin_weight_only_quant_type == torch.quint4x2)
    else:
        quant_mode = QuantMode(0)

    def is_bias(_name):
        return 'bias' in _name

    for model_file in iterate_shard_files(model_dir, tp_rank):
        logger.debug(f'Loading file {str(model_file)}...')
        model_params = load_state_dict(model_file, dtype=dtype)
        for name, param in model_params.items():
            logger.debug(f'Converting weight {name}...')
            tllm_name = rename_hf_to_tllm(name)
            param = param.detach().cpu()

            # TODO: Support SmmothQuant.

            if 'self_attention.query_key_value' in name:
                if not is_bias(name):
                    param = split_qkv_tp(param, num_heads, hidden_size, tp_size,
                                         tp_rank)
                    # TODO: Add KV scalers when quantizing KV cache.
                else:
                    param = split_qkv_bias_tp(param, num_heads, hidden_size,
                                              tp_size, tp_rank)
                add_tllm_weight(weights, tllm_name, param, quant_mode)
            elif 'self_attention.dense' in name:
                if not is_bias(name):
                    param = split_matrix_tp(param, tp_size, tp_rank, dim=1)
                add_tllm_weight(weights, tllm_name, param, quant_mode)
            elif 'mlp.dense_h_to_4h' in name:
                if not is_bias(name):
                    param = split_matrix_tp(param, tp_size, tp_rank, dim=0)
                else:
                    param = split_matrix_tp(param, tp_size, tp_rank, dim=0)
                add_tllm_weight(weights, tllm_name, param, quant_mode)
            elif 'mlp.dense_4h_to_h' in name:
                if not is_bias(name):
                    param = split_matrix_tp(param, tp_size, tp_rank, dim=1)
                add_tllm_weight(weights, tllm_name, param, quant_mode)
            elif 'word_embeddings.' in name:
                # TODO: safetensor doesn't allow to save a shared tensor.
                # Currently, we clone the weight but to save the disk, it
                # would be better to skip saving lm_head weights and
                # handle it at the loading phase.
                lm_head = split_matrix_tp(param, tp_size, tp_rank, dim=0)
                weights['lm_head.weight'] = lm_head.clone()

                if not use_parallel_embedding:
                    weights[tllm_name] = param
                else:
                    assert hf_config.vocab_size % tp_size == 0
                    weights[tllm_name] = split_matrix_tp(param,
                                                         tp_size,
                                                         tp_rank,
                                                         dim=sharding_dim)
            elif contain_any(name,
                             ('input_layernorm', 'post_attention_layernorm',
                              'word_embeddings_layernorm.', 'ln_f.')):
                weights[tllm_name] = param
        del model_params

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


def do_convert_from_ckpt(args):
    return (args.model_dir.exists() and args.smoothquant is None
            and not args.use_weight_only and not args.int8_kv_cache)


def convert(worker_rank, world_size, args, convert_args):
    convert_from_ckpt = do_convert_from_ckpt(args)
    for rank in range(worker_rank, world_size, args.workers):
        if convert_from_ckpt:
            weights = convert_from_hf_checkpoint(rank=rank, **convert_args)
        else:
            weights = convert_hf_bloom(rank=rank, **convert_args)
        safetensors.torch.save_file(weights,
                                    args.output_dir / f'rank{rank}.safetensors')


def main():
    # TODO(qijun): Currently, the convert script depends on a torch op:
    # torch.ops.trtllm.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."

    logger.set_level(args.log_level)
    tik = time.time()

    args.output_dir.mkdir(exist_ok=True, parents=True)

    quant_algo = None
    plugin_weight_only_quant_type = None
    if args.use_weight_only and args.weight_only_precision == 'int8':
        plugin_weight_only_quant_type = torch.int8
        quant_algo = QuantAlgo.W8A16
    elif args.use_weight_only and args.weight_only_precision == 'int4':
        plugin_weight_only_quant_type = torch.quint4x2
        quant_algo = QuantAlgo.W4A16
    elif args.smoothquant:
        if args.per_channel and args.per_token:
            quant_algo = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN
        elif args.per_channel and not args.per_token:
            quant_algo = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN
        elif not args.per_channel and args.per_token:
            quant_algo = QuantAlgo.W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN
        else:
            quant_algo = QuantAlgo.W8A8_SQ_PER_TENSOR_PLUGIN

    kv_cache_quant_algo = None
    if args.int8_kv_cache:
        kv_cache_quant_algo = QuantAlgo.INT8

    hf_config = BloomConfig.from_pretrained(args.model_dir)
    config = {
        'architecture': hf_config.architectures[0],
        'dtype': args.dtype,
        'logits_dtype': 'float32',
        'num_hidden_layers': hf_config.num_hidden_layers,
        'num_attention_heads': hf_config.num_attention_heads,
        'hidden_size': hf_config.hidden_size,
        'norm_epsilon': hf_config.layer_norm_epsilon,
        'vocab_size': hf_config.vocab_size,
        'position_embedding_type': 'alibi',
        'hidden_act': 'gelu',
        'intermediate_size': hf_config.hidden_size * 4,
        'quantization': {
            'quant_algo': quant_algo,
            'kv_cache_quant_algo': kv_cache_quant_algo,
        },
        'mapping': {
            'world_size': world_size,
            'tp_size': args.tp_size,
            'pp_size': args.pp_size,
        },
        'use_parallel_embedding': args.use_parallel_embedding,
        'embedding_sharding_dim': args.embedding_sharding_dim,
    }

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

    # TODO: convert_from_hf_checkpoint is memory efficient but has not
    # supported quantization yet. Will enable once implemented.
    convert_from_ckpt = do_convert_from_ckpt(args)
    if not convert_from_ckpt:
        logger.info(f'Convert by using model')
        hf_bloom = BloomForCausalLM.from_pretrained(args.model_dir,
                                                    torch_dtype="auto",
                                                    device_map="auto",
                                                    trust_remote_code=True)
    else:
        logger.info(f'Convert by using checkpoint')
        hf_bloom = None

    act_range = {}
    bloom_qkv_param = {}
    bloom_smoother = {}

    if args.smoothquant is not None or args.int8_kv_cache:
        os.environ["TOKENIZERS_PARALLELISM"] = os.environ.get(
            "TOKENIZERS_PARALLELISM", "false")
        tokenizer = BloomTokenizerFast.from_pretrained(args.model_dir)
        dataset = load_calib_dataset(args.calib_dataset)

        act_range = capture_activation_range(hf_bloom, tokenizer, dataset)
        if args.smoothquant is not None:
            smooth_bloom_model(hf_bloom, act_range, args.smoothquant,
                               bloom_qkv_param, bloom_smoother)

    convert_args = dict(
        tensor_parallel=args.tp_size,
        dtype=args.dtype,
        use_weight_only=args.use_weight_only,
        plugin_weight_only_quant_type=plugin_weight_only_quant_type,
        use_parallel_embedding=args.use_parallel_embedding,
        sharding_dim=args.embedding_sharding_dim,
        use_smooth_quant=args.smoothquant,
        act_range=act_range,
        bloom_qkv_param=bloom_qkv_param,
        smoother=bloom_smoother,
        per_channel=args.per_channel,
        per_token=args.per_token,
        int8_kv_cache=args.int8_kv_cache,
    )
    if convert_from_ckpt:
        convert_args['model_dir'] = args.model_dir
    else:
        convert_args['hf_bloom'] = hf_bloom

    if args.workers == 1:
        convert(0, world_size, args, convert_args)
    else:
        if args.workers > world_size:
            args.workers = world_size
        logger.info(f'Convert checkpoint using {args.workers} workers.')
        import torch.multiprocessing as mp
        mp.spawn(convert,
                 nprocs=args.workers,
                 args=(world_size, args, convert_args))

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


if __name__ == '__main__':
    main()
