# 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 copy
import functools
import os
import sys
import time
from collections import defaultdict
from pathlib import Path
from typing import List, Optional

import numpy as np
import safetensors
import torch
import torch.nn as nn
from tqdm import tqdm
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
from transformers.models.llama.modeling_llama import LlamaDecoderLayer
from transformers.pytorch_utils import Conv1D

from ..._utils import pad_vocab_size, release_gc, str_dtype_to_torch
from ...logger import logger
from ...quantization import QuantAlgo
from ...quantization.quantize import (qserve_pack_reorder_per_channel,
                                      qserve_pack_reorder_per_group,
                                      qserve_quantize_weight_per_channel,
                                      qserve_quantize_weight_per_group)
from ..convert_utils import (dup_kv_weight, generate_int8,
                             get_tllm_linear_weight, iterate_shard_files,
                             load_calib_dataset, load_state_dict,
                             retrieved_layer_index_from_name, smooth_gemm,
                             smooth_gemm_fc1_gate, split, split_matrix_tp,
                             split_qkv_bias_tp, split_qkv_tp)
from ..modeling_utils import PretrainedConfig
from .config import LLaMAConfig


@torch.no_grad()
def smooth_llama_model(model, scales, alpha, llama_qkv_para, llama_smoother):
    # Smooth the activation and weights with smoother = $\diag{s}$
    for name, module in model.named_modules():
        if not isinstance(
                module,
                LlamaDecoderLayer) and not module.__class__.__name__ in [
                    "InternLMDecoderLayer", "MistralDecoderLayer"
                ]:
            continue
        # qkv_proj
        layer_name_q = name + ".self_attn.q_proj"
        layer_name_k = name + ".self_attn.k_proj"
        layer_name_v = name + ".self_attn.v_proj"
        layer_name_qkv = name + ".self_attn.qkv_proj"

        weight = torch.cat([
            module.self_attn.q_proj.weight, module.self_attn.k_proj.weight,
            module.self_attn.v_proj.weight
        ],
                           dim=0)

        smoother = smooth_gemm(weight, scales[layer_name_q]["x"],
                               module.input_layernorm.weight, None, alpha)

        scales[layer_name_qkv]["x"] = scales[layer_name_q]["x"] / smoother
        scales[layer_name_qkv]["w"] = weight.abs().max(dim=1)[0]
        scales[layer_name_qkv]["y"] = torch.cat([
            scales[layer_name_q]["y"], scales[layer_name_k]["y"],
            scales[layer_name_v]["y"]
        ],
                                                dim=0)

        # see transpose_weights function
        llama_qkv_para[layer_name_qkv] = weight.transpose(0, 1)

        # =================================================================
        layer_name = name + ".self_attn.o_proj"
        smoother = smooth_gemm(module.self_attn.o_proj.weight,
                               scales[layer_name]["x"], None, None, alpha)
        llama_smoother[layer_name] = smoother.float()

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

        # ==================================================================
        fc1_layer_name = name + ".mlp.gate_proj"
        gate_layer_name = name + ".mlp.up_proj"

        smoother = smooth_gemm_fc1_gate(module.mlp.gate_proj.weight,
                                        module.mlp.up_proj.weight,
                                        scales[fc1_layer_name]["x"],
                                        module.post_attention_layernorm.weight,
                                        None, alpha)

        scales[fc1_layer_name]["x"] = scales[fc1_layer_name]["x"] / smoother
        scales[fc1_layer_name]["w"] = module.mlp.gate_proj.weight.abs().max(
            dim=1)[0]

        scales[gate_layer_name]["x"] = scales[gate_layer_name]["x"] / smoother
        scales[gate_layer_name]["w"] = module.mlp.up_proj.weight.abs().max(
            dim=1)[0]

        # ==================================================================
        layer_name = name + ".mlp.down_proj"
        smoother = smooth_gemm(module.mlp.down_proj.weight,
                               scales[layer_name]["x"], None, None, alpha)
        llama_smoother[layer_name] = smoother.float()
        scales[layer_name]["x"] = scales[layer_name]["x"] / smoother
        scales[layer_name]["w"] = module.mlp.down_proj.weight.abs().max(
            dim=1)[0]

        # ==================================================================
        if hasattr(module, 'residual_mlp'):
            fc1_layer_name = name + ".residual_mlp.w1"
            gate_layer_name = name + ".residual_mlp.w3"

            smoother = smooth_gemm_fc1_gate(module.residual_mlp.w1.weight,
                                            module.residual_mlp.w3.weight,
                                            scales[fc1_layer_name]["x"],
                                            module.residual_layernorm.weight,
                                            None, alpha)

            scales[fc1_layer_name]["x"] = scales[fc1_layer_name]["x"] / smoother
            scales[fc1_layer_name]["w"] = module.residual_mlp.w1.weight.abs(
            ).max(dim=1)[0]

            scales[gate_layer_name][
                "x"] = scales[gate_layer_name]["x"] / smoother
            scales[gate_layer_name]["w"] = module.residual_mlp.w3.weight.abs(
            ).max(dim=1)[0]

            # ==================================================================
            layer_name = name + ".residual_mlp.w2"
            smoother = smooth_gemm(module.residual_mlp.w2.weight,
                                   scales[layer_name]["x"], None, None, alpha)
            llama_smoother[layer_name] = smoother.float()
            scales[layer_name]["x"] = scales[layer_name]["x"] / smoother
            scales[layer_name]["w"] = module.residual_mlp.w2.weight.abs().max(
                dim=1)[0]


@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})

    tokenizer.pad_token = tokenizer.eos_token

    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"):
        datapoint = dataset[i:i + 1]
        line = copy.copy(datapoint)
        line[0] = line[0] + ' TL;DR: '
        line[0] = line[0].strip()
        line[0] = line[0].replace(" n't", "n't")
        input_ids = tokenizer(line,
                              return_tensors="pt",
                              max_length=seq_len,
                              padding=True,
                              truncation=True).input_ids.to(device)
        model(input_ids)
    for h in hooks:
        h.remove()
    return act_scales


def get_weight(named_params, prefix, dtype):
    if named_params[prefix + '.weight'].dtype != dtype:
        named_params[prefix +
                     '.weight'].data = named_params[prefix +
                                                    '.weight'].to(dtype)
    return named_params[prefix + '.weight'].detach()


def get_weight_and_scale(named_params,
                         prefix,
                         dtype,
                         mapping=None,
                         split_scale=False):
    if prefix + '.weight_scale' not in named_params:
        return get_weight(named_params, prefix, dtype), None
    else:
        assert named_params[prefix + '.weight'].dtype == torch.float8_e4m3fn
        assert named_params[prefix + '.weight_scale'].dtype == torch.float32
        weight_scale = named_params[prefix + '.weight_scale'].detach()
        if split_scale:
            weight_scale = split(weight_scale,
                                 mapping.tp_size,
                                 mapping.tp_rank,
                                 dim=0)
        return named_params[prefix +
                            '.weight'].detach(), weight_scale.reshape(-1)


def get_bias(named_params, prefix, dtype):
    if named_params[prefix + '.bias'].dtype != dtype:
        named_params[prefix + '.bias'].data = named_params[prefix +
                                                           '.bias'].to(dtype)
    return named_params[prefix + '.bias'].detach()


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


def fp8_per_channel_quant_weight_gpu(weight, clamp_val, rank=0):
    weight = weight.to("cuda:" + str(rank))
    # activation range bound.
    x = weight.to(torch.float32).clamp(clamp_val[0], clamp_val[1])
    xmax = x.abs().max(-1, keepdim=True).values
    # minimum scaling factor.
    torch_weight_scales = (xmax / 448.0).clamp(min=1.0 / (448.0 * 512.0))
    out = x / torch_weight_scales
    torch_weight_scales = torch_weight_scales.reshape(-1)
    out = torch.clamp(out, -448, 448)
    processed_torch_weights = out.to(torch.float8_e4m3fn)

    processed_torch_weights = processed_torch_weights.to(
        torch.float8_e4m3fn).cpu()
    torch_weight_scales = torch_weight_scales.cpu()

    return processed_torch_weights, torch_weight_scales


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,
                              multi_query_mode=False):
    results = {}

    def multi_query_split(data, local_dim, head_size, tp_size, cur_rank):
        q, k, v = torch.split(data, [local_dim, head_size, head_size], dim=-1)
        q_split = torch.chunk(q, tp_size, dim=-1)
        k_split = torch.chunk(k, tp_size, dim=-1)
        v_split = torch.chunk(v, tp_size, dim=-1)
        return [
            torch.concat((q_split[ii], k_split[ii], v_split[ii]), axis=-1)
            for ii in range(tp_size)
        ][cur_rank]

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

    if per_token:
        if per_channel:
            original_weights = torch.Tensor(vals["weight.int8.col"])
        else:
            original_weights = torch.Tensor(vals["weight.int8"])
        local_dim = original_weights.shape[0]
        head_size = (original_weights.shape[1] - local_dim) // 2

        if multi_query_mode:
            cur_weights = multi_query_split(original_weights, local_dim,
                                            head_size, tensor_parallel, rank)
        else:
            cur_weights = torch.chunk(original_weights,
                                      tensor_parallel,
                                      dim=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.Tensor([1.0]).to(torch.float32)

        if per_channel:
            cur_per_channel_value = vals["scale_w_quant_orig.col"]
            if smoother_value is None:
                if multi_query_mode:
                    cur_per_channel_value = multi_query_split(
                        vals["scale_w_quant_orig.col"], local_dim, head_size,
                        tensor_parallel, rank)
                else:
                    cur_per_channel_value = torch.chunk(
                        vals["scale_w_quant_orig.col"],
                        tensor_parallel,
                        dim=cat_dim)[rank]
        else:
            cur_per_channel_value = vals["scale_w_quant_orig"]
            if is_qkv:
                if multi_query_mode:
                    cur_per_channel_value = multi_query_split(
                        vals["scale_w_quant_orig"], local_dim, head_size,
                        tensor_parallel, rank)
                else:
                    cur_per_channel_value = torch.chunk(
                        vals["scale_w_quant_orig"],
                        tensor_parallel,
                        dim=cat_dim)[rank]

        results[prefix + 'per_channel_scale'] = cur_per_channel_value.reshape(
            col_shape).contiguous()
    else:
        if per_channel:
            original_weights = torch.Tensor(vals["weight.int8.col"])
        else:
            original_weights = torch.Tensor(vals["weight.int8"])
        local_dim = original_weights.shape[0]
        head_size = (original_weights.shape[1] - local_dim) // 2

        if multi_query_mode:
            cur_weights = multi_query_split(original_weights, local_dim,
                                            head_size, tensor_parallel, rank)
        else:
            cur_weights = torch.chunk(original_weights,
                                      tensor_parallel,
                                      dim=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 per_channel:
            cur_per_channel_value = vals["scale_y_accum_quant.col"]
            if smoother_value is None:
                if multi_query_mode:
                    cur_per_channel_value = multi_query_split(
                        vals["scale_y_accum_quant.col"], local_dim, head_size,
                        tensor_parallel, rank)
                else:
                    cur_per_channel_value = torch.chunk(
                        vals["scale_y_accum_quant.col"],
                        tensor_parallel,
                        dim=cat_dim)[rank]
        else:
            cur_per_channel_value = vals["scale_y_accum_quant"]
            # QKV is always per_channel
            if is_qkv:
                if multi_query_mode:
                    cur_per_channel_value = multi_query_split(
                        vals["scale_y_accum_quant"], local_dim, head_size,
                        tensor_parallel, rank)
                else:
                    cur_per_channel_value = torch.chunk(
                        vals["scale_y_accum_quant"],
                        tensor_parallel,
                        dim=cat_dim)[rank]

        results[prefix +
                'per_channel_scale'] = torch.Tensor(cur_per_channel_value).to(
                    torch.float32).reshape(col_shape).contiguous()
        results[prefix + 'act_scale'] = torch.Tensor(
            [[vals['scale_y_quant_orig']]]).to(torch.float32).contiguous()
        results[last_prefix] = torch.Tensor([vals['scale_x_orig_quant']
                                             ]).to(torch.float32).contiguous()

    if smoother_value is not None:
        cur_smoother_value = torch.chunk(smoother_value,
                                         tensor_parallel,
                                         dim=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 get_prefix_and_param_name_map(architecture, use_safetensors=False):

    key_postfix = ""
    if use_safetensors:
        key_postfix = ".weight"

    architecture = architecture.lower()
    if "exaone" in architecture:
        model_prefix = "transformer"
        param_name_map = {
            "vocab_embedding": "wte" + key_postfix,  # vocab_embedding
            "lm_head": "lm_head" + key_postfix,  # lm_head
            "ln_f": "ln_f" + key_postfix,  # ln_f
            "attention.qkv": "attn.attention",  # attention.qkv
            "qkv_suffix": "_proj" + key_postfix,  # qkv_suffix
            "attention.dense":
            "attn.attention.out_proj" + key_postfix,  # attention.dense
            "mlp.gate": "mlp.c_fc_1" + key_postfix,  # mlp.gate
            "mlp.proj": "mlp.c_proj" + key_postfix,  # mlp.proj
            "mlp.fc": "mlp.c_fc_0" + key_postfix,  # mlp.fc
            "input_layernorm": "ln_1" + key_postfix,  # input_layernorm
            "post_layernorm": "ln_2" + key_postfix,  # post_layernorm
        }
        layer_prefix = 'h'
    else:  # LLaMA
        model_prefix = "model"
        param_name_map = {
            "vocab_embedding": "embed_tokens" + key_postfix,  # vocab_embedding
            "lm_head": "lm_head" + key_postfix,  # lm_head
            "ln_f": "norm" + key_postfix,  # ln_f
            "attention.qkv": "self_attn",  # attention.qkv
            "qkv_suffix": "_proj" + key_postfix,  # qkv suffix
            "attention.dense":
            "self_attn.o_proj" + key_postfix,  # attention.dense
            "mlp.gate": "mlp.up_proj" + key_postfix,  # mlp.gate
            "mlp.proj": "mlp.down_proj" + key_postfix,  # mlp.proj
            "mlp.fc": "mlp.gate_proj" + key_postfix,  # mlp.fc
            "input_layernorm":
            "input_layernorm" + key_postfix,  # input_layernorm
            "post_layernorm":
            "post_attention_layernorm" + key_postfix,  # post_layernorm
        }
        layer_prefix = 'layers'

    return model_prefix, layer_prefix, param_name_map


def load_hf_llama(model_dir: str, load_model_on_cpu: bool = False):
    if "vila" in model_dir:
        sys.path.append(model_dir + "/../VILA")
        from llava.model import LlavaLlamaConfig, LlavaLlamaModel  # noqa
        from transformers import AutoModel
        model = AutoModel.from_pretrained(
            model_dir,
            device_map='auto',
            trust_remote_code=True,
        )
        return model.llm

    hf_config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True)
    model_cls = AutoModelForCausalLM
    if hf_config.model_type == "llava":
        from transformers import LlavaForConditionalGeneration
        model_cls = LlavaForConditionalGeneration
    if hf_config.model_type == "llava_next":
        from transformers import LlavaNextForConditionalGeneration
        model_cls = LlavaNextForConditionalGeneration
    model = model_cls.from_pretrained(
        model_dir,
        device_map='auto' if not load_model_on_cpu else 'cpu',
        torch_dtype='auto',
        trust_remote_code=True,
    )
    if hf_config.model_type in ["llava", "llava_next"]:
        model = model.language_model
    return model


def load_weights_from_hf_model(hf_model,
                               config: LLaMAConfig,
                               act_range: Optional[dict] = None,
                               qkv_para: Optional[dict] = None,
                               smoother: Optional[dict] = None):
    quant_algo = config.quantization.quant_algo
    use_weight_only = quant_algo in [QuantAlgo.W8A16, QuantAlgo.W4A16]
    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
    use_gemm_woq_plugin = (not config.disable_weight_only_quant_plugin)
    use_fp8_rowwise = quant_algo in [QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN]

    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
    fp8_kv_cache = config.quantization.kv_cache_quant_algo == QuantAlgo.FP8
    if use_smooth_quant or int8_kv_cache:
        assert act_range is not None
        assert qkv_para is not None
        assert smoother is not None

    weights = {}
    tik = time.time()
    model_params = dict(hf_model.named_parameters())
    dtype = getattr(torch, config.dtype)

    mapping = config.mapping
    moe_config = config.moe
    mha_mode = (config.num_key_value_heads == config.num_attention_heads)
    layers_range = config.mapping.pp_layers(config.num_hidden_layers)
    exclude_layers_id = [0, config.num_hidden_layers - 1]

    model_prefix, layer_prefix, param_name_map = get_prefix_and_param_name_map(
        config.architecture)

    def convert_layer(l):
        prefix = f'{model_prefix}.{layer_prefix}.{l}.'
        tllm_prex = f'transformer.layers.{l - layers_range[0]}.'
        q_weight = get_weight(
            model_params, prefix + f'{param_name_map["attention.qkv"]}.q_proj',
            dtype)
        k_weight = get_weight(
            model_params, prefix + f'{param_name_map["attention.qkv"]}.k_proj',
            dtype)
        v_weight = get_weight(
            model_params, prefix + f'{param_name_map["attention.qkv"]}.v_proj',
            dtype)

        # Meta's recipe of not using fp8 rowwise for the first and last layer.
        use_fp8_rowwise_in_layer = use_fp8_rowwise and (
            l not in exclude_layers_id)

        if not mha_mode:
            if config.num_key_value_heads < mapping.tp_size:
                # duplicate the KV heads up to tensor_parallel
                k_weight = dup_kv_weight(k_weight, config.num_key_value_heads,
                                         mapping.tp_size)
                v_weight = dup_kv_weight(v_weight, config.num_key_value_heads,
                                         mapping.tp_size)
            assert (k_weight.shape[0] %
                    (mapping.tp_size * config.head_size)) == 0
            assert (v_weight.shape[0] %
                    (mapping.tp_size * config.head_size)) == 0

            wq = split(q_weight, mapping.tp_size, mapping.tp_rank)
            wk = split(k_weight, mapping.tp_size, mapping.tp_rank)
            wv = split(v_weight, mapping.tp_size, mapping.tp_rank)

            split_v = torch.concat((wq, wk, wv))

        else:
            qkv_weight = torch.cat([q_weight, k_weight, v_weight], dim=0)

            split_v = split_qkv_tp(qkv_weight, config.num_attention_heads,
                                   config.hidden_size, mapping.tp_size,
                                   mapping.tp_rank)

        if prefix + f'{param_name_map["attention.qkv"]}.q_proj.bias' in model_params:
            # only used in Internlm 7B models
            q_bias = get_bias(
                model_params,
                prefix + f'{param_name_map["attention.qkv"]}.q_proj', dtype)
            k_bias = get_bias(
                model_params,
                prefix + f'{param_name_map["attention.qkv"]}.k_proj', dtype)
            v_bias = get_bias(
                model_params,
                prefix + f'{param_name_map["attention.qkv"]}.v_proj', dtype)
            qkv_bias = torch.cat((q_bias, k_bias, v_bias))
            split_bias_v = split_qkv_bias_tp(qkv_bias,
                                             config.num_attention_heads,
                                             config.hidden_size,
                                             mapping.tp_size, mapping.tp_rank)
        else:
            split_bias_v = None

        if use_smooth_quant:
            qkv_weight = qkv_para[prefix +
                                  f'{param_name_map["attention.qkv"]}.qkv_proj']
            qkv_out_dim = qkv_weight.shape[1]

            if not mha_mode:
                local_dim = qkv_weight.shape[0]
                kv_hidden_size = (qkv_weight.shape[-1] - local_dim) // 2
                qkv_weight = qkv_weight.reshape(local_dim,
                                                local_dim + 2 * kv_hidden_size)
            else:
                qkv_weight = qkv_weight.reshape(config.hidden_size, 3,
                                                config.hidden_size)

            int8_weights = generate_int8(
                qkv_weight,
                act_range.get(prefix +
                              f'{param_name_map["attention.qkv"]}.qkv_proj'),
                is_qkv=True,
                multi_query_mode=bool(not mha_mode))

            weights.update(
                get_tllm_linear_sq_weight(int8_weights,
                                          tllm_prex + 'attention.qkv.',
                                          [1, qkv_out_dim // mapping.tp_size],
                                          mapping.tp_size,
                                          is_qkv=True,
                                          bias=split_bias_v,
                                          per_token=per_token,
                                          per_channel=per_channel,
                                          last_prefix=tllm_prex +
                                          'input_layernorm.scale_to_int',
                                          smoother_value=None,
                                          smoother_shape=None,
                                          rank=mapping.tp_rank,
                                          cat_dim=-1,
                                          multi_query_mode=bool(not mha_mode)))
        else:
            weights.update(
                get_tllm_linear_weight(split_v,
                                       tllm_prex + 'attention.qkv.',
                                       split_bias_v,
                                       use_weight_only,
                                       plugin_weight_only_quant_type,
                                       dtype,
                                       use_gemm_woq_plugin,
                                       use_fp8_rowwise=False))

        if int8_kv_cache:
            qkv_y = torch.cat([
                act_range.get(prefix +
                              f'{param_name_map["attention.qkv"]}.q_proj')["y"],
                act_range.get(prefix +
                              f'{param_name_map["attention.qkv"]}.k_proj')["y"],
                act_range.get(prefix +
                              f'{param_name_map["attention.qkv"]}.v_proj')["y"]
            ],
                              dim=0)

            int8_kv_scales = qkv_y.max() / 127.

            kv_cache_weights = {}

            kv_cache_weights[
                tllm_prex +
                'attention.kv_cache_scaling_factor'] = int8_kv_scales.reshape(
                    [1])

            weights.update(kv_cache_weights)
        elif fp8_kv_cache:
            # FIXME: set it to 1.0f for fp8 kv cache.
            weights[tllm_prex +
                    'attention.kv_cache_scaling_factor'] = torch.tensor(
                        [1.0], dtype=torch.float32)

        attn_dense_weight = get_weight(
            model_params, prefix + param_name_map["attention.dense"], dtype)
        split_v = split_matrix_tp(attn_dense_weight,
                                  mapping.tp_size,
                                  mapping.tp_rank,
                                  dim=1)

        if prefix + f'{param_name_map["attention.dense"]}.bias' in model_params:
            attn_dense_bias = get_bias(
                model_params, prefix + param_name_map["attention.dense"], dtype)
        else:
            attn_dense_bias = None
        if use_smooth_quant:
            attn_dense_weight = attn_dense_weight.t()
            proj_out_dim = attn_dense_weight.shape[0]

            int8_weights = generate_int8(
                attn_dense_weight,
                act_range.get(prefix + param_name_map["attention.dense"]))
            weights.update(
                get_tllm_linear_sq_weight(
                    int8_weights,
                    tllm_prex + 'attention.dense.', [1, config.hidden_size],
                    mapping.tp_size,
                    is_qkv=False,
                    bias=attn_dense_bias,
                    per_token=per_token,
                    per_channel=per_channel,
                    last_prefix=tllm_prex +
                    'attention.quantization_scaling_factor',
                    smoother_value=smoother[(
                        prefix + param_name_map["attention.dense"])],
                    smoother_shape=[1, proj_out_dim // mapping.tp_size],
                    rank=mapping.tp_rank,
                    cat_dim=0))
        else:
            weights.update(
                get_tllm_linear_weight(split_v,
                                       tllm_prex + 'attention.dense.',
                                       attn_dense_bias,
                                       use_weight_only,
                                       plugin_weight_only_quant_type,
                                       dtype,
                                       use_gemm_woq_plugin,
                                       use_fp8_rowwise=False))

        if moe_config.has_moe():
            rank_experts = list(range(moe_config.num_experts))
            if mapping.has_moe_ep():
                rank_experts = mapping.ep_experts(moe_config.num_experts)
            architecture = config.architecture.lower()
            if "granite" not in architecture:
                for suffix in ["w1", "w2", "w3"]:
                    model_params[f'model.layers.{l}.block_sparse_moe.experts.{suffix}.weight'] = \
                                torch.stack([model_params[f'model.layers.{l}.block_sparse_moe.experts.{expert}.{suffix}.weight'].detach()
                                            for expert in rank_experts])
                w3 = model_params[
                    f'model.layers.{l}.block_sparse_moe.experts.w3.weight']
                w2 = model_params[
                    f'model.layers.{l}.block_sparse_moe.experts.w2.weight']
                w1 = model_params[
                    f'model.layers.{l}.block_sparse_moe.experts.w1.weight']
            else:
                w2 = model_params[
                    f'model.layers.{l}.block_sparse_moe.output_linear.weight']
                half_size = model_params[
                    f'model.layers.{l}.block_sparse_moe.input_linear.weight'].shape[
                        -2] // 2
                w1, w3 = model_params[
                    f'model.layers.{l}.block_sparse_moe.input_linear.weight']\
                        .split(half_size, dim=-2)
                w1 = w1[rank_experts[0]:rank_experts[-1] + 1]
                w2 = w2[rank_experts[0]:rank_experts[-1] + 1]
                w3 = w3[rank_experts[0]:rank_experts[-1] + 1]

            if mapping.has_moe_tp():
                w3 = split(w3, mapping.moe_tp_size, mapping.moe_tp_rank, dim=1)
                w2 = split(w2, mapping.moe_tp_size, mapping.moe_tp_rank, dim=2)
                w1 = split(w1, mapping.moe_tp_size, mapping.moe_tp_rank, dim=1)

            model_params[
                f'model.layers.{l}.block_sparse_moe.experts.w3w1.weight'] = torch.concat(
                    [w3, w1], dim=-2)

            model_params[
                f'model.layers.{l}.block_sparse_moe.experts.w2.weight'] = w2

            ## block_sparse_moe.experts.w2.weight
            moe_experts_w2_weights = get_weight(
                model_params, prefix + 'block_sparse_moe.experts.w2', dtype)
            weights.update(
                get_tllm_linear_weight(moe_experts_w2_weights,
                                       tllm_prex + 'mlp.proj.', None,
                                       use_weight_only,
                                       plugin_weight_only_quant_type, dtype,
                                       use_gemm_woq_plugin))
            ##block_sparse_moe.experts.w3w1.weight
            moe_experts_w3w1_weights = get_weight(
                model_params, prefix + 'block_sparse_moe.experts.w3w1', dtype)
            weights.update(
                get_tllm_linear_weight(moe_experts_w3w1_weights,
                                       tllm_prex + 'mlp.fc.', None,
                                       use_weight_only,
                                       plugin_weight_only_quant_type, dtype,
                                       use_gemm_woq_plugin))

            if config.residual_mlp:
                residual_mlp_gate_weights = get_weight(
                    model_params, prefix + 'residual_mlp.w3', dtype)
                if use_smooth_quant:
                    residual_mlp_gate_weights = residual_mlp_gate_weights.t()
                    int8_weights = generate_int8(
                        residual_mlp_gate_weights,
                        act_range.get(prefix + 'residual_mlp.w3'))
                    weights.update(
                        get_tllm_linear_sq_weight(
                            int8_weights,
                            tllm_prex + 'residual_mlp.gate.',
                            [1, config.hidden_size // mapping.tp_size],
                            mapping.tp_size,
                            is_qkv=False,
                            per_token=per_token,
                            per_channel=per_channel,
                            last_prefix=tllm_prex +
                            'post_layernorm.scale_to_int',
                            smoother_value=None,
                            smoother_shape=None,
                            rank=mapping.tp_rank,
                            cat_dim=-1))
                else:
                    split_v = split_matrix_tp(residual_mlp_gate_weights,
                                              mapping.tp_size,
                                              mapping.tp_rank,
                                              dim=0)
                    weights.update(
                        get_tllm_linear_weight(split_v,
                                               tllm_prex + 'residual_mlp.gate.',
                                               None, use_weight_only,
                                               plugin_weight_only_quant_type,
                                               dtype, use_gemm_woq_plugin))

                residual_mlp_fc_weight = get_weight(model_params,
                                                    prefix + 'residual_mlp.w1',
                                                    dtype)
                if use_smooth_quant:
                    residual_mlp_fc_weight = residual_mlp_fc_weight.t(
                    )  #verified
                    int8_weights = generate_int8(
                        residual_mlp_fc_weight,
                        act_range.get(prefix + 'residual_mlp.w1'))
                    weights.update(
                        get_tllm_linear_sq_weight(
                            int8_weights,
                            tllm_prex + 'residual_mlp.fc.',
                            [1, config.hidden_size // mapping.tp_size],
                            mapping.tp_size,
                            is_qkv=False,
                            per_token=per_token,
                            per_channel=per_channel,
                            last_prefix=tllm_prex +
                            'post_layernorm.scale_to_int',
                            smoother_value=None,
                            smoother_shape=None,
                            rank=mapping.tp_rank,
                            cat_dim=-1))
                else:
                    split_v = split_matrix_tp(residual_mlp_fc_weight,
                                              mapping.tp_size,
                                              mapping.tp_rank,
                                              dim=0)
                    weights.update(
                        get_tllm_linear_weight(split_v,
                                               tllm_prex + 'residual_mlp.fc.',
                                               None, use_weight_only,
                                               plugin_weight_only_quant_type,
                                               dtype, use_gemm_woq_plugin))

                residual_mlp_proj_weight = get_weight(
                    model_params, prefix + 'residual_mlp.w2', dtype)

                if use_smooth_quant:
                    residual_mlp_proj_weight = residual_mlp_proj_weight.t()
                    int8_weights = generate_int8(
                        residual_mlp_proj_weight,
                        act_range.get(prefix + 'residual_mlp.w2'))
                    weights.update(
                        get_tllm_linear_sq_weight(
                            int8_weights,
                            tllm_prex + 'residual_mlp.proj.',
                            [1, config.hidden_size],
                            mapping.tp_size,
                            is_qkv=False,
                            per_token=per_token,
                            per_channel=per_channel,
                            last_prefix=tllm_prex +
                            'residual_mlp.quantization_scaling_factor',
                            smoother_value=smoother[prefix + 'residual_mlp.w2'],
                            smoother_shape=[
                                1, config.hidden_size // mapping.tp_size
                            ],
                            rank=mapping.tp_rank,
                            cat_dim=0))
                else:
                    split_v = split_matrix_tp(residual_mlp_proj_weight,
                                              mapping.tp_size,
                                              mapping.tp_rank,
                                              dim=1)
                    weights.update(
                        get_tllm_linear_weight(split_v,
                                               tllm_prex + 'residual_mlp.proj.',
                                               None, use_weight_only,
                                               plugin_weight_only_quant_type,
                                               dtype, use_gemm_woq_plugin))

            architecture = config.architecture.lower()
            if "granite" not in architecture:
                moe_experts_gate_weights = get_weight(
                    model_params, prefix + 'block_sparse_moe.gate',
                    torch.float32)
            else:
                moe_experts_gate_weights = get_weight(
                    model_params, prefix + 'block_sparse_moe.router.layer',
                    torch.float32)
            weights.update(
                get_tllm_linear_weight(
                    moe_experts_gate_weights,
                    tllm_prex + 'mlp.router.',
                    None,
                    False,  # Router should never be quantized
                    plugin_weight_only_quant_type,
                    dtype,
                    use_gemm_woq_plugin))
        else:
            mlp_gate_weight, mlp_gate_weight_scale = get_weight_and_scale(
                model_params, prefix + param_name_map["mlp.gate"], dtype,
                mapping, True)
            split_v = split_matrix_tp(mlp_gate_weight,
                                      mapping.tp_size,
                                      mapping.tp_rank,
                                      dim=0)
            if use_smooth_quant:

                mlp_gate_weight = mlp_gate_weight.t()
                int8_weights = generate_int8(
                    # mlp_gate_weight, act_range.get(prefix + 'mlp.up_proj'))
                    mlp_gate_weight,
                    act_range.get(prefix + param_name_map["mlp.gate"]))

                weights.update(
                    get_tllm_linear_sq_weight(
                        int8_weights,
                        tllm_prex + 'mlp.gate.',
                        [1, config.intermediate_size // mapping.tp_size],
                        mapping.tp_size,
                        is_qkv=False,
                        per_token=per_token,
                        per_channel=per_channel,
                        last_prefix=tllm_prex + 'post_layernorm.scale_to_int',
                        smoother_value=None,
                        smoother_shape=None,
                        rank=mapping.tp_rank,
                        cat_dim=-1))
            else:
                weights.update(
                    get_tllm_linear_weight(
                        split_v,
                        tllm_prex + 'mlp.gate.',
                        None,
                        use_weight_only,
                        plugin_weight_only_quant_type,
                        dtype,
                        use_gemm_woq_plugin,
                        use_fp8_rowwise_in_layer,
                        weight_scale=mlp_gate_weight_scale,
                        clamp_value=config.quantization.clamp_val))

            mlp_fc_weight, mlp_fc_weight_scale = get_weight_and_scale(
                model_params, prefix + param_name_map["mlp.fc"], dtype, mapping,
                True)
            split_v = split_matrix_tp(mlp_fc_weight,
                                      mapping.tp_size,
                                      mapping.tp_rank,
                                      dim=0)

            if use_smooth_quant:
                mlp_fc_weight = mlp_fc_weight.t()  #verified
                int8_weights = generate_int8(
                    mlp_fc_weight,
                    act_range.get(prefix + param_name_map["mlp.fc"]))
                weights.update(
                    get_tllm_linear_sq_weight(
                        int8_weights,
                        tllm_prex + 'mlp.fc.',
                        [1, config.intermediate_size // mapping.tp_size],
                        mapping.tp_size,
                        is_qkv=False,
                        per_token=per_token,
                        per_channel=per_channel,
                        last_prefix=tllm_prex + 'post_layernorm.scale_to_int',
                        smoother_value=None,
                        smoother_shape=None,
                        rank=mapping.tp_rank,
                        cat_dim=-1))
            else:
                weights.update(
                    get_tllm_linear_weight(
                        split_v,
                        tllm_prex + 'mlp.fc.',
                        None,
                        use_weight_only,
                        plugin_weight_only_quant_type,
                        dtype,
                        use_gemm_woq_plugin,
                        use_fp8_rowwise_in_layer,
                        weight_scale=mlp_fc_weight_scale,
                        clamp_value=config.quantization.clamp_val))

            mlp_proj_weight, mlp_proj_weight_scale = get_weight_and_scale(
                model_params, prefix + param_name_map["mlp.proj"], dtype)
            split_v = split_matrix_tp(mlp_proj_weight,
                                      mapping.tp_size,
                                      mapping.tp_rank,
                                      dim=1)

            if use_smooth_quant:
                mlp_proj_weight = mlp_proj_weight.t()
                int8_weights = generate_int8(
                    mlp_proj_weight,
                    act_range.get(prefix + param_name_map["mlp.proj"]))
                weights.update(
                    get_tllm_linear_sq_weight(
                        int8_weights,
                        tllm_prex + 'mlp.proj.', [1, config.hidden_size],
                        mapping.tp_size,
                        is_qkv=False,
                        per_token=per_token,
                        per_channel=per_channel,
                        last_prefix=tllm_prex +
                        'mlp.quantization_scaling_factor',
                        smoother_value=smoother[prefix +
                                                param_name_map["mlp.proj"]],
                        smoother_shape=[
                            1, config.intermediate_size // mapping.tp_size
                        ],
                        rank=mapping.tp_rank,
                        cat_dim=0))
            else:
                weights.update(
                    get_tllm_linear_weight(
                        split_v,
                        tllm_prex + 'mlp.proj.',
                        None,
                        use_weight_only,
                        plugin_weight_only_quant_type,
                        dtype,
                        use_gemm_woq_plugin,
                        use_fp8_rowwise_in_layer,
                        weight_scale=mlp_proj_weight_scale,
                        clamp_value=config.quantization.clamp_val))

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

        post_ln_weight = get_weight(model_params,
                                    prefix + param_name_map["post_layernorm"],
                                    dtype)
        weights[tllm_prex + 'post_layernorm.weight'] = post_ln_weight

        if config.residual_mlp:
            residual_ln_weight = get_weight(model_params,
                                            prefix + 'residual_layernorm',
                                            dtype)
            weights[tllm_prex +
                    'residual_layernorm.weight'] = residual_ln_weight

        cur_block_weights = [
            weight_name for weight_name in model_params
            if weight_name.find(prefix) != -1
        ]
        for weight_name in cur_block_weights:
            model_params[weight_name] = None

    for l in layers_range:
        convert_layer(l)
        release_gc()

    vocab_embedding = get_weight(
        model_params, f'{model_prefix}.{param_name_map["vocab_embedding"]}',
        dtype)

    if mapping.is_first_pp_rank():
        if config.use_parallel_embedding:
            weights['transformer.vocab_embedding.weight'] = split_matrix_tp(
                vocab_embedding,
                mapping.tp_size,
                mapping.tp_rank,
                dim=config.embedding_sharding_dim)
        else:
            weights['transformer.vocab_embedding.weight'] = vocab_embedding

    if mapping.is_last_pp_rank():
        if hf_model.config.tie_word_embeddings:
            # lm_head.weight has the same weights as embedding
            lm_head_weights = vocab_embedding.clone()
        else:
            lm_head_weights = get_weight(model_params,
                                         param_name_map["lm_head"], dtype)

        if config.vocab_size % mapping.tp_size != 0:
            # padding
            vocab_size_padded = pad_vocab_size(config.vocab_size,
                                               mapping.tp_size)
            pad_width = vocab_size_padded - config.vocab_size

            lm_head_weights = torch.nn.functional.pad(lm_head_weights,
                                                      (0, 0, 0, pad_width),
                                                      'constant',
                                                      value=0)
        weights['lm_head.weight'] = split_matrix_tp(lm_head_weights,
                                                    mapping.tp_size,
                                                    mapping.tp_rank,
                                                    dim=0)
        ln_f_w = get_weight(model_params,
                            f'{model_prefix}.{param_name_map["ln_f"]}', dtype)
        weights['transformer.ln_f.weight'] = ln_f_w

    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: LLaMAConfig,
             device: str = 'cuda',
             calib_dataset: str = 'cnn_dailymail',
             trust_remote_code: bool = True,
             calib_batches: int = 512,
             calib_max_seq_length: int = 512):
    '''
        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"
    assert hf_model_dir is not None
    ## only load and call smooth quant routine once for all ranks
    hf_config = AutoConfig.from_pretrained(hf_model_dir,
                                           trust_remote_code=trust_remote_code)
    assert "llava" not in hf_config.model_type, "Smooth quant llava/vila/llava_next is not supported yet"
    hf_model = AutoModelForCausalLM.from_pretrained(
        hf_model_dir,
        device_map='auto' if device != 'cpu' else 'cpu',
        torch_dtype='auto' if not use_smooth_quant else torch.float16,
        trust_remote_code=trust_remote_code)

    os.environ["TOKENIZERS_PARALLELISM"] = os.environ.get(
        "TOKENIZERS_PARALLELISM", "false")
    tokenizer = AutoTokenizer.from_pretrained(
        hf_model_dir,
        trust_remote_code=trust_remote_code,
        use_fast=False,
        padding_side='left')

    dataset = load_calib_dataset(calib_dataset)

    if calib_batches == -1:  # use the whole dataset if calib_batches is -1
        calib_batches = len(dataset)

    act_range = capture_activation_range(hf_model,
                                         tokenizer,
                                         dataset,
                                         num_samples=calib_batches,
                                         seq_len=calib_max_seq_length)
    qkv_para, smoother = {}, {}
    if use_smooth_quant:
        smooth_llama_model(hf_model, act_range, quant_config.smoothquant_val,
                           qkv_para, smoother)

    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,
            qkv_para=qkv_para,
            smoother=smoother,
        )
        safetensors.torch.save_file(
            weights, os.path.join(output_dir, f'rank{rank}.safetensors'))
        del weights


class QkvWeightHelper:
    """ A helper utility for loading QKV weights from sharded files. """

    def __init__(self, config: PretrainedConfig):
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.tp_size = config.mapping.tp_size
        self.tp_rank = config.mapping.tp_rank
        self.is_mha = self.num_heads == self.num_kv_heads
        self.head_size = None if not hasattr(config,
                                             "head_size") else config.head_size
        self._qkv_weights = {}
        self.remove_duplicated_kv_heads = getattr(config,
                                                  'remove_duplicated_kv_heads',
                                                  False)

    @staticmethod
    def is_qkv_weight(name):
        for k in ['q_proj', 'k_proj', 'v_proj']:
            if 'self_attn' in name and k in name:
                return True
        return False

    def add_weight(self, i: int, name: str, weight: torch.Tensor):
        if 'q_proj' in name:
            tag = 'q'
        elif 'k_proj' in name:
            tag = 'k'
        elif 'v_proj' in name:
            tag = 'v'
        else:
            raise ValueError(f'Got an unexpected parameter of name {name}')
        if i not in self._qkv_weights:
            self._qkv_weights[i] = {}
        self._qkv_weights[i][tag] = weight

    def is_qkv_prepared(self, layer_idx):
        if layer_idx not in self._qkv_weights:
            return False
        weights = self._qkv_weights[layer_idx]
        return 'q' in weights and 'k' in weights and 'v' in weights

    def split_qkv_weights(self, layer_idx):
        if not self.is_qkv_prepared(layer_idx):
            return None
        weights = self._qkv_weights.pop(layer_idx)  # to prevent memory leak.
        q, k, v = (torch.tensor(weights[t]) for t in ['q', 'k', 'v'])

        if self.remove_duplicated_kv_heads:
            head_size = self.hidden_size // self.num_heads if self.head_size is None else self.head_size
            k = k.reshape(
                [k.shape[0] // head_size // 2, 2, head_size, self.hidden_size])
            v = v.reshape(
                [v.shape[0] // head_size // 2, 2, head_size, self.hidden_size])
            assert (k[:, 0] == k[:, 1]).all()
            assert (v[:, 0] == v[:, 1]).all()
            k = k[:, 0].reshape([-1, self.hidden_size])
            v = v[:, 0].reshape([-1, self.hidden_size])

        if not self.is_mha:
            head_size = self.hidden_size // self.num_heads if self.head_size is None else self.head_size
            if self.num_kv_heads < self.tp_size:
                # duplicate the KV heads up to tensor_parallel
                k = dup_kv_weight(k, self.num_kv_heads, self.tp_size)
                v = dup_kv_weight(v, self.num_kv_heads, self.tp_size)
            assert k.shape[0] % (self.tp_size * head_size) == 0
            assert v.shape[0] % (self.tp_size * head_size) == 0
            wq = split(q, self.tp_size, self.tp_rank)
            wk = split(k, self.tp_size, self.tp_rank)
            wv = split(v, self.tp_size, self.tp_rank)
            fused_qkv = torch.cat((wq, wk, wv), dim=0)
        else:
            qkv = torch.cat([q, k, v], dim=0)
            qkv = qkv.reshape(3, q.shape[0], q.shape[1])
            fused_qkv = split(qkv, self.tp_size, self.tp_rank, dim=1)
            fused_qkv = fused_qkv.reshape(3 * (q.shape[0] // self.tp_size),
                                          q.shape[1])
        return fused_qkv


def load_weights_from_hf_by_shard(model_dir: str, config: LLaMAConfig):
    '''Weights-only quantization is the only supported quantization recipe here.'''
    logger.info('Loading weights from HF LLaMA...')
    quant_algo = config.quantization.quant_algo
    use_weight_only = quant_algo in [QuantAlgo.W8A16, QuantAlgo.W4A16]
    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

    weights = {}
    tik = time.time()
    dtype = getattr(torch, config.dtype)

    mapping = config.mapping
    moe_config = config.moe
    assert not moe_config.has_moe(), "MoE does not support sharded load"
    assert "Exaone" not in config.architecture, "EXAONE model currently not support sharded load"

    from transformers import AutoConfig
    hf_config = AutoConfig.from_pretrained(model_dir)

    quant_mode = config.quant_mode
    if quant_mode.is_int8_weight_only():
        plugin_weight_only_quant_type = torch.int8
    elif quant_mode.is_int4_weight_only():
        plugin_weight_only_quant_type = torch.quint4x2
    elif config.quant_mode.has_fp8_rowwise():
        plugin_weight_only_quant_type = torch.float8_e4m3fn
    else:
        plugin_weight_only_quant_type = None
    use_weight_only = quant_mode.is_weight_only()
    use_fp8_rowwise = quant_mode.has_fp8_rowwise()
    # Meta's recipe of not using fp8 rowwise for the first and last layer.
    exclude_layers_id = [0, config.num_hidden_layers - 1]

    layers_range = mapping.pp_layers(config.num_hidden_layers)

    qkv_weight_helper = QkvWeightHelper(config)

    def convert_to_dtype(name, param, model_params, dtype):
        # fp8 rowwise weights will only load fp8 weights and scales for the mlp layer.
        if ('weight_scale' in name or name.replace('weight', 'weight_scale') in model_params) \
           and use_fp8_rowwise:
            assert 'mlp' in name, "only MLP layers support fp8 rowwise currently."
            return param
        else:
            return param.to(dtype)

    def fp8_rowwise_quantization(name,
                                 param,
                                 model_params,
                                 clamp_value,
                                 split_scale=False):
        # check if weights are already quantized.
        loaded_weight_scale = model_params.get(
            name.replace('weight', 'weight_scale'))
        if loaded_weight_scale is not None:
            assert param.dtype == torch.float8_e4m3fn, "weight data type must be torch.float8_e4m3fn"
            if split_scale:
                assert mapping is not None
                loaded_weight_scale = split(loaded_weight_scale,
                                            mapping.tp_size,
                                            mapping.tp_rank,
                                            dim=0)

            return param, loaded_weight_scale.reshape(-1)
        else:
            return fp8_per_channel_quant_weight_gpu(param, clamp_value)

    for model_file in iterate_shard_files(model_dir,
                                          rank=mapping.tp_rank,
                                          progress_bar=False):
        logger.debug(f'Loading file {str(model_file)}...')
        model_params = load_state_dict(model_file)
        for name, param in model_params.items():
            logger.debug(f'Converting weight {name}...')
            layer_idx = retrieved_layer_index_from_name(name)
            tllm_prex = f'transformer.layers.{layer_idx}.'

            param = convert_to_dtype(name, param, model_params, dtype)

            if layer_idx is None:
                layer = None
            else:
                if layer_idx not in layers_range:
                    continue
                else:
                    tllm_prex = f'transformer.layers.{layer_idx - layers_range[0]}.'

            # Meta's recipe of not using fp8 rowwise for the first and last layer.
            use_fp8_rowwise_in_layer = use_fp8_rowwise and (
                layer_idx not in exclude_layers_id)

            if 'model.embed_tokens.weight' in name:
                if hf_config.tie_word_embeddings:
                    # lm_head.weight has the same weights as embedding
                    if mapping.is_last_pp_rank():

                        if config.vocab_size % mapping.tp_size != 0:
                            # padding
                            vocab_size_padded = pad_vocab_size(
                                config.vocab_size, mapping.tp_size)
                            pad_width = vocab_size_padded - config.vocab_size
                            param = torch.from_numpy(
                                np.pad(param.detach().cpu().numpy(),
                                       ((0, pad_width), (0, 0)),
                                       'constant',
                                       constant_values=0))
                        weights['lm_head.weight'] = split(
                            param, mapping.tp_size, mapping.tp_rank)
                if config.use_parallel_embedding:
                    param = split(param, mapping.tp_size, mapping.tp_rank,
                                  config.embedding_sharding_dim)
                if mapping.is_first_pp_rank():
                    weights['transformer.vocab_embedding.weight'] = param
            elif 'model.norm.weight' in name:
                if mapping.is_last_pp_rank():
                    weights['transformer.ln_f.weight'] = param
            elif 'lm_head.weight' in name:
                if mapping.is_last_pp_rank():
                    if config.vocab_size % mapping.tp_size != 0:
                        # padding
                        vocab_size_padded = pad_vocab_size(
                            config.vocab_size, mapping.tp_size)
                        pad_width = vocab_size_padded - config.vocab_size
                        param = torch.from_numpy(
                            np.pad(param.detach().cpu().numpy(),
                                   ((0, pad_width), (0, 0)),
                                   'constant',
                                   constant_values=0))
                    weights['lm_head.weight'] = split(param, mapping.tp_size,
                                                      mapping.tp_rank)
            elif 'input_layernorm.weight' in name:
                weights[tllm_prex + 'input_layernorm.weight'] = param
            elif 'post_attention_layernorm.weight' in name:
                weights[tllm_prex + 'post_layernorm.weight'] = param
            elif qkv_weight_helper.is_qkv_weight(name):
                qkv_weight_helper.add_weight(layer_idx, name, param)
                if not qkv_weight_helper.is_qkv_prepared(layer_idx):
                    continue
                split_v = qkv_weight_helper.split_qkv_weights(layer_idx)
                if use_weight_only:
                    param = split_v.transpose()
                    processed_torch_weights, torch_weight_scales = \
                        torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix(
                            param, plugin_weight_only_quant_type)
                    weights[tllm_prex +
                            'attention.qkv.weight'] = processed_torch_weights
                    weights[
                        tllm_prex +
                        'attention.qkv.per_channel_scale'] = torch_weight_scales
                else:
                    weights[tllm_prex + 'attention.qkv.weight'] = split_v
            elif 'self_attn.o_proj.weight' in name:
                split_v = split(param, mapping.tp_size, mapping.tp_rank, dim=1)
                if use_weight_only:
                    processed_torch_weights, torch_weight_scales = \
                        torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix(
                            split_v.transpose(), plugin_weight_only_quant_type)
                    weights[tllm_prex +
                            'attention.dense.weight'] = processed_torch_weights
                    weights[
                        tllm_prex +
                        'attention.dense.per_channel_scale'] = torch_weight_scales
                else:
                    weights[tllm_prex + 'attention.dense.weight'] = split_v
            elif name.endswith('mlp.up_proj.weight'):
                split_v = split(param, mapping.tp_size, mapping.tp_rank, dim=0)
                if use_weight_only:
                    processed_torch_weights, torch_weight_scales = \
                        torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix(
                            split_v.transpose(), plugin_weight_only_quant_type)
                    weights[tllm_prex +
                            'mlp.gate.weight'] = processed_torch_weights
                    weights[tllm_prex +
                            'mlp.gate.per_channel_scale'] = torch_weight_scales
                elif use_fp8_rowwise_in_layer:
                    processed_torch_weights, torch_weight_scales = fp8_rowwise_quantization(
                        name, split_v, model_params,
                        config.quantization.clamp_val, True)
                    weights[tllm_prex +
                            'mlp.gate.weight'] = processed_torch_weights.view(
                                plugin_weight_only_quant_type)
                    weights[
                        tllm_prex +
                        'mlp.gate.per_channel_scale'] = torch_weight_scales.to(
                            torch.float32)
                else:
                    weights[tllm_prex + 'mlp.gate.weight'] = split_v
            elif name.endswith('mlp.down_proj.weight'):
                split_v = split(param, mapping.tp_size, mapping.tp_rank, dim=1)
                if use_weight_only:
                    processed_torch_weights, torch_weight_scales = \
                        torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix(
                            split_v.transpose(), plugin_weight_only_quant_type)
                    weights[tllm_prex +
                            'mlp.proj.weight'] = processed_torch_weights
                    weights[tllm_prex +
                            'mlp.proj.per_channel_scale'] = torch_weight_scales
                elif use_fp8_rowwise_in_layer:
                    processed_torch_weights, torch_weight_scales = fp8_rowwise_quantization(
                        name, split_v, model_params,
                        config.quantization.clamp_val)
                    weights[tllm_prex +
                            'mlp.proj.weight'] = processed_torch_weights.view(
                                plugin_weight_only_quant_type)
                    weights[
                        tllm_prex +
                        'mlp.proj.per_channel_scale'] = torch_weight_scales.to(
                            torch.float32)
                else:
                    weights[tllm_prex + 'mlp.proj.weight'] = split_v

            elif name.endswith('mlp.gate_proj.weight'):
                split_v = split(param, mapping.tp_size, mapping.tp_rank, dim=0)
                if use_weight_only:
                    processed_torch_weights, torch_weight_scales = \
                        torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix(
                            split_v.transpose(), plugin_weight_only_quant_type)
                    layer.mlp.fc.weight.value = processed_torch_weights
                    layer.mlp.fc.per_channel_scale.value = torch_weight_scales
                    weights[tllm_prex +
                            'mlp.fc.weight'] = processed_torch_weights
                    weights[tllm_prex +
                            'mlp.fc.per_channel_scale'] = torch_weight_scales
                elif use_fp8_rowwise_in_layer:
                    processed_torch_weights, torch_weight_scales = fp8_rowwise_quantization(
                        name, split_v, model_params,
                        config.quantization.clamp_val, True)
                    weights[tllm_prex +
                            'mlp.fc.weight'] = processed_torch_weights.view(
                                plugin_weight_only_quant_type)
                    weights[
                        tllm_prex +
                        'mlp.fc.per_channel_scale'] = torch_weight_scales.to(
                            torch.float32)
                else:
                    weights[tllm_prex + 'mlp.fc.weight'] = split_v

        del model_params

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


def load_weights_from_hf_safetensors(model_dir: str, config: LLaMAConfig):
    logger.info('Loading weights from Huggingface {} safetensors...'.format(
        config.architecture.split('ForCausalLM')[0]))
    tik = time.time()
    import json
    import os

    import safetensors
    weights = {}

    model_dir = model_dir if model_dir.endswith("/") else model_dir + "/"
    safetensors_map = {}
    has_safetensor_index_json = True
    try:
        with open(model_dir + "model.safetensors.index.json", 'r') as fr:
            sharding_map = json.load(fr)
        for k, v in sharding_map['weight_map'].items():
            safetensors_map[k] = int(v[6:11]) - 1
    except FileNotFoundError:
        has_safetensor_index_json = False

    shard_files = []
    for name in os.listdir(model_dir):
        if name.endswith(".safetensors"):
            if has_safetensor_index_json and name not in sharding_map[
                    'weight_map'].values():
                continue
            shard_files.append(name)
    shard_files.sort()
    safetensors_ptrs = [
        safetensors.safe_open(model_dir + shard_file,
                              framework="pt",
                              device="cpu") for shard_file in shard_files
    ]

    mapping = config.mapping
    num_hidden_layers = config.num_hidden_layers
    vocab_size = config.vocab_size
    pad_vocab = vocab_size % mapping.tp_size != 0
    vocab_size_padded = pad_vocab_size(config.vocab_size, mapping.tp_size)
    dtype = config.dtype

    moe_config = config.moe

    kv_tp_size = None
    kv_tp_rank = None
    if config.num_key_value_heads < mapping.tp_size:
        kv_tp_size = config.num_key_value_heads
        kv_tp_rank = mapping.tp_rank * kv_tp_size // mapping.tp_size

    model_prefix, layer_prefix, param_name_map = get_prefix_and_param_name_map(
        config.architecture, use_safetensors=True)

    torch_dtype = str_dtype_to_torch(dtype)

    def load(key,
             tp_dim=-1,
             no_prefix=0,
             is_expert_weights=False,
             tp_size=None,
             tp_rank=None):
        if not no_prefix:
            key = f'{model_prefix}.' + key
        ptr_idx = safetensors_map[key] if key in safetensors_map else 0

        if key not in safetensors_ptrs[ptr_idx].keys():
            return None

        tensor_slice = safetensors_ptrs[ptr_idx].get_slice(key)
        tensor_shape = tensor_slice.get_shape()
        if tp_dim == -1:
            res = tensor_slice[:]
        elif tp_dim >= 0 and tp_dim < len(tensor_shape):
            if is_expert_weights:
                if tp_size is None:
                    tp_size = mapping.moe_tp_size
                if tp_rank is None:
                    tp_rank = mapping.moe_tp_rank
            else:
                if tp_size is None:
                    tp_size = mapping.tp_size
                if tp_rank is None:
                    tp_rank = mapping.tp_rank
            dim_size = tensor_shape[tp_dim]
            if dim_size % tp_size != 0:
                logger.error(
                    f"Current weight {key}'s shape {tensor_shape} is invalid at dimension {tp_dim} for TP size {tp_size}"
                )
            indices = [slice(None)] * len(tensor_shape)
            indices[tp_dim] = slice(dim_size * tp_rank // tp_size,
                                    dim_size * (tp_rank + 1) // tp_size)
            res = tensor_slice[indices]
        else:
            raise ValueError(
                f"Invalid TP dim {tp_dim} for weight {key}'s shape {tensor_shape}"
            )
        return res.to(torch_dtype).contiguous(
        ) if "block_sparse_moe.gate" not in key and "block_sparse_moe.router" not in key else res.to(
            torch.float32)

    def load_and_set(target,
                     key,
                     tp_dim=-1,
                     no_prefix=0,
                     is_expert_weights=False):
        res = load(key, tp_dim, no_prefix, is_expert_weights)
        weights[target] = res
        if "weight" in key:
            bias = load(key.replace("weight", "bias"), -1, no_prefix,
                        is_expert_weights)
            if bias is not None:
                weights[target.replace("weight", "bias")] = bias

    if mapping.is_first_pp_rank():
        weights['transformer.vocab_embedding.weight'] = load(
            param_name_map["vocab_embedding"], config.embedding_sharding_dim
            if config.use_parallel_embedding else -1)  # vocab_embedding

    if mapping.is_last_pp_rank():
        v = load(param_name_map["lm_head"], -1, 1) if pad_vocab else load(
            param_name_map["lm_head"], 0, 1)  # lm_head
        if v is None:
            v = load(param_name_map["vocab_embedding"],
                     -1 if pad_vocab else 0).clone().detach()

        if pad_vocab:
            v = torch.nn.functional.pad(
                v, (0, 0, 0, vocab_size_padded - vocab_size), 'constant', 0)
            v = split(v, mapping.tp_size, mapping.tp_rank)
        weights['lm_head.weight'] = v
        weights['transformer.ln_f.weight'] = load(
            param_name_map["ln_f"])  # ln_f

    layers_range = mapping.pp_layers(num_hidden_layers)
    for l in layers_range:
        layer_idx = l - layers_range[0]
        prefix = f'{layer_prefix}.{l}.'
        tllm_prex = f'transformer.layers.{layer_idx}'

        # Attention
        qkv_list = []
        for comp in ["q", "k", "v"]:
            tp_size = kv_tp_size if comp != "q" else None
            tp_rank = kv_tp_rank if comp != "q" else None
            weight_part = load(prefix + f'{param_name_map["attention.qkv"]}.' +
                               comp + param_name_map["qkv_suffix"],
                               0,
                               tp_size=tp_size,
                               tp_rank=tp_rank)
            qkv_list.append(weight_part)
            bias_part = load(
                (prefix + f'{param_name_map["attention.qkv"]}.' + comp +
                 param_name_map["qkv_suffix"]).replace("weight", "bias"),
                0,
                tp_size=tp_size,
                tp_rank=tp_rank)
            if bias_part is not None:
                qkv_list.append(bias_part)
        if len(qkv_list) == 3:
            # No bias
            weights[f'{tllm_prex}.attention.qkv.weight'] = torch.cat(
                qkv_list, 0)
        else:
            weights[f'{tllm_prex}.attention.qkv.weight'] = torch.cat(
                qkv_list[::2], 0)
            weights[f'{tllm_prex}.attention.qkv.bias'] = torch.cat(
                qkv_list[1::2], 0)
        load_and_set(f'{tllm_prex}.attention.dense.weight',
                     prefix + param_name_map["attention.dense"],
                     1)  # attention.dense

        # MLP
        if not moe_config.has_moe():
            load_and_set(f'{tllm_prex}.mlp.gate.weight',
                         prefix + param_name_map["mlp.gate"], 0)  # mlp.gate
            load_and_set(f'{tllm_prex}.mlp.proj.weight',
                         prefix + param_name_map["mlp.proj"], 1)  # mlp.proj
            load_and_set(f'{tllm_prex}.mlp.fc.weight',
                         prefix + param_name_map["mlp.fc"], 0)  # mlp.fc

        else:
            architecture = config.architecture.lower()
            if "granite" not in architecture:
                weights[f'{tllm_prex}.mlp.router.weight'] = load(
                    prefix + 'block_sparse_moe.gate.weight')
            else:
                weights[f'{tllm_prex}.mlp.router.weight'] = load(
                    prefix + 'block_sparse_moe.router.layer.weight')
            rank_experts = list(range(moe_config.num_experts))
            if mapping.has_moe_ep():
                rank_experts = mapping.ep_experts(moe_config.num_experts)

            if "granite" not in architecture:
                expert_weight_list = []
                for suffix in range(3):
                    tp_dim = -1
                    if mapping.has_moe_tp():
                        tp_dim = 1 if suffix == 1 else 0
                    expert_weight_list.append(
                        torch.stack(
                            list(
                                load(
                                    prefix +
                                    f'block_sparse_moe.experts.{expert}.w{suffix + 1}.weight',
                                    tp_dim=tp_dim,
                                    is_expert_weights=True)
                                for expert in rank_experts)))

                w1 = expert_weight_list[0]
                w2 = expert_weight_list[1]
                w3 = expert_weight_list[2]
            else:
                w2 = load(prefix + f'block_sparse_moe.output_linear.weight',
                          is_expert_weights=True)  #TODO: correct this
                w13 = load(prefix + f'block_sparse_moe.input_linear.weight',
                           is_expert_weights=True)

                half_size = w13.shape[-2] // 2
                w1, w3 = w13.split(half_size, dim=-2)
                w1 = w1[rank_experts[0]:rank_experts[-1] + 1]
                w2 = w2[rank_experts[0]:rank_experts[-1] + 1]
                w3 = w3[rank_experts[0]:rank_experts[-1] + 1]

            weights[f'{tllm_prex}.mlp.fc.weight'] = \
                torch.concat([w3, w1], dim=-2).contiguous()
            weights[f'{tllm_prex}.mlp.proj.weight'] = w2.contiguous()

        load_and_set(f'{tllm_prex}.input_layernorm.weight', prefix +
                     param_name_map["input_layernorm"])  # input_layernorm
        load_and_set(f'{tllm_prex}.post_layernorm.weight', prefix +
                     param_name_map["post_layernorm"])  # post_layernorm

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

    return weights


def load_weights_from_gptq(quant_ckpt_path: str, config: LLaMAConfig):
    logger.info('Loading weights from groupwise GPTQ LLaMA safetensors...')
    weights = {}
    tik = time.time()

    num_hidden_layers = config.num_hidden_layers
    vocab_size = config.vocab_size
    dtype = config.dtype
    mapping = config.mapping
    quant_algo = config.quantization.quant_algo

    gptq_llama = safetensors.safe_open(quant_ckpt_path,
                                       framework="pt",
                                       device=0)
    gptq_prefix = "model."
    gptq_suffix_list = [".qweight", ".qzeros", ".scales"]
    gptq_key_list = [
        "embed_tokens.weight",  # vocab_embedding
        "lm_head.weight",  # lm_head
        "norm.weight",  # ln_f
        "self_attn.",  # attention.qkv
        "_proj",  # qkv suffix
        "self_attn.o_proj",  # attention.dense
        "mlp.up_proj",  # mlp.gate
        "mlp.down_proj",  # mlp.proj
        "mlp.gate_proj",  # mlp.fc
        "input_layernorm.weight",  # input_layernorm
        "post_attention_layernorm.weight",  # post_layernorm
    ]
    split_sym = "."

    packer = torch.ops.trtllm.pack_int8_tensor_to_packed_int4
    preprocessor = torch.ops.trtllm.preprocess_weights_for_mixed_gemm
    torch_dtype = str_dtype_to_torch(dtype)

    def load(key, no_prefix=0):
        if no_prefix:
            return gptq_llama.get_tensor(key)
        else:
            return gptq_llama.get_tensor(gptq_prefix + key)

    def torch_split(v, dim):
        if v.shape[dim] % mapping.tp_size != 0:
            logger.error(
                "Current weight shape is invalid for mapping.tp_size=" +
                str(mapping.tp_size))
            assert False, "Invalid TP size"
        return v.split(v.shape[dim] // mapping.tp_size,
                       dim=dim)[mapping.tp_rank]

    def unpack_int32_into_int8(w_packed):
        # Unpack inputs packed in int32/float32 into uint4 and store them in int8 format
        w_packed_int4x2 = w_packed.contiguous().view(torch.uint8)
        w_unpacked = torch.zeros(w_packed_int4x2.shape[0],
                                 w_packed_int4x2.shape[1] * 2,
                                 dtype=torch.int8)
        w_unpacked[:, ::2] = w_packed_int4x2 % 16
        w_unpacked[:, 1::2] = w_packed_int4x2 // 16
        return w_unpacked.contiguous()

    def process_and_assign_weight(v: List[torch.Tensor],
                                  tllm_prex: str,
                                  quant_algo: QuantAlgo,
                                  tp_dim: int = -1):
        if tp_dim == -1:
            qweight_int32, qzeros_int32, scales_fp16 = [
                item.cpu() for item in v
            ]
        else:
            qweight_int32, qzeros_int32, scales_fp16 = [
                torch_split(item, tp_dim).cpu() for item in v
            ]

        USE_UINT4_INPUT = 1  # Set to true if checkpoint store UINT4 weights
        USE_UINT8_INPUT = 1  # Set to true if checkpoint store UINT8 weights
        USE_GPTQ_FOR_LLAMA = 1  # GPTQ-for-LLaMA added 1 to zeros

        if quant_algo == QuantAlgo.W4A16_GPTQ:
            # unpack inputs packed in int32 into int4 and store them in int8 format
            qweight_unpacked_int8 = unpack_int32_into_int8(
                qweight_int32.T).T.contiguous() - 8
            qweight_interleaved = preprocessor(
                packer(qweight_unpacked_int8), torch.quint4x2,
                torch.float16).view(torch.float16)
            # zeros = zeros * scales
            qzeros_unpacked_int32 = unpack_int32_into_int8(qzeros_int32)
            if not USE_UINT4_INPUT:
                # Correcting UINT4 values back to INT4 order
                mask_negative = qzeros_unpacked_int32[qzeros_unpacked_int32 < 0]
                mask_positive = qzeros_unpacked_int32[qzeros_unpacked_int32 >=
                                                      0]
                qzeros_unpacked_int32 = qzeros_unpacked_int32 + 16 * mask_negative - 16 * mask_positive
            zeros_x_scales_fp16 = (-qzeros_unpacked_int32 + 8 * USE_UINT4_INPUT
                                   - USE_GPTQ_FOR_LLAMA) * scales_fp16
        else:
            # unpack inputs packed in int32 into int8
            qweight_unpacked_int8 = (
                qweight_int32.T.contiguous().view(torch.uint8).T.contiguous() -
                128).to(torch.int8)
            qweight_interleaved = preprocessor(qweight_unpacked_int8,
                                               torch.int8, torch.float16).view(
                                                   torch.float16)
            qzeros_unpacked_int32 = qzeros_int32.view(torch.uint8)
            zeros_x_scales_fp16 = (-qzeros_unpacked_int32 +
                                   128 * USE_UINT8_INPUT -
                                   USE_GPTQ_FOR_LLAMA) * scales_fp16

        zeros_x_scales_fp16 = zeros_x_scales_fp16.half()

        results = {
            f'{tllm_prex}.weight': qweight_interleaved,
            f'{tllm_prex}.weights_scaling_factor': scales_fp16,
            f'{tllm_prex}.zero': zeros_x_scales_fp16,
        }
        return results

    # Load weights from GPTQ checkpoint into TRT-LLM module
    # 1. vocab_embedding
    v = load(gptq_key_list[0])
    if mapping.is_first_pp_rank():
        # tensorrt_llm_llama.vocab_embedding.weight.value = v.to(
        #     torch_dtype).cpu().numpy()
        weights['transformer.vocab_embedding.weight'] = v.to(torch_dtype)
    # 2. lm_head
    v = load(gptq_key_list[1], "no_prefix")
    if mapping.is_last_pp_rank():
        # tensorrt_llm_llama.lm_head.weight.value = torch_split(
        #     v, 0).to(torch_dtype).cpu().numpy()
        if vocab_size % mapping.tp_size != 0:
            # padding
            vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size)
            pad_width = vocab_size_padded - vocab_size
            v = torch.from_numpy(
                np.pad(v.detach().cpu().numpy(), ((0, pad_width), (0, 0)),
                       'constant',
                       constant_values=0))
        weights['lm_head.weight'] = torch_split(v, 0).to(torch_dtype)

    # 3. ln_f
    v = load(gptq_key_list[2])
    if mapping.is_last_pp_rank():
        # tensorrt_llm_llama.ln_f.weight.value = v.to(torch_dtype).cpu().numpy()
        weights['transformer.ln_f.weight'] = v.to(torch_dtype)
    # 4. Weights inside each layer
    layers_range = mapping.pp_layers(num_hidden_layers)
    for l in layers_range:
        layer_idx = l - layers_range[0]
        prefix = "layers" + split_sym + str(layer_idx) + split_sym
        logger.info(f'Process weights in layer: {layer_idx}')
        # layer = tensorrt_llm_llama.layers[layer_idx]
        tllm_prex = f'transformer.layers.{l-layers_range[0]}'
        # 4.1 attention.qkv
        qkv_weight_list = []
        for suf in gptq_suffix_list:
            qkv_list = []
            for comp in ["q", "k", "v"]:
                comp_part = load(prefix + gptq_key_list[3] + comp +
                                 gptq_key_list[4] + suf)
                comp_part = torch_split(comp_part, 1)
                qkv_list.append(comp_part)
            qkv_weight_list.append(torch.cat(qkv_list, dim=1))

        # process_and_assign_weight(layer.attention.qkv, qkv_weight_list)
        weights.update(
            process_and_assign_weight(qkv_weight_list,
                                      f'{tllm_prex}.attention.qkv', quant_algo))
        # 4.2 attention.dense
        v = [load(prefix + gptq_key_list[5] + suf) for suf in gptq_suffix_list]
        # process_and_assign_weight(layer.attention.dense, v, 0)
        weights.update(
            process_and_assign_weight(v,
                                      f'{tllm_prex}.attention.dense',
                                      quant_algo,
                                      tp_dim=0))
        # 4.3 mlp.gate
        v = [load(prefix + gptq_key_list[6] + suf) for suf in gptq_suffix_list]
        # process_and_assign_weight(layer.mlp.gate, v, 1)
        weights.update(
            process_and_assign_weight(v,
                                      f'{tllm_prex}.mlp.gate',
                                      quant_algo,
                                      tp_dim=1))
        # 4.4 mlp.proj
        v = [load(prefix + gptq_key_list[7] + suf) for suf in gptq_suffix_list]
        # process_and_assign_weight(layer.mlp.proj, v, 0)
        weights.update(
            process_and_assign_weight(v,
                                      f'{tllm_prex}.mlp.proj',
                                      quant_algo,
                                      tp_dim=0))
        # 4.5 mlp.fc
        v = [load(prefix + gptq_key_list[8] + suf) for suf in gptq_suffix_list]
        # process_and_assign_weight(layer.mlp.fc, v, 1)
        weights.update(
            process_and_assign_weight(v,
                                      f'{tllm_prex}.mlp.fc',
                                      quant_algo,
                                      tp_dim=1))
        # 4.6 input_layernorm
        v = load(prefix + gptq_key_list[9])
        # layer.input_layernorm.weight.value = v.to(torch_dtype).cpu().numpy()
        weights[f'{tllm_prex}.input_layernorm.weight'] = v.to(torch_dtype)

        # 4.7 post_layernorm
        v = load(prefix + gptq_key_list[10])
        # layer.post_layernorm.weight.value = v.to(torch_dtype).cpu().numpy()
        weights[f'{tllm_prex}.post_layernorm.weight'] = v.to(torch_dtype)

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

    return weights


def load_weights_from_deepcompressor(quant_ckpt_path: str, config: LLaMAConfig):
    logger.info(
        'Loading weights from DeepCompressor torch checkpoint for QServe W4A8 inference...'
    )
    weights = {}
    tik = time.time()

    per_group = config.quant_mode.has_per_group_scaling()
    group_size = 128 if per_group else -1

    num_hidden_layers = config.num_hidden_layers
    vocab_size = config.vocab_size
    dtype = config.dtype
    mapping = config.mapping
    torch_dtype = str_dtype_to_torch(dtype)
    assert torch_dtype == torch.float16, "Currently QServe only supports float16"

    # weight
    fake_quant_weights = torch.load(quant_ckpt_path + '/model.pt',
                                    map_location='cpu',
                                    weights_only=True)
    # scale.0, scale.1, zero
    quant_params = torch.load(quant_ckpt_path + '/scale.pt',
                              map_location='cpu',
                              weights_only=True)

    def load(key):
        if 'zero' in key:
            v = quant_params[key]
            # https://github.com/mit-han-lab/qserve/blob/64ee627dfd747510809998d3592439f05a71ba31/scripts/ckpt_converter/checkpoint_converter.py#L99
            if v.min() < 0:
                v = v + 8
            return v
        if 'scale' in key:
            return quant_params[key]
        return fake_quant_weights[key]

    if per_group:
        deepcompressor_suffix = [
            'weight', 'weight.scale.0', 'weight.scale.1', 'weight.scaled_zero'
        ]
        qserve_suffix = ['weight', 's1_scales', 's2_scales', 's2_szeros']
    else:
        deepcompressor_suffix = [
            'weight', 'weight.scale.0', 'weight.scaled_zero'
        ]
        qserve_suffix = ['weight', 's1_scales', 's1_szeros']

    def tp_split_tensor(v: torch.Tensor, dim):
        if v.shape[dim] % mapping.tp_size != 0:
            logger.error(
                f"Current weight shape is invalid for mapping.tp_size={mapping.tp_size}"
            )
            assert False, "Invalid TP size"
        return v.split(v.shape[dim] // mapping.tp_size,
                       dim=dim)[mapping.tp_rank].contiguous()

    def tp_split_weight_and_params(v: List[torch.Tensor], column_linear: bool):
        if per_group:
            weight, s1_scales, s2_scales, s2_szeros = v
            # weight (out_features, in_features)
            # weight.scale.0 (out_features, 1, 1, 1)
            # weight.scale.1 (out_features, 1, in_features/group_size, 1)
            # weight.scaled_zero (out_features, 1, in_features/group_size, 1)
            if column_linear:
                weight = tp_split_tensor(weight, 0)
                s1_scales = tp_split_tensor(s1_scales, 0)
                s2_scales = tp_split_tensor(s2_scales, 0)
                s2_szeros = tp_split_tensor(s2_szeros, 0)
            else:
                weight = tp_split_tensor(weight, 1)
                s1_scales = s1_scales
                s2_scales = tp_split_tensor(s2_scales, 2)
                s2_szeros = tp_split_tensor(s2_szeros, 2)
            return [weight, s1_scales, s2_scales, s2_szeros]
        else:
            weight, s1_scales, s1_szeros = v
            # weight (out_features, in_features)
            # weight.scale.0 (out_features, 1, 1, 1)
            # weight.zero (out_features, 1, 1, 1)
            if column_linear:
                weight = tp_split_tensor(weight, 0)
                s1_scales = tp_split_tensor(s1_scales, 0)
                s1_szeros = tp_split_tensor(s1_szeros, 0)
            else:
                weight = tp_split_tensor(weight, 1)
                s1_scales = s1_scales
                s1_szeros = s1_szeros
            return [weight, s1_scales, s1_szeros]

    def process_weight_and_params(v: List[torch.Tensor], tllm_prex: str):
        if per_group:
            weight, s1_scales, s2_scales, s2_szeros = v
            qweight = qserve_quantize_weight_per_group(weight, s1_scales,
                                                       s2_scales, s2_szeros,
                                                       group_size)
            qweight, s1_scales, s2_scales, s2_zeros = qserve_pack_reorder_per_group(
                qweight, s1_scales, s2_scales, s2_szeros, group_size)

            return {
                # Note: Linear modules in TRTLLM do not use the name 'qweight'
                f'{tllm_prex}.{qserve_suffix[0]}': qweight,
                f'{tllm_prex}.{qserve_suffix[1]}': s1_scales,
                f'{tllm_prex}.{qserve_suffix[2]}': s2_scales,
                f'{tllm_prex}.{qserve_suffix[3]}': s2_zeros,
            }
        else:
            weight, s1_scales, s1_szeros = v
            qweight = qserve_quantize_weight_per_channel(
                weight, s1_scales, s1_szeros)
            qweight, s1_scales, s1_szeros = qserve_pack_reorder_per_channel(
                qweight, s1_scales, s1_szeros)

            return {
                # Note: Linear modules in TRTLLM use the name 'weight' instead of 'qweight'
                f'{tllm_prex}.{qserve_suffix[0]}': qweight,
                f'{tllm_prex}.{qserve_suffix[1]}': s1_scales,
                f'{tllm_prex}.{qserve_suffix[2]}': s1_szeros,
            }

    # Load weights
    # 1. vocab_embedding
    v = load('model.embed_tokens.weight')
    if mapping.is_first_pp_rank():
        weights['transformer.vocab_embedding.weight'] = v.to(torch_dtype)

    # 2. lm_head
    v = load('lm_head.weight')
    if mapping.is_last_pp_rank():
        if vocab_size % mapping.tp_size != 0:
            # padding
            vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size)
            pad_width = vocab_size_padded - vocab_size
            v = torch.nn.functional.pad(v, (0, 0, 0, pad_width))
        weights['lm_head.weight'] = tp_split_tensor(v, 0).to(torch_dtype)

    # 3. ln_f
    v = load('model.norm.weight')
    if mapping.is_last_pp_rank():
        weights['transformer.ln_f.weight'] = v.to(torch_dtype)

    # 4. Weights inside each layer
    layers_range = mapping.pp_layers(num_hidden_layers)
    for layer_idx in layers_range:
        prefix = f'model.layers.{layer_idx}'
        logger.info(f'Processing weights in layer: {layer_idx}')
        tllm_prex = f'transformer.layers.{layer_idx - layers_range[0]}'

        # 4.1 attention.qkv
        qkv_list = []
        for comp in ["q", "k", "v"]:
            v = [
                load(f'{prefix}.self_attn.{comp}_proj.{suffix}')
                for suffix in deepcompressor_suffix
            ]
            v = tp_split_weight_and_params(v, column_linear=True)
            qkv_list.append(v)
        # Concat qkv
        q, k, v = qkv_list
        qkv = [
            torch.concat((q[i], k[i], v[i]), dim=0)
            for i in range(len(deepcompressor_suffix))
        ]
        weights.update(
            process_weight_and_params(qkv, f'{tllm_prex}.attention.qkv'))

        # 4.2 attention.dense
        v = [
            load(f'{prefix}.self_attn.o_proj.{suffix}')
            for suffix in deepcompressor_suffix
        ]
        v = tp_split_weight_and_params(v, column_linear=False)
        weights.update(
            process_weight_and_params(v, f'{tllm_prex}.attention.dense'))

        # TODO: The naming here is tricky.
        # The implementation of GatedMLP is act(fc(x)) * gate(x).
        # However, the common convention is act(gate_proj(x)) * up_proj(x).

        # 4.3 mlp.gate
        v = [
            load(f'{prefix}.mlp.up_proj.{suffix}')
            for suffix in deepcompressor_suffix
        ]
        v = tp_split_weight_and_params(v, column_linear=True)
        weights.update(process_weight_and_params(v, f'{tllm_prex}.mlp.gate'))

        # 4.4 mlp.fc
        v = [
            load(f'{prefix}.mlp.gate_proj.{suffix}')
            for suffix in deepcompressor_suffix
        ]
        v = tp_split_weight_and_params(v, column_linear=True)
        weights.update(process_weight_and_params(v, f'{tllm_prex}.mlp.fc'))

        # 4.5 mlp.proj
        v = [
            load(f'{prefix}.mlp.down_proj.{suffix}')
            for suffix in deepcompressor_suffix
        ]
        v = tp_split_weight_and_params(v, column_linear=False)
        weights.update(process_weight_and_params(v, f'{tllm_prex}.mlp.proj'))

        # 4.6 input_layernorm
        v = load(f'{prefix}.input_layernorm.weight')
        weights[f'{tllm_prex}.input_layernorm.weight'] = v.to(torch_dtype)

        # 4.7 post_layernorm
        v = load(f'{prefix}.post_attention_layernorm.weight')
        weights[f'{tllm_prex}.post_layernorm.weight'] = v.to(torch_dtype)

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

    # TODO: All the RMSNorm weight, including ln_f, input_layernorm, post_layernorm, are actually all 1s
    # Could implement a simplified module without weight
    return weights


def load_torch_meta_ckpt(meta_ckpt_path: Path):
    '''
        meta_ckpt_path: The format of meta_ckpt_path is like <xxx>/consolidated.xx There are two possible cases:
            1. A file like <xxx>/consolidated.xx.pth, loading it by torch.load directly
            2. A folder like <xxx>/consolidated.xx/, need to load all weights in the folder.
    '''
    file_path = meta_ckpt_path.parent / (meta_ckpt_path.name + ".pth")
    if file_path.exists() and file_path.is_file():
        return torch.load(file_path, map_location="cpu")
    else:
        folder_path = meta_ckpt_path
        assert folder_path.exists() and folder_path.is_dir()

        ckpts = list(Path(folder_path).glob("consolidated-*.pth"))

        all_weights = {}
        for ckpt in ckpts:
            _weight = torch.load(ckpt, map_location="cpu")
            all_weights = all_weights | _weight
            del _weight
        return all_weights


def load_weights_from_meta_ckpt(meta_ckpt_dir: str, config: LLaMAConfig):
    torch_dtype = str_dtype_to_torch(config.dtype)
    mapping = config.mapping
    use_fp8_rowwise = config.quant_mode.has_fp8_rowwise()
    if config.quant_mode.has_any_quant() and not use_fp8_rowwise:
        logger.error(
            "Meta ckpts only support fp8_rowwise quantization currently.")
    weights = {}
    # Meta's recipe of not using fp8 rowwise for the first and last layer.
    exclude_layers_id = [0, config.num_hidden_layers - 1]

    def gather_ckpts(ckpts):
        gathered = {}
        for k in ckpts[0]:
            d = 0
            # TODO not sure should we consider tok here.
            if any([n in k for n in ["wo", "w2"]]):
                d = 1
            if "norm" in k or "rope" in k:  # no TP
                gathered[k] = ckpts[0][k].clone()
            else:
                gathered[k] = torch.cat([pt[k] for pt in ckpts], dim=d).clone()
        return gathered

    def split_ckpt(ckpt, ranks_per_ckpt, ckpt_rank):
        split_ckpt = {}
        for k, v in ckpt.items():
            d = 0
            if any(n in k for n in
                   ["wo", "feed_forward.w2", "tok", "feed_forward.gate"]):
                d = 1
            if "norm" in k or "rope" in k:  # no TP
                split_ckpt[k] = v.clone()
            elif config.num_key_value_heads < mapping.tp_size and any(
                    n in k for n in ["wk", "wv"]):
                assert mapping.tp_size % config.num_key_value_heads == 0
                # special case: we need to duplicate KV head
                tmp = dup_kv_weight(v, config.num_key_value_heads,
                                    mapping.tp_size)
                split_ckpt[k] = torch.split(tmp,
                                            tmp.shape[d] // ranks_per_ckpt,
                                            dim=d)[ckpt_rank].clone()
            else:
                split_ckpt[k] = torch.split(v,
                                            v.shape[d] // ranks_per_ckpt,
                                            dim=d)[ckpt_rank].clone()
        return split_ckpt

    def get_current_weights(num_ckpts):
        if num_ckpts > mapping.tp_size:
            # combine ckpts
            assert (num_ckpts % mapping.tp_size) == 0
            nf = num_ckpts // mapping.tp_size
            fs = nf * mapping.tp_rank
            file_ids = list(range(fs, fs + nf))
            ckpts = []
            for f in file_ids:
                ckpt = load_torch_meta_ckpt(
                    Path(meta_ckpt_dir, f"consolidated.{f:02d}"))
                ckpts.append(ckpt)
            return gather_ckpts(ckpts)
        elif num_ckpts < mapping.tp_size:
            # split ckpt
            assert (mapping.tp_size % num_ckpts) == 0
            ranks_per_ckpt = mapping.tp_size // num_ckpts
            ckpt_fid = mapping.tp_rank // ranks_per_ckpt
            ckpt_rank = mapping.tp_rank % ranks_per_ckpt
            nH_per_ckpt = config.num_attention_heads // num_ckpts
            assert (nH_per_ckpt % ranks_per_ckpt) == 0
            ckpt = load_torch_meta_ckpt(
                Path(meta_ckpt_dir, f"consolidated.{ckpt_fid:02d}"))
            return split_ckpt(ckpt, ranks_per_ckpt, ckpt_rank)

        # num_ckpts == tensor_parallel, 1:1 mapping from files to TP
        return load_torch_meta_ckpt(
            Path(meta_ckpt_dir, f"consolidated.{mapping.tp_rank:02d}"))

    def permute(w, nH, d, dH):
        # due to MQA's wk, nH*dH != d could be true
        return w.view(nH, dH // 2, 2, d).transpose(1, 2).reshape(nH * dH, d)

    def extract_layer_idx(name):
        ss = name.split('.')
        for s in ss:
            if s.isdigit():
                return s
        return None

    if not hasattr(load_weights_from_meta_ckpt, "saved_embed"):
        load_weights_from_meta_ckpt.saved_embed = None

    def combine_embeddings(embeds, num_ckpts):
        if len(embeds) == 1:
            return embeds[0]
        assert [
            embeds[i].shape == embeds[i + 1].shape
            for i in range(len(embeds) - 1)
        ]
        if embeds[0].shape[0] == config.vocab_size // num_ckpts:
            merge_dim = 0
        elif embeds[0].shape[1] == config.hidden_size // num_ckpts:
            merge_dim = 1
        else:
            logger.error("Unable to infer embedding split dimension")
            assert False, "Unable to infer embedding split dimension"
        return torch.cat(embeds, dim=merge_dim)

    def gather_embedding(cur_embed, name: str, num_ckpts):
        if mapping.tp_size == 1:
            # even if num_ckpts > 1, get_current_weights will already have it gathered
            return cur_embed
        if load_weights_from_meta_ckpt.saved_embed is None:
            embeds = [None] * num_ckpts
            for i in range(num_ckpts):
                ckpt = load_torch_meta_ckpt(
                    Path(meta_ckpt_dir, f"consolidated.{i:02d}"))
                embeds[i] = ckpt[name]
            embed = combine_embeddings(embeds, num_ckpts).to(torch_dtype)
            load_weights_from_meta_ckpt.saved_embed = embed

        return load_weights_from_meta_ckpt.saved_embed

    logger.info('Loading weights from Meta LLaMA checkpoints ...')
    tik = time.time()

    num_kv_heads = config.num_key_value_heads
    mha_mode = (num_kv_heads == config.num_attention_heads)

    ckpts = list(Path(meta_ckpt_dir).glob("consolidated.*"))
    num_ckpts = len(ckpts)
    # llama/llama2 doesn't have MQA. So, simplifying loader logic by not worrying about it.
    assert num_kv_heads > 1 or num_kv_heads >= num_ckpts, \
        f"We don't know how the {num_kv_heads} KV heads are distributed among {num_ckpts} checkpoints."

    tik = time.time()
    ckpt = get_current_weights(num_ckpts)
    tok = time.time()
    t = time.strftime('%H:%M:%S', time.gmtime(tok - tik))
    logger.info(f'[{mapping.rank}] get_current_weights. Total time: {t}')

    head_size = config.hidden_size // config.num_attention_heads
    layers_range = mapping.pp_layers(config.num_hidden_layers)

    for l in layers_range:
        prefix = f'layers.{l}.attention.'
        q_weight = permute(ckpt[prefix + 'wq.weight'].clone(),
                           nH=(config.num_attention_heads // mapping.tp_size),
                           d=config.hidden_size,
                           dH=head_size)
        if num_kv_heads < mapping.tp_size and num_ckpts >= mapping.tp_size:
            assert mapping.tp_size % num_kv_heads == 0
            assert False, "Not supported yet"
        k_weight = permute(ckpt[prefix + 'wk.weight'].clone(),
                           nH=((num_kv_heads + mapping.tp_size - 1) //
                               mapping.tp_size),
                           d=config.hidden_size,
                           dH=head_size)
        v_weight = ckpt[prefix + 'wv.weight'].clone()

        qkv_weight = torch.cat([q_weight, k_weight, v_weight], dim=0)
        ckpt[prefix + 'qkv.weight'] = qkv_weight

    for k, v in tqdm(ckpt.items()):
        dtype = torch_dtype if 'feed_forward.gate' not in k else torch.float32

        v = v.to(dtype)
        if "tok_embeddings" in k:
            if not config.use_parallel_embedding:
                v = gather_embedding(v, k, num_ckpts)
            elif config.embedding_sharding_dim == 0:
                # this needs a gather and then resplit along different dims
                v = gather_embedding(v, k, num_ckpts)
                v = split(v, mapping.tp_size, mapping.tp_rank, 0)
            if mapping.is_first_pp_rank():
                weights['transformer.vocab_embedding.weight'] = v
        elif "output" in k:
            if mapping.is_last_pp_rank():
                if config.vocab_size % mapping.tp_size != 0:
                    # padding
                    vocab_size_padded = pad_vocab_size(config.vocab_size,
                                                       mapping.tp_size)
                    pad_width = vocab_size_padded - config.vocab_size
                    v = torch.from_numpy(
                        np.pad(v.detach().cpu().numpy(),
                               ((0, pad_width), (0, 0)),
                               'constant',
                               constant_values=0))
                weights['lm_head.weight'] = v.detach().clone()
        elif k == "norm.weight":
            if mapping.is_last_pp_rank():
                weights['transformer.ln_f.weight'] = v
        else:
            # layer specific weights
            layer_idx = extract_layer_idx(k)
            if layer_idx is None or int(layer_idx) not in layers_range:
                continue

            # Meta's recipe of not using fp8 rowwise for the first and last layer.
            use_fp8_rowwise_in_layer = use_fp8_rowwise and (
                int(layer_idx) not in exclude_layers_id)
            idx = int(layer_idx) - layers_range[0]
            tllm_prex = f'transformer.layers.{idx}.'

            if 'attention_norm.weight' in k:
                weights[tllm_prex + 'input_layernorm.weight'] = v
            elif 'ffn_norm.weight' in k:
                weights[tllm_prex + 'post_layernorm.weight'] = v
            elif 'feed_forward.w3.weight' in k:
                if use_fp8_rowwise_in_layer:
                    processed_torch_weights, torch_weight_scales = fp8_per_channel_quant_weight_gpu(
                        v, config.quantization.clamp_val)
                    weights[tllm_prex +
                            'mlp.gate.weight'] = processed_torch_weights
                    weights[tllm_prex +
                            'mlp.gate.per_channel_scale'] = torch_weight_scales
                else:
                    weights[tllm_prex + 'mlp.gate.weight'] = v
            elif 'feed_forward.w2.weight' in k:
                if use_fp8_rowwise_in_layer:
                    processed_torch_weights, torch_weight_scales = fp8_per_channel_quant_weight_gpu(
                        v, config.quantization.clamp_val)
                    weights[tllm_prex +
                            'mlp.proj.weight'] = processed_torch_weights
                    weights[tllm_prex +
                            'mlp.proj.per_channel_scale'] = torch_weight_scales
                else:
                    weights[tllm_prex + 'mlp.proj.weight'] = v
            elif 'feed_forward.w1.weight' in k:
                if use_fp8_rowwise_in_layer:
                    processed_torch_weights, torch_weight_scales = fp8_per_channel_quant_weight_gpu(
                        v, config.quantization.clamp_val)
                    weights[tllm_prex +
                            'mlp.fc.weight'] = processed_torch_weights
                    weights[tllm_prex +
                            'mlp.fc.per_channel_scale'] = torch_weight_scales
                else:
                    weights[tllm_prex + 'mlp.fc.weight'] = v
            elif 'attention.wo.weight' in k:
                weights[tllm_prex + 'attention.dense.weight'] = v
            elif 'attention.qkv.weight' in k:
                weights[tllm_prex + 'attention.qkv.weight'] = v
            elif 'feed_forward.gate' in k:
                weights[tllm_prex + 'mlp.router.weight'] = v

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