使其可部署在 24GB 顯存的單張 NVIDIA A10 Tensor Core GPU

概述

CodeFuse是由螞蟻集團(tuán)開發(fā)的代碼語言大模型，旨在支持整個(gè)軟件開發(fā)生命周期，涵蓋設(shè)計(jì)、需求、編碼、測試、部署、運(yùn)維等關(guān)鍵階段。

為了在下游任務(wù)上獲得更好的精度，CodeFuse 提出了多任務(wù)微調(diào)框架（MFTCoder），能夠解決數(shù)據(jù)不平衡和不同收斂速度的問題。

通過對(duì)比多個(gè)預(yù)訓(xùn)練基座模型的精度表現(xiàn)，我們發(fā)現(xiàn)利用 MFTCoder [1,2]微調(diào)后的模型顯著優(yōu)于原始基座模型。其中，尤為值得關(guān)注的是采用了 MFTCoder 框架，并利用多任務(wù)數(shù)據(jù)集進(jìn)行微調(diào)的 CodeFuse-CodeLlama-34B [3] 模型，在 HumanEval 評(píng)估數(shù)據(jù)集中取得了當(dāng)時(shí)的最好結(jié)果。具體來說，基于 CodeLlama-34b-Python 模型進(jìn)行微調(diào)的 CodeFuse-CodeLlama-34B 在 HumanEval-python 上實(shí)現(xiàn)了 74.4% 的 pass@1（貪婪解碼）。以下是完整的代碼能力評(píng)估結(jié)果 ：

在代碼補(bǔ)全、text2code、代碼翻譯、單測生成以及代碼生成任務(wù)上，CodeFuse-CodeLlama-34B 全面超過 GPT-3.5；CodeFuse-CodeLlama-34B 能夠在單測生成和代碼補(bǔ)全（HumanEval ）任務(wù)上超過 GPT-4。同時(shí)，上述微調(diào)模型、MFTCoder 訓(xùn)練框架和高質(zhì)量代碼數(shù)據(jù)集已經(jīng)開源。

然而，CodeFuse-CodeLlama-34B 的部署遇到了如下兩個(gè)挑戰(zhàn)：

1）數(shù)據(jù)類型為 fp16 的 34B 模型，顯存占用為 68 GB，至少需要 3 張 A10 才能加載模型，部署成本很高；

2）在模型推理的生成階段，通常伴隨著長條形的矩陣運(yùn)算，此時(shí)計(jì)算量較小，不足以掩蓋 GPU 的訪存延遲，即 memory bound 問題，此時(shí)程序的性能受限于 GPU 帶寬。

為了解決上述問題，我們利用 GPTQ 量化技術(shù)，在降低了部署成本的同時(shí)，也緩解了 GPU 的帶寬壓力 ，從而顯著提升了推理速度。最終，CodeFuse-CodeLlama-34B 的 int4 量化模型可以部署在單張 A10 顯卡上，推理速度可以達(dá)到 20 tokens/s (batch_size=1)。同時(shí)，相較于 fp16 數(shù)據(jù)精度的模型，通過算法上的優(yōu)化，int4 量化引入的精度下降可以控制在 1% 以內(nèi)。下面，我們從模型量化和測試兩個(gè)方面展示我們是如何實(shí)現(xiàn) CodeFuse-CodeLlama-34B 模型的 int4 量化部署的。另外，TensorRT-LLM 也支持了 CodeFuse 中基于 MFTCoder 訓(xùn)練的開源模型部署。

CodeFuse-CodeLlama-34B int4 量化

這里我們使用 GPTQ [4]技術(shù)對(duì)模型進(jìn)行 int4 量化。GPTQ 是對(duì)逐層量化范式經(jīng)典框架 OBQ（Optimal Brain Quantization）[5]的高效實(shí)現(xiàn)，能夠利用單張 A100-80G 在 4 小時(shí)內(nèi)完成 OPT-175B 模型的量化，并且可以獲得較好的準(zhǔn)確率。

另外，我們這里采用了靜態(tài)量化方式，即通過矯正數(shù)據(jù)離線地進(jìn)行量化，得到諸如縮放因子和零點(diǎn)的量化參數(shù)，在推理時(shí)不再進(jìn)行量化參數(shù)的更新。與之對(duì)應(yīng)的是動(dòng)態(tài)量化，會(huì)在模型推理的同時(shí)根據(jù)輸入進(jìn)行量化參數(shù)的調(diào)整。最后，我們這里進(jìn)行的是 int4-weight-only 量化，即只對(duì)權(quán)重進(jìn)行量化而不對(duì)層輸入進(jìn)行量化，即 W4A16 量化。

GPTQ 算法

為了量化權(quán)重，OBQ 框架對(duì)層重建損失函數(shù)進(jìn)行二階泰勒級(jí)數(shù)展開，同時(shí)假設(shè)在未量化的權(quán)重值處一階梯度為零，從而得到如下優(yōu)化問題：

其中，是所有未量化權(quán)重對(duì)應(yīng)的 Hessian 矩陣。那么，量化誤差以及權(quán)重更新值分別為

上面的兩個(gè)公式意味著所有未量化權(quán)重需要通過更新以補(bǔ)償量化帶來的量化誤差。同時(shí)，層重建損失函數(shù)可以按照輸出通道（output channel, OC）分解為獨(dú)立的子問題，例如：

其中Hessian 矩陣為。為了充分利用 GPU 的能力，GPTQ 做了如下三個(gè)改進(jìn)：

所有輸出通道共享相同的量化順序，從而使得行間共享同一份 Hessian 矩陣，大大減少了算法計(jì)算量。

使用一次 Cholesky 分解代替了在 GPTQ 每次迭代中對(duì)整個(gè) Hessian 矩陣的逆矩陣的高斯消元迭代更新方式。既大大減少了計(jì)算量，又得以利用成熟 GPU 矩陣庫中的 Cholesky 算法，且避免了迭代更新方式在矩陣運(yùn)算中所帶來的數(shù)值不穩(wěn)定問題。

通過將整個(gè)計(jì)算過程由對(duì)單個(gè)輸入通道進(jìn)行更新，等效轉(zhuǎn)變?yōu)閯澐?batch 并逐 batch 更新的方式，避免了每次量化對(duì)整個(gè) Hessian 與權(quán)重矩陣的 GPU 讀寫操作，大大降低了 GPU 訪存數(shù)量。

上述的改進(jìn)使得 GPTQ 可以有效提升 GPU 利用率，從而能夠?qū)Υ竽Ｐ瓦M(jìn)行高效量化。

int4-weight-only 量化

這里我們利用開源工具 AutoGPTQ進(jìn)行量化，工具超參數(shù)如下;

利用 AutoGPTQ 進(jìn)行模型加載和推理的例子如下：

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("")
  tokenizer.eos_token_id = tokenizer.convert_tokens_to_ids("")


  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('
') else f'{prompt}
'


  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)

在做靜態(tài)量化時(shí)，GPTQ 使用矯正數(shù)據(jù)集作為輸入計(jì)算 Hessian 矩陣，從而更新未量化權(quán)重進(jìn)而補(bǔ)償量化帶來的誤差。如果推理階段的輸入和矯正數(shù)據(jù)集有偏差（bias），那么量化時(shí)用矯正數(shù)據(jù)得到的 Hessian 矩陣就無法完全反映推理輸入，這會(huì)導(dǎo)致 GPTQ 的誤差補(bǔ)償失效（失效的程度和偏差成正比），出現(xiàn)量化模型在推理輸入上量化誤差變大的情況，進(jìn)而導(dǎo)致量化模型的精度下降。

為了解決上述問題，對(duì)于微調(diào)模型，我們使用了一種數(shù)據(jù)分布對(duì)齊技術(shù)減少模型量化帶來的損失。通過抽取訓(xùn)練數(shù)據(jù)（CodeFuse 開源的高質(zhì)量代碼數(shù)據(jù)集 evol）中的 Question 作為引導(dǎo)方式，利用原始模型生成 Answer，將 Question 和 Answer 拼接起來作為矯正數(shù)據(jù)；最終在 HumanEval Benchmarks 的 Python pass@1 取得了 73.8% 的準(zhǔn)確率，相較于 bf16 模型僅有 0.6% 的精度損失。同時(shí)，在 CMNLI 和 C-Eval 兩個(gè)數(shù)據(jù)集的精度損失也比較少。

構(gòu)建 TensorRT 引擎

在通過 AutoGPTQ 可以得到 safetensors 格式的 int4 量化模型[6]后，我們的目標(biāo)是構(gòu)建單卡 TensorRT 引擎，同時(shí)保證 activation 是 fp16 的數(shù)據(jù)精度。通過 examples/llama/build.py 進(jìn)行 TensorRT 引擎構(gòu)建時(shí)，需要關(guān)注如下參數(shù)：

dtype：設(shè)置為 fp16

use_gpt_attention_plugin：設(shè)置為 fp16，構(gòu)建引擎時(shí)利用 gpt a ttention plugin 并且數(shù)據(jù)精度為 fp16

use_gemm_plugin：設(shè)置為 fp16，構(gòu)建引擎時(shí)利用 gemm_plugin 并且數(shù)據(jù)精度為 fp16

use_weight_only：觸發(fā) weight only 量化

weight_only_precision：設(shè)置為 int4 _gptq，表示構(gòu)建 W4A16 的 GPTQ 量化模型引擎

per_group：gptq 為group-wise 量化，所以需要觸發(fā) per-group

max_batch_size: TensorRT 引擎最大允許 batch size

max_input_len：TensorRT 引擎最大允許輸入長度

max_output_len：TensorRT 引擎最大允許輸出長度

綜上，我們?cè)趩慰?A10/A100上構(gòu)建 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 時(shí)，不同的輸入輸出長度和量化精度情況下，TensorRT-LLM 在 A10/A100 上的推理速度表現(xiàn)?？梢钥吹?，在 A100 上，TensorRT-LLM 的 int4 相對(duì) fp16，最高能夠帶來 2.4 倍的加速，相對(duì) int8 最高也能帶來 1.7 倍的加速。

注意：以上性能測試均基于 TensorRT-LLM 的 0.6.1 版本

顯存占用和結(jié)果測試

我們測量了模型加載后占用的顯存占用情況，以及輸入 2048/1024 tokens 并輸出 1024/2048 tokens 時(shí)的顯存使用情況；同時(shí)我們也測試了量化前后的精度情況，如下表所示：

可見，4bit 量化后，顯存占用大幅縮小，在一張 A10（24GB 顯存）上就能部署 34B 的大模型，具備非常好的實(shí)用性。

模型演示

我們通過終端命令行[7]以及網(wǎng)頁聊天機(jī)器人 [8]兩種不同的方式，展示我們最終的推理效果，具體細(xì)節(jié)可以訪問開源的鏈接。

Cli Demo

WebuiDemo

總結(jié)

在這篇文章中，我們介紹了如何使用 TensorRT-LLM 來加速 CodeFuse 的推理性能。具體而言，我們按照順序展示了如何使用 GPTQ Int4 量化方法、增強(qiáng) GPTQ 量化算法精度的自動(dòng)對(duì)齊技術(shù)、TensorRT-LLM int4 量化模型的使用方法以及相應(yīng)的評(píng)估過程。通過 TensorRT-LLM 的支持，CodeFuse 實(shí)現(xiàn)了較低的推理延遲和優(yōu)化的部署成本。

審核編輯：劉清