运用NVIDIA TensorRT-LLM支撑CodeFuse-CodeLlama-34B上的int4量化和推理优化实践

本文首发于 NVIDIA

一、概述

CodeFuse(github.com/codefuse-ai) 是由蚂蚁集团开发的代码言语大模型,旨在支撑整个软件开发生命周期,涵盖规划、需求、编码、测验、布置、运维等关键阶段。

为了在下游使命上取得更好的精度,CodeFuse 提出了多使命微调结构(MFTCoder),可以处理数据不平衡和不同收敛速度的问题。

经过比照多个预练习基座模型的精度表现,咱们发现运用 MFTCoder[1,2]微调后的模型明显优于原始基座模型。其中,尤为值得关注的是采用了 MFTCoder 结构,并运用多使命数据集进行微调的 CodeFuse-CodeLlama-34B[3]模型,在 HumanEval 评价数据会集取得了当时的最好成果。详细来说,根据 CodeLlama-34b-Python 模型进行微调的 CodeFuse-CodeLlama-34B 在 HumanEval-python 上完结了 74.4% 的 pass@1(贪婪解码)。以下是完整的代码才能评价成果 :

运用NVIDIA TensorRT-LLM支撑CodeFuse-CodeLlama-34B上的int4量化和推理优化实践

在代码补全、text2code、代码翻译、单测生成以及代码生成使命上,CodeFuse-CodeLlama-34B 全面超越 GPT-3.5;CodeFuse-CodeLlama-34B 可以在单测生成和代码补全(HumanEval )使命上超越 GPT-4。一起,上述微调模型、MFTCoder 练习结构和高质量代码数据集已经开源(github:github.com/codefuse-ai…

然而,CodeFuse-CodeLlama-34B 的布置遇到了如下两个挑战:

1)数据类型为 fp16 的 34B 模型,显存占用为 68 GB,至少需求 3 张 A10 才能加载模型,布置本钱很高;

2)在模型推理的生成阶段,通常伴随着长条形的矩阵运算,此刻核算量较小,不足以掩盖 GPU 的访存推迟,即 memory bound 问题,此刻程序的功能受限于 GPU 带宽。

为了处理上述问题,咱们运用 GPTQ 量化技术,在降低了布置本钱的一起,也缓解了 GPU 的带宽压力 ,从而明显提高了推理速度。终究,CodeFuse-CodeLlama-34B 的 int4 量化模型可以布置在单张 A10 显卡上,推理速度可以达到 20 tokens/s (batch_size=1)。一起,相较于 fp16 数据精度的模型,经过算法上的优化,int4 量化引入的精度下降可以控制在 1% 以内。下面,咱们从模型量化和测验两个方面展现咱们是怎么完结 CodeFuse-CodeLlama-34B 模型的 int4 量化布置的。别的,TensorRT-LLM也支撑了 CodeFuse 中根据 MFTCoder 练习的开源模型布置。

二、CodeFuse-CodeLlama-34B int4 量化

这儿咱们运用 GPTQ [4] 技术对模型进行 int4 量化。GPTQ 是对逐层量化范式经典结构 OBQ(Optimal Brain Quantization)[5] 的高效完结,可以运用单张 A100-80G 在 4 小时内完结 OPT-175B 模型的量化,而且可以取得较好的准确率。

别的,咱们这儿采用了静态量化办法,即经过纠正数据离线地进行量化,得到诸如缩放因子和零点的量化参数,在推理时不再进行量化参数的更新。与之对应的是动态量化,会在模型推理的一起根据输入进行量化参数的调整。最后,咱们这儿进行的是 int4-weight-only 量化,即只对权重进行量化而不对层输入进行量化,即 W4A16 量化。

GPTQ 算法

运用NVIDIA TensorRT-LLM支撑CodeFuse-CodeLlama-34B上的int4量化和推理优化实践

  1. 所有输出通道同享相同的量化次序,从而使得行间同享同一份 Hessian 矩阵,大大削减了算法核算量。
  2. 运用一次 Cholesky 分化替代了在 GPTQ 每次迭代中对整个 Hessian 矩阵的逆矩阵的高斯消元迭代更新办法。既大大削减了核算量,又得以运用老练 GPU 矩阵库中的 Cholesky 算法,且避免了迭代更新办法在矩阵运算中所带来的数值不稳定问题。
  3. 经过将整个核算进程由对单个输入通道进行更新,等效转变为区分 batch 并逐 batch 更新的办法,避免了每次量化对整个 Hessian 与权重矩阵的 GPU 读写操作,大大降低了 GPU 访存数量。

上述的改进使得 GPTQ 可以有效提高 GPU 运用率,从而可以对大模型进行高效量化。

三、int4-weight-only 量化

这儿咱们运用开源东西 AutoGPTQ(github.com/PanQiWei/Au…) 进行量化,东西超参数如下;

运用NVIDIA TensorRT-LLM支撑CodeFuse-CodeLlama-34B上的int4量化和推理优化实践

运用 AutoGPTQ 进行模型加载和推理的比如如下:

import os
import torch
import time
from modelscope import AutoTokenizer, snapshot_download
from auto_gptq import AutoGPTQForCausalLM
os.environ["TOKENIZERS_PARALLELISM"] = "false"
def load_model_tokenizer(model_path):
    """
    Load model and tokenizer based on the given model name or local path of downloaded model.
    """
    tokenizer = AutoTokenizer.from_pretrained(model_path, 
                                              trust_remote_code=True, 
                                              use_fast=False,
                                              lagecy=False)
    tokenizer.padding_side = "left"
    tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids("<unk>")
    tokenizer.eos_token_id = tokenizer.convert_tokens_to_ids("</s>")
    model = AutoGPTQForCausalLM.from_quantized(model_path, 
                                                inject_fused_attention=False,
                                                inject_fused_mlp=False,
                                                use_cuda_fp16=True,
                                                disable_exllama=False,
                                                device_map='auto'   # Support multi-gpus
                                              )
    return model, tokenizer
def inference(model, tokenizer, prompt):
    """
    Uset the given model and tokenizer to generate an answer for the speicifed prompt.
    """
    st = time.time()
    inputs = prompt if prompt.endswith('n') else f'{prompt}n'
    input_ids = tokenizer.encode(inputs, 
                                  return_tensors="pt", 
                                  padding=True, 
                                  add_special_tokens=False).to("cuda")
    with torch.no_grad():
        generated_ids = model.generate(
            input_ids=input_ids,
            top_p=0.95,
            temperature=0.1,
            do_sample=True,
            max_new_tokens=512,
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=tokenizer.pad_token_id              
        )
    print(f'generated tokens num is {len(generated_ids[0][input_ids.size(1):])}')
    outputs = tokenizer.batch_decode(generated_ids, skip_special_tokens=True) 
    print(f'generate text is {outputs[0][len(inputs): ]}')
    latency = time.time() - st
    print('latency is {} seconds'.format(latency))
if __name__ == "__main__":
    model_dir = snapshot_download('codefuse-ai/CodeFuse-CodeLlama-34B-4bits', revision='v1.0.0')
    prompt = 'Please write a QuickSort program in Python'
    model, tokenizer = load_model_tokenizer(model_dir)
    inference(model, tokenizer, prompt)

在做静态量化时,GPTQ 运用纠正数据集作为输入核算 Hessian 矩阵,从而更新未量化权重从而补偿量化带来的差错。如果推理阶段的输入和纠正数据集有差错(bias),那么量化时用纠正数据得到的 Hessian 矩阵就无法完全反映推理输入,这会导致 GPTQ 的差错补偿失效(失效的程度和差错成正比),呈现量化模型在推理输入上量化差错变大的状况,从而导致量化模型的精度下降。

为了处理上述问题,关于微调模型,咱们运用了一种数据散布对齐技术削减模型量化带来的损失。经过抽取练习数据(CodeFuse 开源的高质量代码数据集 evol)中的 Question 作为引导办法,运用原始模型生成 Answer,将 Question 和 Answer 拼接起来作为纠正数据;终究在 HumanEval Benchmarks 的 Python pass@1 取得了 73.8% 的准确率,相较于 bf16 模型仅有 0.6% 的精度损失。一起,在 CMNLI 和 C-Eval 两个数据集的精度损失也比较少。

运用NVIDIA TensorRT-LLM支撑CodeFuse-CodeLlama-34B上的int4量化和推理优化实践

四、构建 TensorRT 引擎

在经过 AutoGPTQ 可以得到 safetensors 格局的 int4 量化模型[6]后,咱们的方针是构建单卡 TensorRT 引擎,一起确保 activation 是 fp16 的数据精度。经过 examples/llama/build.py 进行 TensorRT 引擎构建时,需求关注如下参数:

  • dtype:设置为 fp16
  • use_gpt_attention_plugin:设置为 fp16,构建引擎时运用 gpt a ttention plugin 而且数据精度为 fp16
  • use_gemm_plugin:设置为 fp16,构建引擎时运用 gemm_plugin 而且数据精度为 fp16
  • use_weight_only:触发 weight only 量化
  • weight_only_precision:设置为 int4 _gptq,表明构建 W4A16 的 GPTQ 量化模型引擎
  • per_group:gptq 为group-wise 量化,所以需求触发 per-group
  • max_batch_size: TensorRT 引擎最大允许 batch size
  • max_input_len:TensorRT 引擎最大允许输入长度
  • max_output_len:TensorRT 引擎最大允许输出长度

综上,咱们在单卡 A10/A100 上构建 TensorRT 引擎的指令如下:

python build.py --model_dir  "${model_dir}" 
                --quant_safetensors_path "${quant_safetensors_path}" 
                --dtype float16 
                --use_gpt_attention_plugin float16 
                --use_gemm_plugin float16 
                --use_weight_only 
                --weight_only_precision int4_gptq 
                --max_batch_size 1 
                --max_input_len 2048 
                --max_output_len 1024 
                --per_group 
                --output_dir "${engin_dir}" 2>&1  | tee dev_build.log

五、测验

功能

下面,咱们主要测验了 batch size 为 1 时,不同的输入输出长度和量化精度状况下,TensorRT-LLM 在 A10/A100 上的推理速度表现。可以看到,在 A100 上,TensorRT-LLM 的 int4 相对 fp16,最高可以带来 2.4 倍的加快,相对 int8 最高也能带来 1.7 倍的加快。

运用NVIDIA TensorRT-LLM支撑CodeFuse-CodeLlama-34B上的int4量化和推理优化实践

注意:以上功能测验均根据 TensorRT-LLM 的 0.6.1 版别

显存占用和成果测验

咱们测量了模型加载后占用的显存占用状况,以及输入 2048/1024 tokens 并输出 1024/2048 tokens 时的显存运用状况;一起咱们也测验了量化前后的精度状况,如下表所示:

运用NVIDIA TensorRT-LLM支撑CodeFuse-CodeLlama-34B上的int4量化和推理优化实践

可见,4bit 量化后,显存占用大幅缩小,在一张 A10(24GB 显存)上就能布置 34B 的大模型,具备非常好的实用性。

六、模型演示

咱们经过终端指令行[7]以及网页谈天机器人[8]两种不同的办法,展现咱们终究的推理作用,详细细节可以访问开源的链接。

Cli Demo

运用NVIDIA TensorRT-LLM支撑CodeFuse-CodeLlama-34B上的int4量化和推理优化实践

Webui Demo

运用NVIDIA TensorRT-LLM支撑CodeFuse-CodeLlama-34B上的int4量化和推理优化实践

七、总结

在这篇文章中,咱们介绍了怎么运用TensorRT-LLM来加快 CodeFuse 的推理功能。详细而言,咱们依照次序展现了怎么运用 GPTQ Int4 量化办法、增强 GPTQ 量化算法精度的主动对齐技术、TensorRT-LLM int4 量化模型的运用办法以及相应的评价进程。经过 TensorRT-LLM 的支撑,CodeFuse 完结了较低的推理推迟和优化的布置本钱。欢迎我们关注 CodeFuse 获取最新发布的更高准确率的微调大模型。

参考资料:

[1] Liu, B., Chen, C., Liao, C., Gong, Z., Wang, H., Lei, Z., Liang, M., Chen, D., Shen, M., Zhou, H., Yu, H., & Li, J. (2023). MFTCoder: Boosting Code LLMs with Multitask Fine-Tuning. ArXiv, abs/2311.02303.

[2] Zhang, Z., Chen, C., Liu, B., Liao, C., Gong, Z., Yu, H., Li, J., & Wang, R. (2023). Unifying the Perspectives of NLP and Software Engineering: A Survey on Language Models for Code.

[3] huggingface.co/codefuse-ai…

[4] Frantar, E., Ashkboos, S., Hoefler, T., & Alistarh, D. (2022). GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers. ArXiv, abs/2210.17323.

[5] Frantar, E., Singh, S. P., Alistarh, D. (2022). Optimal Brain Compression: A Framework for Accurate Post-Training Quantization and Pruning. Advances in Neural Information Processing Systems, 35, 4475-4488.

[6] huggingface.co/codefuse-ai…

[7] Codefuse-ai: github.com/codefuse-ai

[8] Codefuse-chatbot: github.com/codefuse-ai…