TensorRT-LLMs/tensorrt_llm/models/chatglm/convert.py

import copy
import functools
import os
import time
from collections import defaultdict
from typing import Dict, Optional, Tuple

import numpy as np
import safetensors
import torch
from tqdm import tqdm
from transformers import AutoModel, AutoTokenizer

from tensorrt_llm._utils import pad_vocab_size
from tensorrt_llm.models import ChatGLMConfig
from tensorrt_llm.models.convert_utils import load_calib_dataset
from tensorrt_llm.quantization import QuantAlgo

from .config import GLM_ARCH1_VERSIONS, GLM_ARCH2_VERSIONS


def split(weight: torch.Tensor,
          tp_size: int,
          rank: int = 0,
          dim: int = 0) -> torch.Tensor:
    if tp_size == 1:
        return weight
    elif weight.ndim == 1:
        return torch.chunk(weight, tp_size)[rank].contiguous()
    else:
        return torch.chunk(weight, tp_size, dim=dim)[rank].contiguous()


def tile_kv_weight_bias(v: torch.Tensor, kv_num_head: int, tp_size: int):
    head_size = v.shape[0] // kv_num_head
    reps = tp_size // kv_num_head
    if v.ndim == 1:
        v = v.reshape(kv_num_head, head_size)[:, None, :]
        v = v.expand(kv_num_head, reps, head_size).reshape(-1).clone()
    else:
        hidden_size = v.shape[1]
        v = v.reshape(kv_num_head, head_size, hidden_size)[:, None, :, :]
        v = v.expand(kv_num_head, reps, head_size,
                     hidden_size).reshape(-1, hidden_size).clone()
    return v


def split_qkv(v: torch.Tensor, tp_size: int, rank: int, hidden_size: int,
              num_heads: int, num_kv_heads: int):
    head_size = hidden_size // num_heads
    if tp_size == 1:
        return v

    assert v.shape[0] == hidden_size + head_size * num_kv_heads * 2
    query = v[:hidden_size]
    key = v[hidden_size:hidden_size + head_size * num_kv_heads]
    value = v[hidden_size + head_size * num_kv_heads:hidden_size +
              head_size * num_kv_heads * 2]

    if num_kv_heads < tp_size:
        key = tile_kv_weight_bias(key, num_kv_heads, tp_size)
        value = tile_kv_weight_bias(value, num_kv_heads, tp_size)
    assert (key.shape[0] % (tp_size * head_size)) == 0
    assert (value.shape[0] % (tp_size * head_size)) == 0

    q_tmp = torch.chunk(query, tp_size, dim=0)[rank]
    k_tmp = torch.chunk(key, tp_size, dim=0)[rank]
    v_tmp = torch.chunk(value, tp_size, dim=0)[rank]
    return torch.concatenate([q_tmp, k_tmp, v_tmp], dim=0).contiguous()


def split_embedding(
    param: torch.Tensor,
    tp_size: int,
    tp_rank: int,
    use_parallel_embedding: bool = False,
    sharding_dim: int = 0,
) -> torch.Tensor:
    if param is None:
        return None
    if not use_parallel_embedding:
        return param

    vocab_size, hidden_size = param.size()
    if sharding_dim == 0:
        if vocab_size % tp_size != 0:
            vocab_size_padded = pad_vocab_size(vocab_size, tp_size)
            pad_width = vocab_size_padded - vocab_size
            param = torch.nn.functional.pad(param, (0, 0, 0, pad_width),
                                            value=0)
        else:
            assert hidden_size % tp_size == 0
    return split(param, tp_size, tp_rank, dim=sharding_dim)


def get_weight(params: Dict[str, torch.Tensor], prefix: str,
               dtype: torch.dtype) -> torch.Tensor:
    if f'{prefix}.weight' not in params:
        return None
    return params[f'{prefix}.weight'].to(dtype).detach().cpu()


def get_bias(params: Dict[str, torch.Tensor], prefix: str,
             dtype: torch.dtype) -> torch.Tensor:
    if f'{prefix}.bias' not in params:
        return None
    return params[f'{prefix}.bias'].to(dtype).detach().cpu()


def get_weight_and_bias(params: Dict[str, torch.Tensor], prefix: str,
                        dtype: torch.dtype) -> Tuple[torch.Tensor]:
    return get_weight(params, prefix, dtype), get_bias(params, prefix, dtype)


def swap_and_split_mlp(weight: torch.Tensor,
                       tp_size: int,
                       tp_rank: int,
                       dim: int = 0) -> torch.Tensor:
    """Swap the positions of gate and fc weights, and split weights for tensor parallel.
    """
    gate_weight, fc_weight = torch.chunk(weight, 2, dim=dim)
    fc_w = split(fc_weight, tp_size, tp_rank, dim=dim)
    gate_w = split(gate_weight, tp_size, tp_rank, dim=dim)
    return torch.cat([fc_w, gate_w], dim=dim).contiguous()


def get_tllm_linear_weight(
    weight: torch.Tensor,
    prefix: str,
    bias: Optional[torch.Tensor] = None,
    use_weight_only: bool = False,
    plugin_weight_only_quant_type: torch.dtype = torch.int8
) -> Dict[str, torch.Tensor]:
    results = {}
    if use_weight_only:
        v = weight.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[f'{prefix}.weight'] = processed_torch_weights
        results[f'{prefix}.per_channel_scale'] = torch_weight_scales
    else:
        results[f'{prefix}.weight'] = weight.contiguous()

    if bias is not None:
        results[f'{prefix}.bias'] = bias

    return results


@torch.no_grad()
def apply_smoothing(
    scales,
    gemm_weights,
    norm_weights=None,
    norm_bias=None,
    dtype=torch.float32,
    norm_1p=False,
):
    if not isinstance(gemm_weights, list):
        gemm_weights = [gemm_weights]

    if norm_weights is not None:
        assert norm_weights.numel() == scales.numel()
        norm_weights.div_(scales).to(dtype)
    if norm_bias is not None:
        assert norm_bias.numel() == scales.numel()
        norm_bias.div_(scales).to(dtype)
    if norm_1p:
        norm_weights += (1 / scales) - 1

    for gemm in gemm_weights:
        gemm.mul_(scales.view(1, -1)).to(dtype)


@torch.no_grad()
def smooth_gemm(
    gemm_weights,
    act_scales,
    norm_weights=None,
    norm_bias=None,
    alpha=0.5,
    weight_scales=None,
):
    if not isinstance(gemm_weights, list):
        gemm_weights = [gemm_weights]
    orig_dtype = gemm_weights[0].dtype

    for gemm in gemm_weights:
        # gemm_weights are expected to be transposed
        assert gemm.shape[1] == act_scales.numel()

    if weight_scales is None:
        weight_scales = torch.cat(
            [gemm.abs().max(dim=0, keepdim=True)[0] for gemm in gemm_weights],
            dim=0)
        weight_scales = weight_scales.max(dim=0)[0]
    weight_scales.to(float).clamp(min=1e-5)
    scales = (act_scales.to(gemm_weights[0].device).to(float).pow(alpha) /
              weight_scales.pow(1 - alpha)).clamp(min=1e-5)

    apply_smoothing(scales, gemm_weights, norm_weights, norm_bias, orig_dtype)

    return scales


@torch.no_grad()
def capture_activation_range(
    model,
    tokenizer,
    dataset,
    num_samples=64,
    seq_len=512,
):

    model.eval()
    device = next(model.parameters()).device
    scales = defaultdict(lambda: {"x": None, "y": None, "w": None})

    def stat_tensor(name, tensor, key):
        tensor = tensor.view(-1, tensor.shape[-1]).detach()
        comming_max = tensor.abs().max(dim=0)[0].float()
        if scales[name][key] is None:
            scales[name][key] = comming_max
        else:
            scales[name][key] = torch.max(scales[name][key], comming_max)

    def stat_input_hook(m, x, y, name):
        if isinstance(x, tuple):
            x = x[0]
        stat_tensor(name, x, "x")
        stat_tensor(name, y, "y")
        # TODO: we don't need to do it every forward because inference does not change weight
        if scales[name]["w"] is None:
            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, torch.nn.Linear):
            hooks.append(
                m.register_forward_hook(
                    functools.partial(stat_input_hook, name=name)))

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

    for h in hooks:
        h.remove()

    return scales


def generate_int8(weights, act_range, is_qkv=False, multi_query_mode=True):
    """
     This function has two purposes:
      - compute quantized weights, scaled either per-tensor or per-column
      - compute scaling factors

      Depending on the GEMM API (CUTLASS/CUBLAS) the required scaling factors differ.
      CUTLASS uses two sets of scaling factors. One for the activation X, one for the weight W.
      CUBLAS only has one (we can't do per-row scaling). So we must provide pre-multiplied scaling factor.

      Here is the list of what we need (T means per-tensor, C per-column):
        - scale_x_orig_quant puts fp activation into the quantized range (i.e. [-128, 127], for int8). Used before the GEMM. (T)
        - scale_y_quant_orig puts quantized activation into the fp range. Used if the GEMM outputs int8. (T)
        - scale_w_quant_orig puts weights from quant range to fp range (used with CUTLASS) (T, C)
        - scale_y_accum_quant puts the GEMM result (XW) from accumulation range (int32)
          to quant range (int8) (used for CUBLAS) (T, C)

      Note that we don't do anything special about row-parallel GEMM. Theoretically, we could have per-GPU scaling factors too,
      but then the model would change depending on the number of GPUs used.

      For QKV projection, the behavior is special. Even if we have a single matrix to perform QKV projection, we consider it
      as three different matrices: Q, K, and V. So per-tensor actually means one scaling factor for each Q, K and V.
      For our GEMM implementation to respect this behavior, we use per-column mode and replicate values along columns.
    """

    # For ChatGLM2/3-6B models (num_kv_head == 2), we regard multi_query_mode == True to reuse code from gpt and baichuan examples.
    if act_range["w"].dtype == torch.bfloat16:
        act_range["w"] = act_range["w"].to(torch.float32)

    def to_np(tensor):
        return tensor.cpu().numpy().astype(np.float32)

    if is_qkv and multi_query_mode:
        hidden_dim, local_dim = weights.shape

        kv_dim = (local_dim - hidden_dim) // 2
        scale_w_q = act_range["w"][0:hidden_dim]
        scale_w_k = act_range["w"][hidden_dim:hidden_dim + kv_dim]
        scale_w_v = act_range["w"][-kv_dim:]

        scale_w_qkv_t = torch.concat([
            scale_w_q.max(dim=0, keepdim=True)[0],
            scale_w_k.max(dim=0, keepdim=True)[0],
            scale_w_v.max(dim=0, keepdim=True)[0]
        ])
        scale_w_orig_quant_t = 127. / to_np(scale_w_qkv_t)
        scale_w_orig_quant_c = 127. / to_np(act_range["w"])
    elif is_qkv and not multi_query_mode:
        scale_w_orig_quant_t = 127. / to_np(act_range["w"].reshape(3, -1).max(
            dim=-1, keepdims=True)[0])
        scale_w_orig_quant_c = 127. / to_np(act_range["w"].reshape(3, -1))

    else:
        scale_w_orig_quant_t = 127. / to_np(act_range["w"].max())
        scale_w_orig_quant_c = 127. / to_np(act_range["w"])
    scale_w_quant_orig_t = 1.0 / scale_w_orig_quant_t
    scale_w_quant_orig_c = 1.0 / scale_w_orig_quant_c

    # compute the rest of needed scaling factors
    act_range_x_max = act_range["x"].max().to(torch.float32).item()
    act_range_y_max = act_range["y"].max().to(torch.float32).item()
    scale_x_orig_quant_t = np.array(127. / act_range_x_max).astype(np.float32)
    scale_y_orig_quant_t = np.array(127. / act_range_y_max).astype(np.float32)
    scale_y_quant_orig_t = np.array(act_range_y_max / 127.).astype(np.float32)
    scale_y_accum_quant_t = scale_y_orig_quant_t / (scale_x_orig_quant_t *
                                                    scale_w_orig_quant_t)
    scale_y_accum_quant_c = scale_y_orig_quant_t / (scale_x_orig_quant_t *
                                                    scale_w_orig_quant_c)
    if is_qkv and not multi_query_mode:
        scale_y_accum_quant_t = np.broadcast_to(scale_y_accum_quant_t,
                                                scale_w_orig_quant_c.shape)
        scale_w_quant_orig_t = np.broadcast_to(scale_w_quant_orig_t,
                                               scale_w_orig_quant_c.shape)
    if is_qkv and multi_query_mode:
        scale_q_y_accum_t = np.broadcast_to(scale_y_accum_quant_t[0],
                                            scale_w_q.shape)
        scale_k_y_accum_t = np.broadcast_to(scale_y_accum_quant_t[1],
                                            scale_w_k.shape)
        scale_v_y_accum_t = np.broadcast_to(scale_y_accum_quant_t[2],
                                            scale_w_v.shape)
        scale_y_accum_quant_t = np.concatenate(
            [scale_q_y_accum_t, scale_k_y_accum_t, scale_v_y_accum_t])
        scale_w_quant_orig_t = np.concatenate([
            np.broadcast_to(scale_w_quant_orig_t[0], scale_w_q.shape),
            np.broadcast_to(scale_w_quant_orig_t[1], scale_w_k.shape),
            np.broadcast_to(scale_w_quant_orig_t[2], scale_w_v.shape)
        ])

    to_i8 = lambda x: x.round().clip(-127, 127).astype(np.int8)

    if is_qkv and multi_query_mode:
        scale_w_orig_quant_t_expand = np.ones([weights.shape[-1]])
        scale_w_orig_quant_t_expand[:hidden_dim] = scale_w_orig_quant_t[0]
        scale_w_orig_quant_t_expand[hidden_dim:hidden_dim +
                                    kv_dim] = scale_w_orig_quant_t[1]
        scale_w_orig_quant_t_expand[-kv_dim:] = scale_w_orig_quant_t[2]
        weight_int8 = to_i8(weights * scale_w_orig_quant_t_expand)
    else:
        weight_int8 = to_i8(weights * scale_w_orig_quant_t)
    return {
        "weight.int8": weight_int8,
        "weight.int8.col": to_i8(weights * scale_w_orig_quant_c),
        "scale_x_orig_quant": scale_x_orig_quant_t.astype(np.float32),
        "scale_w_quant_orig": scale_w_quant_orig_t.astype(np.float32),
        "scale_w_quant_orig.col": scale_w_quant_orig_c.astype(np.float32),
        "scale_y_accum_quant": scale_y_accum_quant_t.astype(np.float32),
        "scale_y_accum_quant.col": scale_y_accum_quant_c.astype(np.float32),
        "scale_y_quant_orig": scale_y_quant_orig_t.astype(np.float32),
    }


@torch.no_grad()
def smooth_chatglm_model(
    model,
    act_range,
    alpha,
    model_smoother,
):
    for name, module in model.named_modules():
        if not module._get_name() == "GLMBlock":
            continue

        # QKV multiplication weight
        layer_name = name + '.self_attention.query_key_value'
        print(f'Smoothing module: {layer_name}')
        weight = module.self_attention.query_key_value.weight
        smoother = smooth_gemm(
            weight,
            act_range[layer_name]["x"],
            module.input_layernorm.weight,
            None,
            alpha,
        )
        act_range[layer_name]["x"] = act_range[layer_name]["x"] / smoother
        act_range[layer_name]["w"] = weight.abs().max(dim=1)[0]

        # Dense multiplication weight
        layer_name = name + ".self_attention.dense"
        print(f'Smoothing module: {layer_name}')
        weight = module.self_attention.dense.weight
        smoother = smooth_gemm(
            weight,
            act_range[layer_name]["x"],
            None,
            None,
            alpha,
        )
        model_smoother[layer_name] = smoother.float()
        act_range[layer_name]["x"] = act_range[layer_name]["x"] / smoother
        act_range[layer_name]["w"] = weight.abs().max(dim=1)[0]

        # Multilayer perceptron h -> 4h weight
        layer_name = name + ".mlp.dense_h_to_4h"
        print(f'Smoothing module: {layer_name}')
        weight = module.mlp.dense_h_to_4h.weight
        smoother = smooth_gemm(
            weight,
            act_range[layer_name]["x"],
            module.post_attention_layernorm.weight,
            None,
            alpha,
        )
        act_range[layer_name]["x"] = act_range[layer_name]["x"] / smoother
        act_range[layer_name]["w"] = weight.abs().max(dim=1)[0]

        # Multilayer perceptron 4h -> h weight
        layer_name = name + ".mlp.dense_4h_to_h"
        print(f'Smoothing module: {layer_name}')
        weight = module.mlp.dense_4h_to_h.weight
        smoother = smooth_gemm(
            weight,
            act_range[layer_name]["x"],
            None,
            None,
            alpha,
        )
        model_smoother[layer_name] = smoother.float()
        act_range[layer_name]["x"] = act_range[layer_name]["x"] / smoother
        act_range[layer_name]["w"] = weight.abs().max(dim=1)[0]


def get_tllm_linear_sq_weight(vals,
                              prefix,
                              shape,
                              is_qkv=False,
                              per_token=False,
                              per_channel=False,
                              last_prefix=None,
                              smoother_value=None,
                              smoother_shape=None):
    results = {}
    col_shape = shape if (is_qkv or per_channel) else [1, 1]

    if per_token:
        if per_channel:
            original_weights = np.array(vals["weight.int8.col"])
        else:
            original_weights = np.array(vals["weight.int8"])

        cur_weights = original_weights
        if is_qkv:
            hidden_dim = cur_weights.shape[0]
            cur_weights = cur_weights.reshape(hidden_dim, -1)
        results[prefix +
                'weight'] = torch.from_numpy(cur_weights).t().contiguous()
        if smoother_value is None:
            results[last_prefix] = torch.from_numpy(
                np.array([1.0], dtype=np.float16))

        if per_channel:
            cur_per_channel_value = vals["scale_w_quant_orig.col"]
        else:
            cur_per_channel_value = vals["scale_w_quant_orig"]

        results[prefix + 'per_channel_scale'] = torch.from_numpy(
            np.array(cur_per_channel_value,
                     dtype=np.float32).reshape(col_shape)).contiguous()
    else:
        if per_channel:
            original_weights = np.array(vals["weight.int8.col"])
        else:
            original_weights = np.array(vals["weight.int8"])
        cur_weights = original_weights

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

        if per_channel:
            cur_per_channel_value = vals["scale_y_accum_quant.col"]
        else:
            cur_per_channel_value = vals["scale_y_accum_quant"]

        results[prefix + 'per_channel_scale'] = torch.from_numpy(
            np.array([cur_per_channel_value],
                     dtype=np.float32).reshape(col_shape)).contiguous()

        results[last_prefix] = torch.from_numpy(
            np.array([vals['scale_x_orig_quant']],
                     dtype=np.float32)).contiguous()

        results[prefix + 'act_scale'] = torch.from_numpy(
            np.array([[vals["scale_y_quant_orig"]]],
                     dtype=np.float32)).contiguous()

    if smoother_value is not None:
        results[prefix + 'smoother'] = smoother_value.reshape(
            smoother_shape).contiguous().to(torch.float32)

    return results


def smooth_quant(hf_model: AutoModel, tokenizer: AutoTokenizer, dataset,
                 smoothquant: float):
    act_range = {}
    # smoother for query_key_value.dense and mlp.proj
    model_smoother = {}
    act_range = capture_activation_range(hf_model,
                                         tokenizer,
                                         dataset,
                                         num_samples=64)
    if smoothquant is not None:
        smooth_chatglm_model(hf_model, act_range, smoothquant, model_smoother)
    return act_range, model_smoother


def load_weights_from_hf_model(hf_model: AutoModel,
                               config: ChatGLMConfig,
                               act_range: Optional[dict] = None,
                               smoother: Optional[dict] = None):
    weights = {}
    tik = time.time()

    model_params = dict(hf_model.named_parameters())
    dtype = getattr(torch, config.dtype)
    num_attention_heads = config.num_attention_heads
    hidden_size = config.hidden_size
    num_kv_heads = config.num_key_value_heads
    num_hidden_layers = config.num_hidden_layers

    chatglm_version = config.chatglm_version
    mapping = config.mapping
    use_parallel_embedding = config.use_parallel_embedding
    sharding_dim = config.embedding_sharding_dim
    share_embedding_table = config.share_embedding_table

    quant_algo = config.quantization.quant_algo
    use_weight_only = quant_algo in [QuantAlgo.W8A16, QuantAlgo.W4A16]
    use_smooth_quant = config.quantization.use_plugin_sq
    per_channel = use_smooth_quant and 'PER_CHANNEL' in quant_algo
    per_token = use_smooth_quant and 'PER_TOKEN' in quant_algo
    int8_kv_cache = config.quantization.kv_cache_quant_algo == QuantAlgo.INT8

    if use_weight_only:
        if quant_algo == QuantAlgo.W8A16:
            plugin_weight_only_quant_type = torch.int8
        elif quant_algo == QuantAlgo.W4A16:
            plugin_weight_only_quant_type = torch.quint4x2
    else:
        plugin_weight_only_quant_type = None

    layers_range = mapping.pp_layers(num_hidden_layers)
    for l in layers_range:
        if chatglm_version in GLM_ARCH1_VERSIONS:
            prefix = f'transformer.layers.{l}'
        elif chatglm_version in GLM_ARCH2_VERSIONS:
            prefix = f'transformer.encoder.layers.{l}'
        tllm_prex = f'transformer.layers.{l-layers_range[0]}'

        # Attention QKV
        attention_attr_name = ''
        if chatglm_version in GLM_ARCH1_VERSIONS:
            attention_attr_name = 'attention'
        elif chatglm_version in GLM_ARCH2_VERSIONS:
            attention_attr_name = 'self_attention'
        qkv_weight, qkv_bias = get_weight_and_bias(
            model_params, f'{prefix}.{attention_attr_name}.query_key_value',
            dtype)

        if use_smooth_quant:
            qkv_act_range = act_range.get(
                f'{prefix}.{attention_attr_name}.query_key_value')
            qkv_vals_int8 = generate_int8(qkv_weight.t().numpy(),
                                          qkv_act_range,
                                          is_qkv=True,
                                          multi_query_mode=True)
            weights.update(
                get_tllm_linear_sq_weight(
                    vals=qkv_vals_int8,
                    prefix=f'{tllm_prex}.attention.qkv.',
                    shape=[1, qkv_weight.size(0)],
                    is_qkv=True,
                    per_token=per_token,
                    per_channel=per_channel,
                    last_prefix=f'{tllm_prex}.input_layernorm.scale_to_int',
                    smoother_value=None,
                    smoother_shape=None))
            if qkv_bias is not None:
                qkv_b = split_qkv(qkv_bias,
                                  mapping.tp_size,
                                  mapping.tp_rank,
                                  hidden_size,
                                  num_attention_heads,
                                  num_kv_heads=num_kv_heads)
                weights[f'{tllm_prex}.attention.qkv.bias'] = qkv_b
        else:
            qkv_w = split_qkv(qkv_weight,
                              mapping.tp_size,
                              mapping.tp_rank,
                              hidden_size,
                              num_attention_heads,
                              num_kv_heads=num_kv_heads)
            if qkv_bias is None:
                qkv_b = None
            else:
                qkv_b = split_qkv(qkv_bias,
                                  mapping.tp_size,
                                  mapping.tp_rank,
                                  hidden_size,
                                  num_attention_heads,
                                  num_kv_heads=num_kv_heads)

            weights.update(
                get_tllm_linear_weight(qkv_w, f'{tllm_prex}.attention.qkv',
                                       qkv_b, use_weight_only,
                                       plugin_weight_only_quant_type))

        if int8_kv_cache:
            qkv_act_range = act_range.get(
                f'{prefix}.{attention_attr_name}.query_key_value')
            qkv_vals_int8 = generate_int8(qkv_weight.t().numpy(),
                                          qkv_act_range,
                                          is_qkv=True,
                                          multi_query_mode=True)
            weights[
                f'{tllm_prex}.attention.kv_cache_scaling_factor'] = torch.from_numpy(
                    np.array([qkv_vals_int8['scale_y_quant_orig']],
                             dtype=np.float32)).contiguous()

        # Attention dense
        attn_dense_weight, attn_dense_bias = get_weight_and_bias(
            model_params, f'{prefix}.{attention_attr_name}.dense', dtype)

        if use_smooth_quant:
            dense_act_range = act_range.get(
                f'{prefix}.{attention_attr_name}.dense')
            dense_smoother = smoother.get(
                f'{prefix}.{attention_attr_name}.dense')
            dense_vals_int8 = generate_int8(attn_dense_weight.t().numpy(),
                                            dense_act_range,
                                            is_qkv=False,
                                            multi_query_mode=True)
            weights.update(
                get_tllm_linear_sq_weight(
                    vals=dense_vals_int8,
                    prefix=f'{tllm_prex}.attention.dense.',
                    shape=[1, hidden_size],
                    is_qkv=False,
                    per_token=per_token,
                    per_channel=per_channel,
                    last_prefix=
                    f'{tllm_prex}.attention.quantization_scaling_factor',
                    smoother_value=dense_smoother,
                    smoother_shape=[1, hidden_size]))
            if attn_dense_bias is not None:
                weights[f'{tllm_prex}.attention.dense.bias'] = attn_dense_bias
        else:
            attn_dense_w = split(attn_dense_weight,
                                 mapping.tp_size,
                                 mapping.tp_rank,
                                 dim=1)
            weights.update(
                get_tllm_linear_weight(attn_dense_w,
                                       f'{tllm_prex}.attention.dense',
                                       attn_dense_bias, use_weight_only,
                                       plugin_weight_only_quant_type))

        # MLP FC
        mlp_fc_weight, mlp_fc_bias = get_weight_and_bias(
            model_params, f'{prefix}.mlp.dense_h_to_4h', dtype)

        if use_smooth_quant:
            fc_act_range = act_range.get(f'{prefix}.mlp.dense_h_to_4h')
            fc_vals_int8 = generate_int8(mlp_fc_weight.t().numpy(),
                                         fc_act_range,
                                         is_qkv=False,
                                         multi_query_mode=True)
            cur_weights = get_tllm_linear_sq_weight(
                vals=fc_vals_int8,
                prefix=f'{tllm_prex}.mlp.fc.',
                shape=[1, mlp_fc_weight.size(0)],
                is_qkv=False,
                per_token=per_token,
                per_channel=per_channel,
                last_prefix=f'{tllm_prex}.post_layernorm.scale_to_int',
                smoother_value=None,
                smoother_shape=None,
            )
            cur_weights[f'{tllm_prex}.mlp.fc.weight'] = swap_and_split_mlp(
                cur_weights[f'{tllm_prex}.mlp.fc.weight'],
                mapping.tp_size,
                mapping.tp_rank,
                dim=0,
            )
            if per_channel:
                cur_weights[
                    f'{tllm_prex}.mlp.fc.per_channel_scale'] = swap_and_split_mlp(
                        cur_weights[f'{tllm_prex}.mlp.fc.per_channel_scale'],
                        mapping.tp_size,
                        mapping.tp_rank,
                        dim=1,
                    )
            weights.update(cur_weights)
            if chatglm_version in GLM_ARCH1_VERSIONS:
                if mlp_fc_bias is not None:
                    mlp_fc_b = split(mlp_fc_bias,
                                     mapping.tp_size,
                                     mapping.tp_rank,
                                     dim=0)
                    weights[f'{tllm_prex}.mlp.fc.bias'] = mlp_fc_b
            elif chatglm_version in GLM_ARCH2_VERSIONS:
                if mlp_fc_bias is not None:
                    mlp_fc_b = swap_and_split_mlp(mlp_fc_bias, mapping.tp_size,
                                                  mapping.tp_rank)
                    weights[f'{tllm_prex}.mlp.fc.bias'] = mlp_fc_b
        else:
            if chatglm_version in GLM_ARCH1_VERSIONS:
                mlp_fc_w = split(mlp_fc_weight,
                                 mapping.tp_size,
                                 mapping.tp_rank,
                                 dim=0)
                if mlp_fc_bias is None:
                    mlp_fc_b = None
                else:
                    mlp_fc_b = split(mlp_fc_bias,
                                     mapping.tp_size,
                                     mapping.tp_rank,
                                     dim=0)
            elif chatglm_version in GLM_ARCH2_VERSIONS:
                mlp_fc_w = swap_and_split_mlp(mlp_fc_weight, mapping.tp_size,
                                              mapping.tp_rank)
                if mlp_fc_bias is None:
                    mlp_fc_b = None
                else:
                    mlp_fc_b = swap_and_split_mlp(mlp_fc_bias, mapping.tp_size,
                                                  mapping.tp_rank)
            weights.update(
                get_tllm_linear_weight(mlp_fc_w, f'{tllm_prex}.mlp.fc',
                                       mlp_fc_b, use_weight_only,
                                       plugin_weight_only_quant_type))

        # MLP Proj
        mlp_proj_weight, mlp_proj_bias = get_weight_and_bias(
            model_params, f'{prefix}.mlp.dense_4h_to_h', dtype)

        if use_smooth_quant:
            proj_act_range = act_range.get(f'{prefix}.mlp.dense_4h_to_h')
            proj_smoother = smoother.get(f'{prefix}.mlp.dense_4h_to_h')
            proj_vals_int8 = generate_int8(mlp_proj_weight.t().numpy(),
                                           proj_act_range,
                                           is_qkv=False,
                                           multi_query_mode=True)
            weights.update(
                get_tllm_linear_sq_weight(
                    vals=proj_vals_int8,
                    prefix=f'{tllm_prex}.mlp.proj.',
                    shape=[1, hidden_size],
                    is_qkv=False,
                    per_token=per_token,
                    per_channel=per_channel,
                    last_prefix=f'{tllm_prex}.mlp.quantization_scaling_factor',
                    smoother_value=proj_smoother,
                    smoother_shape=[1, config.intermediate_size]))
            if mlp_proj_bias is not None:
                weights[f'{tllm_prex}.mlp.proj.bias'] = mlp_proj_bias
        else:
            mlp_proj_w = split(mlp_proj_weight,
                               mapping.tp_size,
                               mapping.tp_rank,
                               dim=1)
            weights.update(
                get_tllm_linear_weight(mlp_proj_w, f'{tllm_prex}.mlp.proj',
                                       mlp_proj_bias, use_weight_only,
                                       plugin_weight_only_quant_type))

        input_ln_weight, input_ln_bias = get_weight_and_bias(
            model_params, f'{prefix}.input_layernorm', dtype)
        weights[f'{tllm_prex}.input_layernorm.weight'] = input_ln_weight
        if input_ln_bias is not None:
            weights[f'{tllm_prex}.input_layernorm.bias'] = input_ln_bias

        post_ln_weight, post_ln_bias = get_weight_and_bias(
            model_params, f'{prefix}.post_attention_layernorm', dtype)
        weights[f'{tllm_prex}.post_layernorm.weight'] = post_ln_weight
        if post_ln_bias is not None:
            weights[f'{tllm_prex}.post_layernorm.bias'] = post_ln_bias

    if mapping.is_first_pp_rank():
        if chatglm_version == 'glm':
            embed_w = get_weight(model_params, 'word_embeddings', dtype)
            pos_embed_w = get_weight(model_params,
                                     'transformer.position_embeddings', dtype)
            weights['transformer.position_embedding.weight'] = split_embedding(
                pos_embed_w,
                tp_size=mapping.tp_size,
                tp_rank=mapping.tp_rank,
                use_parallel_embedding=use_parallel_embedding,
                sharding_dim=sharding_dim)
            block_embed_w = get_weight(model_params,
                                       'transformer.block_position_embeddings',
                                       dtype)
            weights['transformer.block_embedding.weight'] = split_embedding(
                block_embed_w,
                tp_size=mapping.tp_size,
                tp_rank=mapping.tp_rank,
                use_parallel_embedding=use_parallel_embedding,
                sharding_dim=sharding_dim)
        elif chatglm_version == 'chatglm':
            embed_w = get_weight(model_params, 'transformer.word_embeddings',
                                 dtype)
        elif chatglm_version in GLM_ARCH2_VERSIONS:
            embed_w = get_weight(model_params,
                                 'transformer.embedding.word_embeddings', dtype)

        weights['transformer.vocab_embedding.weight'] = split_embedding(
            embed_w,
            tp_size=mapping.tp_size,
            tp_rank=mapping.tp_rank,
            use_parallel_embedding=use_parallel_embedding,
            sharding_dim=sharding_dim)

    if mapping.is_last_pp_rank():
        if chatglm_version == 'glm':
            lm_head_weight = get_weight(model_params, 'word_embeddings',
                                        dtype).clone()
        elif chatglm_version == 'chatglm':
            lm_head_weight = get_weight(model_params,
                                        'transformer.word_embeddings',
                                        dtype).clone()
        elif chatglm_version in GLM_ARCH2_VERSIONS:
            lm_head_weight = get_weight(model_params,
                                        'transformer.output_layer', dtype)
            assert not share_embedding_table

        if not share_embedding_table:
            weights['lm_head.weight'] = split(lm_head_weight,
                                              mapping.tp_size,
                                              mapping.tp_rank,
                                              dim=0)

        if chatglm_version in GLM_ARCH1_VERSIONS:
            ln_f_w, ln_f_b = get_weight_and_bias(model_params,
                                                 'transformer.final_layernorm',
                                                 dtype)
        elif chatglm_version in GLM_ARCH2_VERSIONS:
            ln_f_w, ln_f_b = get_weight_and_bias(
                model_params, 'transformer.encoder.final_layernorm', dtype)
        weights['transformer.ln_f.weight'] = ln_f_w
        if ln_f_b is not None:
            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 quantize(hf_model_dir: str,
             output_dir: str,
             config: ChatGLMConfig,
             calib_dataset: str = 'cnn_dailymail',
             device: str = 'auto'):
    '''
        Quantize the save the model as TRT-LLM checkpoint to output_dir
    '''
    os.makedirs(output_dir, exist_ok=True)
    config.to_json_file(os.path.join(output_dir, 'config.json'))

    mapping = config.mapping
    assert mapping.rank == 0, "quantize should be called at rank 0 only"

    quant_config = config.quantization
    use_smooth_quant = quant_config.use_plugin_sq
    int8_kv_cache = quant_config.kv_cache_quant_algo == QuantAlgo.INT8

    assert use_smooth_quant or int8_kv_cache, "Call from_hugging_face when there is no quantization"
    if use_smooth_quant:
        assert quant_config.smoothquant_val is not None, "A smooth value must be specified when using smooth quant"

    assert hf_model_dir is not None
    ## only load and call smooth quant routine once for all ranks
    if config.chatglm_version == 'glm':
        device_map = 'cuda' if device != "cpu" else 'cpu'
    else:
        device_map = 'auto' if device != "cpu" else 'cpu'
    hf_model = AutoModel.from_pretrained(
        hf_model_dir,
        trust_remote_code=True,
        torch_dtype='auto' if config.chatglm_version != 'glm' else getattr(
            torch, config.dtype),
        device_map=device_map)

    os.environ["TOKENIZERS_PARALLELISM"] = os.environ.get(
        "TOKENIZERS_PARALLELISM", "false")
    tokenizer = AutoTokenizer.from_pretrained(
        hf_model_dir,
        trust_remote_code=True,
    )
    dataset = load_calib_dataset(calib_dataset)
    act_range, smoother = smooth_quant(hf_model, tokenizer, dataset,
                                       quant_config.smoothquant_val)

    for rank in range(mapping.world_size):
        # To avoid changing the mapping arg in-place, also the given mapping from caller is rank agnostic, since quantize is called from only one rank
        config = copy.deepcopy(config)
        config.set_rank(rank)
        weights = load_weights_from_hf_model(
            hf_model,
            config=config,
            act_range=act_range,
            smoother=smoother,
        )
        safetensors.torch.save_file(
            weights, os.path.join(output_dir, f'rank{rank}.safetensors'))
        del weights