first upload code and model

app.py

@@ -0,0 +1,66 @@
+import torch
+
+import gradio as gr
+import argparse
+from utils import load_hyperparam, load_model
+from models.tokenize import Tokenizer
+from models.llama import *
+from generate import LmGeneration
+
+
+args = None
+lm_generation = None
+
+
+def init_args():
+    global args
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    args = parser.parse_args()
+    args.load_model_path = './model_file/chatllama_7b.bin'
+    args.config_path = './config/llama_7b.json'
+    args.spm_model_path = './model_file/tokenizer.model'
+    args.batch_size = 1
+    args.seq_length = 1024
+    args.world_size = 1
+    args.use_int8 = False
+    args.top_p = 0
+    args.repetition_penalty_range = 1024
+    args.repetition_penalty_slope = 0
+    args.repetition_penalty = 1.15
+
+    args = load_hyperparam(args)
+
+    args.tokenizer = Tokenizer(model_path=args.spm_model_path)
+    args.vocab_size = args.tokenizer.sp_model.vocab_size()
+
+
+def init_model():
+    global lm_generation
+    torch.set_default_tensor_type(torch.HalfTensor)
+    model = LLaMa(args)
+    torch.set_default_tensor_type(torch.FloatTensor)
+    model = load_model(model, args.load_model_path)
+    model.eval()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    lm_generation = LmGeneration(model, args.tokenizer)
+
+
+def chat(prompt, top_k, temperature):
+    args.top_k = int(top_k)
+    args.temperature = temperature
+    response = lm_generation.generate(args, [prompt])
+    return response[0]
+
+
+if __name__ == '__main__':
+    init_args()
+    init_model()
+    demo = gr.Interface(
+        fn=chat,
+        inputs=["text", gr.Slider(1, 60, value=40, step=1), gr.Slider(0.1, 2.0, value=1.2, step=0.1)],
+        outputs="text",
+    )
+    demo.launch()
+

config/llama_7b.json

@@ -0,0 +1,21 @@
+{
+  "emb_size": 4096,
+  "feedforward_size": 11008,
+  "hidden_size": 4096,
+  "hidden_act": "silu",
+  "heads_num": 32,
+  "layers_num": 32,
+  "dropout": 0.1,
+  "data_processor": "lm",
+  "max_seq_length": 2048,
+  "embedding": ["word"],
+  "remove_transformer_bias": true,
+  "remove_embedding_layernorm": true,
+  "rotary_position_embedding": true,
+  "encoder": "transformer",
+  "feed_forward": "gated",
+  "mask": "causal",
+  "layernorm_positioning": "pre",
+  "layernorm": "rms",
+  "target": ["lm"]
+}

generate.py

@@ -0,0 +1,144 @@
+import torch
+import torch.nn.functional as F
+
+
+def apply_temperature(scores, tempt):
+    if tempt > 0:
+        scores = scores / tempt
+    return scores
+
+
+def apply_top_p(scores, top_p, filter_value=-float("Inf"), min_tokens_to_keep=1):
+    if top_p > 0 and top_p < 1:
+        sorted_logits, sorted_indices = torch.sort(scores, descending=False)
+        cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
+
+        # Remove tokens with cumulative top_p above the threshold (token with 0 are kept)
+        sorted_indices_to_remove = cumulative_probs <= (1 - top_p)
+        if min_tokens_to_keep > 1:
+            # Keep at least min_tokens_to_keep
+            sorted_indices_to_remove[..., -min_tokens_to_keep:] = 0
+
+        # scatter sorted tensors to original indexing
+        indices_to_remove = sorted_indices_to_remove.scatter(
+            1, sorted_indices, sorted_indices_to_remove
+        )
+        scores = scores.masked_fill(indices_to_remove, filter_value)
+    return scores
+
+
+def apply_top_k(logits, top_k):
+    top_k = min(top_k, logits.size(-1))  # Safety check
+    if top_k > 0:
+        # Remove all tokens with a probability less than the last token of the top-k
+        indices_to_remove = logits < torch.topk(logits.float(), top_k)[0][..., -1, None]
+        logits[indices_to_remove] = -float("Inf")
+
+    return logits
+
+def apply_advanced_repetition_penalty(
+    input_ids, scores, penalty_range, penalty_slope, penalty
+):
+    penalty_range = int(penalty_range)
+    clipped_penalty_range = min(input_ids.shape[-1], penalty_range)
+
+    if penalty != 1.0:
+        if penalty_range > 0:
+            if clipped_penalty_range < input_ids.shape[1]:
+                input_ids = input_ids[..., -clipped_penalty_range:]
+
+            if penalty_slope != 0:
+                _penalty = (
+                    torch.arange(
+                        penalty_range, dtype=scores.dtype, device=scores.device
+                    )
+                    / (penalty_range - 1)
+                ) * 2.0 - 1
+                _penalty = (penalty_slope * _penalty) / (
+                    1 + torch.abs(_penalty) * (penalty_slope - 1)
+                )
+                _penalty = 1 + ((_penalty + 1) / 2).unsqueeze(0) * (penalty - 1)
+                penalty = _penalty[..., -clipped_penalty_range:]
+
+        score = torch.gather(scores, 1, input_ids)
+        score = torch.where(score <= 0, score * penalty, score / penalty)
+        scores.scatter_(1, input_ids, score)
+
+    return scores
+
+
+class LmGeneration:
+    def __init__(self, model, tokenizer):
+        self.model = model
+        self.tokenizer = tokenizer
+
+    def generate(self, args, prompts, cut_off=None, cut_off_times=1):
+        if cut_off is not None:
+            cut_off_times = [cut_off_times for i in range(len(prompts))]
+        batch = len(prompts)
+        assert batch <= args.batch_size
+
+        prompt_tokens = [args.tokenizer.encode(x, bos=True, eos=False) for x in prompts]
+
+        min_prompt_len = min([len(x) for x in prompt_tokens])
+        # max_prompt_len = max([len(x) for x in prompt_tokens])
+
+        total_len = args.seq_length
+
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        tokens = torch.full((batch, total_len), self.tokenizer.pad_id).to(device).long()
+        for idx, t in enumerate(prompt_tokens):
+            tokens[idx, : len(t)] = torch.tensor(t).long()
+        mask = tokens != self.tokenizer.pad_id
+        start_pos = min_prompt_len
+        prev_pos = 0
+        continue_exsample = [i for i in range(batch)]
+        with torch.no_grad():
+            for cur_pos in range(start_pos, total_len):
+                print(cur_pos)
+                logits = self.model.forward(tokens[continue_exsample, prev_pos:cur_pos], prev_pos, continue_exsample).float()
+                next_token_scores = apply_top_k(logits, top_k=args.top_k)
+                next_token_scores = apply_top_p(next_token_scores, args.top_p)
+                next_token_scores = apply_temperature(next_token_scores, args.temperature)
+                next_token_scores = apply_advanced_repetition_penalty(
+                    tokens[continue_exsample, :cur_pos],
+                    next_token_scores,
+                    args.repetition_penalty_range,
+                    args.repetition_penalty_slope,
+                    args.repetition_penalty
+                )
+                scores = F.softmax(next_token_scores, dim=-1)
+                next_token = torch.multinomial(scores, num_samples=1).squeeze(1)
+                next_token = next_token.reshape(-1)
+                next_token = torch.where(
+                    mask[continue_exsample, cur_pos], tokens[continue_exsample, cur_pos], next_token
+                )
+                tokens[continue_exsample, cur_pos] = next_token
+                prev_pos = cur_pos
+                # remove eos examples.
+                continue_exsample = []
+                for i, t in enumerate(tokens.tolist()):
+                    try:
+                        t.index(self.tokenizer.eos_id)
+                    except ValueError:
+                        if cut_off is not None:
+                            if cut_off == self.tokenizer.decode(t[:cur_pos + 1])[-len(cut_off):]:
+                                if cut_off_times[i] == 1:
+                                    continue
+                                else:
+                                    cut_off_times[i] -= 1
+                        continue_exsample.append(i)
+                if len(continue_exsample) == 0:
+                    break
+
+        decoder = []
+        for i, t in enumerate(tokens.tolist()):
+            t = t[: args.seq_length]
+            try:
+                t = t[: t.index(self.tokenizer.pad_id)]
+                t = t[: t.index(self.tokenizer.eos_id)]
+            except ValueError:
+                pass
+            decoder.append(self.tokenizer.decode(t))
+
+        return decoder

model_file/chatllama_7b.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b5bb1fb1ddf737e7f1fcbe0284ecd384dbe8f243d843b82fcdf59fd00e9b3c61
+size 13476956615

model_file/tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9e556afd44213b6bd1be2b850ebbbd98f5481437a8021afaf58ee7fb1818d347
+size 499723

models/llama.py

@@ -0,0 +1,197 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from models.norm import RMSNorm
+from models.rope import precompute_freqs_cis, apply_rotary_emb
+import bitsandbytes as bnb
+import math
+
+
+class NormalLinear(nn.Linear):
+    def reset_parameters(self) -> None:
+        pass
+
+
+class BnbInt8Linear(bnb.nn.Linear8bitLt):
+    def __init__(self, *args, **kwargs):
+        super().__init__(has_fp16_weights=False, threshold=6.0, *args, **kwargs)
+
+    def reset_parameters(self) -> None:
+        pass
+
+
+def get_linear_layer(use_int8):
+    if use_int8:
+        return BnbInt8Linear
+    return NormalLinear
+
+
+class WordEmbedding(nn.Module):
+    def __init__(self, args):
+        super(WordEmbedding, self).__init__()
+        self.embedding = nn.Embedding(args.vocab_size, args.emb_size)
+
+    def forward(self, src):
+        emb = self.embedding(src)
+        return emb
+
+
+class MultiHeadedAttention(nn.Module):
+    def __init__(self, args, hidden_size, heads_num, attention_head_size, has_bias=True, use_int8=True):
+        super(MultiHeadedAttention, self).__init__()
+        self.heads_num = heads_num
+
+        self.per_head_size = attention_head_size
+        self.inner_hidden_size = heads_num * attention_head_size
+
+        Linear = get_linear_layer(use_int8)
+        self.linear_layers = nn.ModuleList(
+            [Linear(hidden_size, self.inner_hidden_size, bias=has_bias) for _ in range(3)]
+        )
+
+        self.final_linear = Linear(self.inner_hidden_size, hidden_size, bias=has_bias)
+
+        # add cache to reduce compute source.
+        self.cache_k = torch.zeros(
+            (args.batch_size, args.seq_length, self.heads_num, self.per_head_size)
+        )
+        self.cache_v = torch.zeros(
+            (args.batch_size, args.seq_length, self.heads_num, self.per_head_size)
+        )
+
+    def forward(self, key, value, query, start_pos, continue_exsample, mask, freqs_cis):
+        batch_size, seq_length, _ = query.size()
+        heads_num = self.heads_num
+        per_head_size = self.per_head_size
+        query, key, value = [l(x).view(batch_size, -1, heads_num, per_head_size) \
+                             for l, x in zip(self.linear_layers, (query, key, value))]
+        query, key = apply_rotary_emb(query, key, freqs_cis=freqs_cis)
+        if self.cache_k.device != key.device:
+            self.cache_k = self.cache_k.to(key)
+        if self.cache_v.device != value.device:
+            self.cache_v = self.cache_v.to(value)
+
+        self.cache_k[continue_exsample, start_pos: start_pos + seq_length] = key
+        self.cache_v[continue_exsample, start_pos: start_pos + seq_length] = value
+
+        key = self.cache_k[continue_exsample, : start_pos + seq_length]
+        value = self.cache_v[continue_exsample, : start_pos + seq_length]
+
+        query, key, value = [x.transpose(1, 2) for x in (query, key, value)]
+
+        scores = torch.matmul(query, key.transpose(-2, -1))
+        scores = scores / math.sqrt(float(per_head_size))
+        if mask is not None:
+            scores += mask
+        # probs = nn.Softmax(dim=-1)(scores)
+        probs = F.softmax(scores.float(), dim=-1).type_as(query)
+        output = torch.matmul(probs, value).transpose(1, 2).\
+            contiguous().view(batch_size, seq_length, -1)
+        return self.final_linear(output)
+
+
+class GatedFeedForward(nn.Module):
+    def __init__(self, hidden_size, feedforward_size, has_bias=True, use_int8=True):
+        super(GatedFeedForward, self).__init__()
+        Linear = get_linear_layer(use_int8)
+        self.linear_gate = Linear(hidden_size, feedforward_size, bias=has_bias)
+        self.linear_1 = Linear(hidden_size, feedforward_size, bias=has_bias)
+        self.linear_2 = Linear(feedforward_size, hidden_size, bias=has_bias)
+        self.act = F.silu
+
+    def forward(self, x):
+        # gate = self.act(self.linear_gate(x))
+        gate = self.act(self.linear_gate(x)).type_as(x)
+        inter_linear = self.linear_1(x)
+        inter = gate * inter_linear
+        output = self.linear_2(inter)
+        return output
+
+
+class TransformerLayer(nn.Module):
+    def __init__(self, args):
+        super(TransformerLayer, self).__init__()
+
+        if hasattr(args, "attention_head_size"):
+            attention_head_size = args.attention_head_size
+        else:
+            attention_head_size = args.hidden_size // args.heads_num
+
+        has_bias = bool(1 - args.remove_transformer_bias)
+        # Multi-head Attention
+        self.self_attn = MultiHeadedAttention(
+            args, args.hidden_size, args.heads_num, attention_head_size, has_bias=has_bias,
+            use_int8=args.use_int8
+        )
+
+        # FFN
+        self.feed_forward = GatedFeedForward(
+            args.hidden_size, args.feedforward_size, has_bias, use_int8=args.use_int8
+        )
+
+        self.layer_norm_1 = RMSNorm(args.hidden_size)
+        self.layer_norm_2 = RMSNorm(args.hidden_size)
+
+    def forward(self, hidden, start_pos, continue_exsample, mask, freqs_cis=None):
+        inter = self.layer_norm_1(hidden)
+        inter = self.self_attn(inter, inter, inter, start_pos, continue_exsample, mask, freqs_cis)
+        hidden = hidden + inter
+        output = self.layer_norm_2(hidden)
+        output = self.feed_forward(output) + hidden
+        return output
+
+
+class TransformerEncoder(nn.Module):
+    def __init__(self, args):
+        super(TransformerEncoder, self).__init__()
+        self.mask = args.mask
+        self.layers_num = args.layers_num
+
+        self.transformer = nn.ModuleList(
+            [TransformerLayer(args) for _ in range(self.layers_num)]
+        )
+
+        self.layer_norm = RMSNorm(args.hidden_size)
+        self.freqs_cis = precompute_freqs_cis(args.hidden_size // args.heads_num, args.max_seq_length * 2)
+
+    def forward(self, emb, start_pos, continue_exsample):
+        batch_size, seq_length, _ = emb.size()
+        mask = None
+        if seq_length > 1:
+            mask = torch.ones(seq_length, seq_length, device=emb.device)
+            mask = torch.tril(mask)
+            mask = (1.0 - mask) * -10000
+            mask = mask.repeat(batch_size, 1, 1, 1)
+
+        hidden = emb
+        freqs_cis = self.freqs_cis[start_pos: start_pos + seq_length].to(hidden.device)
+
+        for i in range(self.layers_num):
+            hidden = self.transformer[i](hidden, start_pos, continue_exsample, mask, freqs_cis=freqs_cis)
+        return self.layer_norm(hidden)
+
+
+class LmOutput(nn.Module):
+    def __init__(self, args):
+        super(LmOutput, self).__init__()
+        # update: lm output not use int8
+        Linear = get_linear_layer(False)
+        self.lm = Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def forward(self, x):
+        return self.lm(x[:, -1, :])
+
+
+class LLaMa(nn.Module):
+    def __init__(self, args):
+        super(LLaMa, self).__init__()
+        self.embedding = WordEmbedding(args)
+        self.encoder = TransformerEncoder(args)
+        self.target = LmOutput(args)
+
+    #@torch.inference_mode()
+    def forward(self, src, start_pos, continue_exsample):
+        emb = self.embedding(src)
+        output = self.encoder(emb, start_pos, continue_exsample)
+        output = self.target(output)
+        return output

models/norm.py

@@ -0,0 +1,16 @@
+from torch import nn
+import torch
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float()).type_as(x)
+        return output * self.weight

models/rope.py

@@ -0,0 +1,30 @@
+import torch
+from typing import Tuple
+
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+    t = torch.arange(end, device=freqs.device)  # type: ignore
+    freqs = torch.outer(t, freqs).float()  # type: ignore
+    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
+    return freqs_cis
+
+
+def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
+    ndim = x.ndim
+    assert 0 <= 1 < ndim
+    assert freqs_cis.shape == (x.shape[1], x.shape[-1])
+    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+    return freqs_cis.view(*shape)
+
+
+def apply_rotary_emb(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    freqs_cis: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
+    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
+    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
+    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
+    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
+    return xq_out.type_as(xq), xk_out.type_as(xk)

models/tokenize.py

@@ -0,0 +1,40 @@
+# copy from
+# https://github.com/tloen/llama-int8/blob/ce74669c767e42b5082391dd0cfcb621ba40c7f9/llama/tokenizer.py
+
+from sentencepiece import SentencePieceProcessor
+from logging import getLogger
+from typing import List
+import os
+
+
+logger = getLogger()
+
+
+class Tokenizer:
+    def __init__(self, model_path: str):
+        # reload tokenizer
+        assert os.path.isfile(model_path), model_path
+        self.sp_model = SentencePieceProcessor(model_file=model_path)
+        logger.info(f"Reloaded SentencePiece model from {model_path}")
+
+        # BOS / EOS token IDs
+        self.n_words: int = self.sp_model.vocab_size()
+        self.bos_id: int = self.sp_model.bos_id()
+        self.eos_id: int = self.sp_model.eos_id()
+        self.pad_id: int = self.sp_model.pad_id()
+        logger.info(
+            f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}"
+        )
+        assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()
+
+    def encode(self, s: str, bos: bool, eos: bool) -> List[int]:
+        assert type(s) is str
+        t = self.sp_model.encode(s)
+        if bos:
+            t = [self.bos_id] + t
+        if eos:
+            t = t + [self.eos_id]
+        return t
+
+    def decode(self, t: List[int]) -> str:
+        return self.sp_model.decode(t)

requirements.txt

@@ -0,0 +1,4 @@
+torch==1.9.0
+bitsandbytes==0.37.2
+sentencepiece
+argparse

utils.py

@@ -0,0 +1,143 @@
+import json
+import sys
+from argparse import Namespace
+import torch
+import os
+
+
+def load_hyperparam(default_args):
+    """
+    Load arguments form argparse and config file
+    Priority: default options < config file < command line args
+    """
+    with open(default_args.config_path, mode="r", encoding="utf-8") as f:
+        config_args_dict = json.load(f)
+
+    default_args_dict = vars(default_args)
+
+    command_line_args_dict = {k: default_args_dict[k] for k in [
+        a[2:] for a in sys.argv if (a[:2] == "--" and "local_rank" not in a)
+    ]}
+    default_args_dict.update(config_args_dict)
+    default_args_dict.update(command_line_args_dict)
+    args = Namespace(**default_args_dict)
+
+    return args
+
+
+def _load_state_dict_into_model(model_to_load, model_path, start_prefix=""):
+    # Convert old format to new format if needed from a PyTorch state_dict
+
+    # copy state_dict so _load_from_state_dict can modify it
+    state_dict = torch.load(model_path, map_location="cpu")
+    metadata = getattr(state_dict, "_metadata", None)
+    state_dict = state_dict.copy()
+    state_dict['target.lm.weight'] = state_dict['target.lm.output_layer.weight']
+    del state_dict['target.lm.output_layer.weight']
+    state_dict['embedding.embedding.weight'] = state_dict['embedding.word.embedding.weight']
+    del state_dict['embedding.word.embedding.weight']
+
+    if metadata is not None:
+        metadata['embedding.embedding'] = metadata['embedding.word.embedding']
+        metadata['target.lm'] = metadata['target.lm.output_layer']
+        if metadata.get('embedding.dropout', None) is not None:
+            del metadata['embedding.dropout']
+        del metadata['embedding.word']
+        del metadata['embedding.word.embedding']
+        del metadata['target.lm.output_layer']
+        del metadata['target.lm.softmax']
+        del metadata['target.lm.criterion']
+        state_dict._metadata = metadata
+
+    error_msgs = []
+
+    # PyTorch's `_load_from_state_dict` does not copy parameters in a module's descendants
+    # so we need to apply the function recursively.
+    def load(module, state_dict, prefix=""):
+        local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
+        args = (state_dict, prefix, local_metadata, True, [], [], error_msgs)
+        # Parameters of module and children will start with prefix. We can exit early if there are none in this
+        # state_dict
+        if len([key for key in state_dict if key.startswith(prefix)]) > 0:
+            import deepspeed
+            # In sharded models, each shard has only part of the full state_dict, so only gather
+            # parameters that are in the current state_dict.
+            named_parameters = dict(module.named_parameters(prefix=prefix[:-1], recurse=False))
+            params_to_gather = [named_parameters[k] for k in state_dict.keys() if k in named_parameters]
+            if len(params_to_gather) > 0:
+                # because zero3 puts placeholders in model params, this context
+                # manager gathers (unpartitions) the params of the current layer, then loads from
+                # the state dict and then re-partitions them again
+                with deepspeed.zero.GatheredParameters(params_to_gather, modifier_rank=0):
+                    if torch.distributed.get_rank() == 0:
+                        module._load_from_state_dict(*args)
+
+        for name, child in module._modules.items():
+            if child is not None:
+                load(child, state_dict, prefix + name + ".")
+
+    load(model_to_load, state_dict, prefix=start_prefix)
+    # Delete `state_dict` so it could be collected by GC earlier. Note that `state_dict` is a copy of the argument, so
+    # it's safe to delete it.
+    del state_dict
+
+    return model_to_load
+
+
+def convert_normal_parameter_to_int8(model, threshold=6.0, modules_to_not_convert=None, current_key_name=None):
+    import bitsandbytes as bnb
+    modules_to_not_convert = ["lm"] if modules_to_not_convert is None else modules_to_not_convert
+    for name, module in model.named_children():
+        if current_key_name is None:
+            current_key_name = []
+        current_key_name.append(name)
+
+        if len(list(module.children())) > 0:
+            convert_normal_parameter_to_int8(module, threshold, modules_to_not_convert, current_key_name)
+
+        if isinstance(module, bnb.nn.Linear8bitLt) and name not in modules_to_not_convert:
+            # Check if the current key is not in the `modules_to_not_convert`
+            if not any(key in ".".join(current_key_name) for key in modules_to_not_convert):
+                model._modules[name].weight = bnb.nn.Int8Params(
+                    module.weight.data,
+                    requires_grad=False,
+                    has_fp16_weights=False
+                )
+                # Force requires grad to False to avoid unexpected errors
+                model._modules[name].requires_grad_(False)
+        # Remove the last key for recursion
+        current_key_name.pop(-1)
+    return model
+
+
+def load_model(model, model_path):
+    if os.path.isdir(model_path):
+        index_filename = os.path.join(model_path, 'pytorch_model.bin.index.json')
+        with open(index_filename, "r") as f:
+            index = json.loads(f.read())
+        shard_filenames = sorted(set(index["weight_map"].values()))
+        shard_filenames = [os.path.join(model_path, f) for f in shard_filenames]
+        for shard_file in shard_filenames:
+            shard_checkpoint = torch.load(shard_file, map_location='cpu')
+            for name, parameter in model.named_parameters():
+                if shard_checkpoint.get(name, None) is not None:
+                    if 'target' in name:
+                        parameter.data = shard_checkpoint['target.lm.output_layer.weight']
+                    elif 'embedding' in name:
+                        parameter.data = shard_checkpoint['embedding.word.embedding.weight']
+                    else:
+                        parameter.data = shard_checkpoint[name]
+                    parameter.requires_grad = False
+            del shard_checkpoint
+    else:
+        checkpoint = torch.load(model_path, map_location='cpu')
+        for parameter_name, parameter in model.named_parameters():
+            if 'target' in parameter_name:
+                parameter.data = checkpoint['target.lm.output_layer.weight']
+            elif 'embedding' in parameter_name:
+                parameter.data = checkpoint['embedding.word.embedding.weight']
+            else:
+                parameter.data = checkpoint[parameter_name]
+            parameter.requires_grad = False
+        del checkpoint
+    return model