本文首发于 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（贪婪解码）。以下是完整的代码才能评价成果 ：
在代码补全、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 算法
- 所有输出通道同享相同的量化次序，从而使得行间同享同一份 Hessian 矩阵，大大削减了算法核算量。
- 运用一次 Cholesky 分化替代了在 GPTQ 每次迭代中对整个 Hessian 矩阵的逆矩阵的高斯消元迭代更新办法。既大大削减了核算量，又得以运用老练 GPU 矩阵库中的 Cholesky 算法，且避免了迭代更新办法在矩阵运算中所带来的数值不稳定问题。
- 经过将整个核算进程由对单个输入通道进行更新，等效转变为区分 batch 并逐 batch 更新的办法，避免了每次量化对整个 Hessian 与权重矩阵的 GPU 读写操作，大大降低了 GPU 访存数量。
上述的改进使得 GPTQ 可以有效提高 GPU 运用率，从而可以对大模型进行高效量化。
三、int4-weight-only 量化
这儿咱们运用开源东西 AutoGPTQ（github.com/PanQiWei/Au…) 进行量化，东西超参数如下;
运用 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 两个数据集的精度损失也比较少。
四、构建 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 倍的加快。
注意：以上功能测验均根据 TensorRT-LLM 的 0.6.1 版别
显存占用和成果测验
咱们测量了模型加载后占用的显存占用状况，以及输入 2048/1024 tokens 并输出 1024/2048 tokens 时的显存运用状况；一起咱们也测验了量化前后的精度状况，如下表所示：
可见，4bit 量化后，显存占用大幅缩小，在一张 A10（24GB 显存）上就能布置 34B 的大模型，具备非常好的实用性。
六、模型演示
咱们经过终端指令行[7]以及网页谈天机器人[8]两种不同的办法，展现咱们终究的推理作用，详细细节可以访问开源的链接。
Cli Demo
Webui Demo
七、总结
在这篇文章中，咱们介绍了怎么运用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…
