Repository: InternLM/lmdeploy

Issue:
[Bug] int4模型triton server请求异常
### Checklist

- [ ] 1. I have searched related issues but cannot get the expected help.
- [ ] 2. The bug has not been fixed in the latest version.

### Describe the bug

使用模型：https://github.com/ymcui/Chinese-LLaMA-Alpaca这个项目开源的Chinese-Alpaca-Pro-13B
部署triton-server后，使用客户端请求：
返回结果直接是这样子的：
![1692258909165](https://github.com/InternLM/lmdeploy/assets/26675984/dcdd0ad0-cdcb-48f5-b940-16a19f27ac6b)


检查triton server日志发现/workspace/models/model_repository/turbomind/1/weights/tok_embeddings.weight has 403548160, but request 403537920, loading model fails!这个有点异常，不清楚正常吗

后面我换成llama2-7b-chat使用int4是可以正常推理的；
使用Chinese-Alpaca-Pro-13B基于fp16也是可以正常推理的；

麻烦检查下啊

### Reproduction
量化步骤：
python3 -m lmdeploy.lite.apis.calibrate --model ../cuiyiming-chinese-alpaca-13b/ --calib_dataset 'c4' --calib_samples 128 --calib_seqlen 2048 --work_dir ./cym-llama-13b-alpaca-int4
python3 -m lmdeploy.lite.apis.auto_awq --model ./cuiyiming-chinese-alpaca-13b/ --w_bits 4 --w_group_size 128 --work_dir ./cym-llama-13b-alpaca-int4
转换成turbomind格式：
python3 -m lmdeploy.serve.turbomind.deploy \
    --model-name llama \
    --model-path ./cym-llama-13b-alpaca-int4 \
    --model-format awq \
    --group-size 128 \
    -d workspace-cym-llama-13b-alpaca-int4
triton-server部署：
bash workspace-cym-llama-13b-alpaca-int4/service_docker_up.sh
trition server部署日志如下：
Triton Server Version 2.29.0

Copyright (c) 2018-2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.

Various files include modifications (c) NVIDIA CORPORATION & AFFILIATES.  All rights reserved.

This container image and its contents are governed by the NVIDIA Deep Learning Container License.
By pulling and using the container, you accept the terms and conditions of this license:
https://developer.nvidia.com/ngc/nvidia-deep-learning-container-license

NOTE: CUDA Forward Compatibility mode ENABLED.
  Using CUDA 11.8 driver version 520.61.05 with kernel driver version 510.73.05.
  See https://docs.nvidia.com/deploy/cuda-compatibility/ for details.

I0817 07:36:15.292952 1 pinned_memory_manager.cc:240] Pinned memory pool is created at '0x7f07f8000000' with size 268435456
I0817 07:36:15.293596 1 cuda_memory_manager.cc:105] CUDA memory pool is created on device 0 with size 67108864
I0817 07:36:15.317642 1 model_lifecycle.cc:459] loading: preprocessing:1
I0817 07:36:15.317689 1 model_lifecycle.cc:459] loading: postprocessing:1
I0817 07:36:15.317729 1 model_lifecycle.cc:459] loading: turbomind:1
I0817 07:36:15.530570 1 libfastertransformer.cc:1746] TRITONBACKEND_Initialize: turbomind
I0817 07:36:15.530593 1 libfastertransformer.cc:1753] Triton TRITONBACKEND API version: 1.10
I0817 07:36:15.530598 1 libfastertransformer.cc:1757] 'turbomind' TRITONBACKEND API version: 1.10
I0817 07:36:20.097221 1 libfastertransformer.cc:1784] TRITONBACKEND_ModelInitialize: turbomind (version 1)
I0817 07:36:20.097900 1 libfastertransformer.cc:307] Instance group type: KIND_CPU count: 48
num_nodes=1
tp_pp_size=1
gpu_size=1
world_size=1
model_instance_size=1
I0817 07:36:20.097925 1 libfastertransformer.cc:346] Sequence Batching: disabled
I0817 07:36:20.097929 1 libfastertransformer.cc:357] Dynamic Batching: disabled
[TM][INFO] Barrier(1)
Model:
head_num: 40
kv_head_num: 40
size_per_head: 128
inter_size: 13824
num_layer: 40
vocab_size: 39408
attn_bias: 0
max_batch_size: 32
max_context_token_num: 8224
session_len: 2056
step_length: 1
cache_max_entry_count: 48
cache_chunk_size: 1
use_context_fmha: 1
start_id: 1
tensor_para_size: 1
pipeline_para_size: 1
enable_custom_all_reduce: 0
model_name: llama
model_dir: /workspace/models/model_repository/turbomind/1/weights
quant_policy: 0
group_size: 128
I0817 07:36:20.098636 1 libfastertransformer.cc:387] Before Loading Weights:
after allocation    : free: 23.37 GB, total: 23.69 GB, used:  0.32 GB
[TM][WARNING] file /workspace/models/model_repository/turbomind/1/weights/tok_embeddings.weight has 403548160, but request 403537920, loading model fails!
[TM][WARNING] file /workspace/models/model_repository/turbomind/1/weights/output.weight has 403548160, but request 403537920, loading model fails!
I0817 07:36:24.537778 1 libfastertransformer.cc:396] After Loading Weights:

客户端请求：
python3 -m lmdeploy.serve.client 0.0.0.0:33337 1

### Error traceback

_No response_


Discussion:
看起来是码表的大小有问题，能不能提供下具体模型的链接？
使用的镜像是哪个呀？
我又重新试了下，还是针对上面说的模型导出fp16之后部署triton server，镜像使用openmmlab/lmdeploy:v0.0.2，推理结果正常，如下所示：
![1692268586486](https://github.com/InternLM/lmdeploy/assets/26675984/5cbaf704-394f-4b70-8061-745ae58c8eea)
但是使用镜像openmmlab/lmdeploy:latest或者openmmlab/lmdeploy:v0.0.5，推理结果异常了：
![1692268646782](https://github.com/InternLM/lmdeploy/assets/26675984/95d314cb-d863-40a0-848c-2ef5c14d2a7c)

使用镜像openmmlab/lmdeploy:latest或者openmmlab/lmdeploy:v0.0.5部署triton server的日志都会出现这么几句话：
[TM][WARNING] file /workspace/models/model_repository/turbomind/1/weights/tok_embeddings.weight has 403548160, but request 403537920, loading model fails!
[TM][WARNING] file /workspace/models/model_repository/turbomind/1/weights/tok_embeddings.weight has 403548160, but request 403537920, loading model fails!
[TM][WARNING] file /workspace/models/model_repository/turbomind/1/weights/output.weight has 403548160, but request 403537920, loading model fails!
[TM][WARNING] file /workspace/models/model_repository/turbomind/1/weights/output.weight has 403548160, but request 403537920, loading model fails!
I0817 10:29:13.045892 1 libfastertransformer.cc:396] After Loading Weights:
after allocation    : free: 10.66 GB, total: 23.69 GB, used: 13.03 GB
> 看起来是码表的大小有问题，能不能提供下具体模型的链接？

模型有点大，在内网，需要的话我可以传到百度网盘分享下
我参考

> > 看起来是码表的大小有问题，能不能提供下具体模型的链接？
> 
> 模型有点大，在内网，需要的话我可以传到百度网盘分享下

麻烦分享下吧
> 

传到hf hub了，https://huggingface.co/StephenLau/MyLlama
针对这个模型导出fp16之后部署triton server，镜像使用openmmlab/lmdeploy:v0.0.2，推理结果正常；但是使用最新的镜像openmmlab/lmdeploy:latest 推理结果异常
你在用deploy.py的时候，传入的model_name是什么？
lmdeploy现在没有alpaca的对话模板
> 

model-name=llama
对话模板，我自己改了一下，在lmdeploy/model.py中，如下：
![1692330827376](https://github.com/InternLM/lmdeploy/assets/26675984/0ee8d55a-f34e-44ce-8a32-3e3ef85367ea)

@MODELS.register_module(name='llama')
class BaseModel:
    """Base model."""

    def __init__(self):
        self.session_len = 2048
        #self.top_p = 0.9
        self.top_p = 1
        self.top_k = 40
        #self.top_k = -1
        #self.temperature = 0.2
        self.temperature = 0
        self.repetition_penalty = 1.2

    @staticmethod
    def get_prompt(prompt, sequence_start=True):
        """Return the prompt that is concatenated with other elements in the
        chat template.

        Args:
            prompt (str): user's input prompt
            sequence_start (bool): indicator for the first round chat of a
               session sequence
        Returns:
            str: the concatenated prompt
        """
        #return prompt

        prompt_input = (
           "Below is an instruction that describes a task. "
           "Write a response that appropriately completes the request.\n\n"
           "### Instruction:\n{instruction}\n\n### Response: "
        )
        if sequence_start==True:
            return prompt_input.format_map({'instruction': prompt})
        else:
            #prompt = 'Human: {}\nAssistant: '.format(prompt)
            prompt = '\n### Instruction:\n{}\n\n### Response: '.format(prompt)
            #return prompt_input.format_map({'instruction': prompt})
            return prompt
你的模型，词表中的数量是：
vocab_size: 39408

但是token_embeddings的权重长度是403548160。403548160 / 39408 = 10240.25 除不尽。我再看看 lmdeploy:v0.0.2

我用官方的chinese-llama-alpaca-13B，它的词表是49954，在 openmmlab/lmdeploy:latest 中运行是没有问题的。
感觉和 #258 是类似的问题。在 #230  中修改了。
btw，可以给lmdeploy提交个pr吗？把alpaca的对话模板加上
> btw，可以给lmdeploy提交个pr吗？把alpaca的对话模板加上

好的，等我调试完这个alpaca这套模型就提一下
> 你的模型，词表中的数量是： vocab_size: 39408
> 
> 但是token_embeddings的权重长度是403548160。403548160 / 39408 = 10240.25 除不尽。我再看看 lmdeploy:v0.0.2
> 
> 我用官方的chinese-llama-alpaca-13B，它的词表是49954，在 openmmlab/lmdeploy:latest 中运行是没有问题的。

模型在lmdeploy:v0.0.2可以正常triton请求推理，openmmlab/lmdeploy:latest 就不行，这个有复现吗？
这个复现了。在 PR https://github.com/InternLM/lmdeploy/pull/230 中修改了。等这个PR合入后，我再更新下docker image。
着急的话，你可以先用这个PR试试。我验证过了
> 这个复现了。在 PR #230 中修改了。等这个PR合入后，我再更新下docker image。 着急的话，你可以先用这个PR试试。我验证过了

Great！docker更新后在这里通知下哈。
这种问题出现的具体原因是什么啊
语言不太好描述。直接看一下 PR #202，对照 lmdeploy/src/turbomind/utils/memory_utils.cu 前后的改动
docker image更新好了，重新pull一下 openmmlab/lmdeploy:latest
记得提交下 alpaca 的对话模板哈，谢谢先
> docker image更新好了，重新pull一下 openmmlab/lmdeploy:latest

使用上面的llama1模型验证了下，int4这套方案可以正常推理。
后续为了对比fp16与int4的结果，我又使用了llama2-13b进行了验证：
int4可以正常推理。
fp16部署异常，信息如下：
使用的lmdeploy=0.0.5
转换fp16: python3 -m lmdeploy.serve.turbomind.deploy llama2 ../pretrained-models/llama-2-13b-3e/ hf --tp 2
使用openmmlab/lmdeploy:latest部署triton server
部署日志：
[TM][INFO] Barrier(2)
Model:
head_num: 40
kv_head_num: 40
size_per_head: 128
inter_size: 13824
num_layer: 40
vocab_size: 65165
attn_bias: 0
max_batch_size: 32
max_context_token_num: 16416
session_len: 4104
step_length: 1
cache_max_entry_count: 48
cache_chunk_size: 1
use_context_fmha: 1
start_id: 1
tensor_para_size: 2
pipeline_para_size: 1
enable_custom_all_reduce: 0
model_name: llama2
model_dir: /workspace/models/model_repository/turbomind/1/weights
quant_policy: 0
group_size: 0
I0821 07:36:42.354299 1 libfastertransformer.cc:388] Before Loading Weights:
after allocation    : free: 23.37 GB, total: 23.69 GB, used:  0.32 GB
[TM][WARNING] file /workspace/models/model_repository/turbomind/1/weights/tok_embeddings.weight has 667310080, but request 667289600, loading model fails!
[TM][WARNING] file /workspace/models/model_repository/turbomind/1/weights/tok_embeddings.weight has 667310080, but request 667289600, loading model fails!
[TM][WARNING] file /workspace/models/model_repository/turbomind/1/weights/output.weight has 667310080, but request 667289600, loading model fails!
[TM][WARNING] file /workspace/models/model_repository/turbomind/1/weights/output.weight has 667310080, but request 667289600, loading model fails!
I0821 07:36:52.777463 1 libfastertransformer.cc:397] After Loading Weights:
after allocation    : free: 10.17 GB, total: 23.69 GB, used: 13.52 GB
W0821 07:36:52.777599 1 libfastertransformer.cc:517] skipping model configuration auto-complete for 'turbomind': not supported for turbomind backend
I0821 07:36:52.778134 1 python_be.cc:1856] TRITONBACKEND_ModelInstanceInitialize: preprocessing_0_0 (CPU device 0)
I0821 07:36:53.951062 1 python_be.cc:1856] TRITONBACKEND_ModelInstanceInitialize: postprocessing_0_0 (CPU device 0)
I0821 07:36:55.492449 1 libfastertransformer.cc:422] Before Loading Model:
after allocation    : free: 10.06 GB, total: 23.69 GB, used: 13.63 GB
after allocation    : free: 10.06 GB, total: 23.69 GB, used: 13.63 GB
[WARNING] gemm_config.in is not found; using default GEMM algo
[WARNING] gemm_config.in is not found; using default GEMM algo
terminate called recursively
terminate called after throwing an instance of 'std::runtime_error'
[d319215a1c06:00001] *** Process received signal ***
  what():  [d319215a1c06:00001] Signal: Aborted (6)
[d319215a1c06:00001] Signal code:  (-6)
[TM][ERROR]  Assertion fail: /opt/tritonserver/lmdeploy/src/turbomind/models/llama/LlamaV2.cc:93

[d319215a1c06:00001] [ 0] /usr/lib/x86_64-linux-gnu/libpthread.so.0[d319215a1c06:1    :0:636] Caught signal 11 (Segmentation fault: Sent by the kernel at address (nil))
(+0x14420)[0x7f80bdd24420]
[d319215a1c06:00001] [ 1] /usr/lib/x86_64-linux-gnu/libc.so.6(gsignal+0xcb)[0x7f80bc5af00b]
[d319215a1c06:00001] [ 2] /usr/lib/x86_64-linux-gnu/libc.so.6(abort+0x12b)[0x7f80bc58e859]
[d319215a1c06:00001] [ 3] /usr/lib/x86_64-linux-gnu/libstdc++.so.6(+0xac53a)[0x7f80bc97653a]
[d319215a1c06:00001] [ 4] /usr/lib/x86_64-linux-gnu/libstdc++.so.6(+0xaa38c)[0x7f80bc97438c]
[d319215a1c06:00001] [ 5] /usr/lib/x86_64-linux-gnu/libstdc++.so.6(+0xaa3f7)[0x7f80bc9743f7]
[d319215a1c06:00001] [ 6] /usr/lib/x86_64-linux-gnu/libstdc++.so.6(+0xaa6a9)[0x7f80bc9746a9]
[d319215a1c06:00001] [ 7] /opt/tritonserver/backends/turbomind/libtransformer-shared.so(_ZN9turbomind17throwRuntimeErrorEPKciRKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE+0x23d)[0x7f7ff37e553d]
[d319215a1c06:00001] [ 8] /opt/tritonserver/backends/turbomind/libtransformer-shared.so(_ZN9turbomind7LlamaV2I6__halfEC2EmmmmmmRKNS_20LlamaAttentionParamsEfiiiiiiiiibSt10shared_ptrINS2_11SharedStateEEPNS_11LlamaWeightIS1_EENS_9NcclParamEP11CUstream_stPNS_15cublasMMWrapperEPNS_10IAllocatorEbP14cudaDeviceProp+0x8f4)[0x7f7ff381b0b4]
[d319215a1c06:00001] [ 9] /opt/tritonserver/backends/turbomind/libtransformer-shared.so(_ZN16LlamaTritonModelI6__halfE25createSharedModelInstanceEiiSt4pairISt6vectorIN9turbomind9NcclParamESaIS5_EES7_ESt10shared_ptrINS4_18AbstractCustomCommEE+0x5b2)[0x7f7ff388f5e2]
[d319215a1c06:00001] [10] /opt/tritonserver/backends/turbomind/libtransformer-shared.so(_ZN16LlamaTritonModelI6__halfE19createModelInstanceEiiP11CUstream_stSt4pairISt6vectorIN9turbomind9NcclParamESaIS7_EES9_ESt10shared_ptrINS6_18AbstractCustomCommEE+0x441)[0x7f7ff3894a91]
[d319215a1c06:00001] [11] /opt/tritonserver/backends/turbomind/libtriton_turbomind.so(+0x1548a)[0x7f807028348a]
[d319215a1c06:00001] [12] /opt/tritonserver/backends/turbomind/libtriton_turbomind.so(+0x15d43)[0x7f8070283d43]
[d319215a1c06:00001] [13] /opt/tritonserver/backends/turbomind/libtriton_turbomind.so(+0x32a91)[0x7f80702a0a91]
[d319215a1c06:00001] [14] /usr/lib/x86_64-linux-gnu/libstdc++.so.6(+0xd6de4)[0x7f80bc9a0de4]
[d319215a1c06:00001] [15] /usr/lib/x86_64-linux-gnu/libpthread.so.0(+0x8609)[0x7f80bdd18609]
[d319215a1c06:00001] [16] /usr/lib/x86_64-linux-gnu/libc.so.6(clone+0x43)[0x7f80bc68b133]
[d319215a1c06:00001] *** End of error message ***

检查日志发现，针对llama2的fp16部署还存在2处bug:
1.和最初出现的信息一致：/weights/tok_embeddings.weight has 667310080, but request 667289600, loading model fails!
2.Assertion fail: /opt/tritonserver/lmdeploy/src/turbomind/models/llama/LlamaV2.cc:93

麻烦再审核下啊
使用 lmdeploy:latest的吗？转FP16的时候，运行 deploy.py 有没有加上参数 --tp 2呢？
> 使用 lmdeploy:latest的吗？转FP16的时候，运行 deploy.py 有没有加上参数 --tp 2呢？

是的，使用lmdeploy:latest；
转fp16命令：python3 -m lmdeploy.serve.turbomind.deploy llama2 ../pretrained-models/llama-2-13b-3e/ hf --tp 2
我使用官方的llama-2 13B，70B是没有问题。

你的模型，vocab_size是65165。理论上token_embeddings.weight的bytes应该为 65165 * 2 * 5120 = 667289600。
但是实际上读的权重字节数是：667310080，对应的vocab_size为65167

方便的话，能不能共享下你的模型，我们检查下。lmdeploy:latest 已经是从tok_embeddings.weight的shape计算vocab_size了
> 方便的话，能不能共享下你的模型，我们检查下。lmdeploy:latest 已经是从tok_embeddings.weight的shape计算vocab_size了

稍等，我传到hub model
> > 方便的话，能不能共享下你的模型，我们检查下。lmdeploy:latest 已经是从tok_embeddings.weight的shape计算vocab_size了
> 
> 稍等，我传到hub model

模型地址：https://huggingface.co/StephenLau/MyLlama-2-13b
对应的pytorch推理代码：这个地址下的[multi_inference_hf2.py](https://huggingface.co/StephenLau/MyLlama-2-13b/blob/main/multi_inference_hf2.py)
麻烦验证下。感谢！
1. 我用你给的模型，deploy.py后，config中的信息如下。vocab_size: 65167，但你上面给的数据是 65165。请用 git pull 重新拉取下 openmmlab/lmdeploy:latest
2. 你使用了 --tp 2，使用tp的话，lmdeploy 目前要求 vocab_size_ 是 tp 的整数倍。这里我们可以改一下

```
Model: 
head_num: 40
kv_head_num: 40
size_per_head: 128
inter_size: 13824
num_layer: 40
vocab_size: 65167
attn_bias: 0
max_batch_size: 32
max_context_token_num: 16416
session_len: 4104
step_length: 1
cache_max_entry_count: 48
cache_chunk_size: 1
use_context_fmha: 1
start_id: 1
tensor_para_size: 2
pipeline_para_size: 1
enable_custom_all_reduce: 0
model_name: llama2
model_dir: /workspace/lmdeploy/workspace/llama2-13b-chat-266/model_repository/turbomind/1/weights
quant_policy: 0
group_size: 0
```


Resolving patch:
diff --git a/lmdeploy/serve/turbomind/deploy.py b/lmdeploy/serve/turbomind/deploy.py
--- a/lmdeploy/serve/turbomind/deploy.py
+++ b/lmdeploy/serve/turbomind/deploy.py
@@ -147,11 +147,22 @@ def save_bin(param: torch.Tensor, name):
     attn_bias = False
     inter_size = 0
 
+    tok_embeddings = model_params['tok_embeddings.weight']
+    _vocab_size, dim = tok_embeddings.shape
+    head_num = dim // size_per_head
+    if _vocab_size % tp != 0:
+        # Resolve https://github.com/InternLM/lmdeploy/issues/266
+        # Pad tok_embeddings and output weights, making their shape divisible by TP # noqa: E501
+        pad_size = (_vocab_size + tp - 1) // tp * tp - _vocab_size
+        # Pad weight at the bottom of dim 0
+        model_params['tok_embeddings.weight'] = torch.nn.functional.pad(
+            tok_embeddings, (0, 0, 0, pad_size), 'constant', 0)
+        # Pad output weight at the bottom of dim 0
+        model_params['output.weight'] = torch.nn.functional.pad(
+            model_params['output.weight'], (0, 0, 0, pad_size), 'constant', 0)
+
     # reverse the splitting axes since the weights are transposed above
     for param_name, param_data in model_params.items():
-        if param_name == 'tok_embeddings.weight':
-            _vocab_size, dim = param_data.shape
-            head_num = dim // size_per_head
         split_dim = None
         key, ext = param_name.split('.')[-2:]
         if key == 'w_qkv' and ext == 'bias':


Test patch: