云原生 AI 工程化實踐之 FasterTransformer 加速 LLM 推理

作者：顏廷帥（瀚廷）

01 背景

OpenAI 在 3 月 15 日發布了備受矚目的 GPT4，它在司法考試和程序編程領域的驚人表現讓大家對大語言模型的熱情達到了頂點。人們紛紛議論我們是否已經跨入通用人工智能的時代。與此同時，基于大語言模型的應用也如雨后春筍般出現，為我們帶來了協同辦公、客服對話、語言翻譯、內容生成等方面前所未有的暢快體驗。

然而，當我們享受著大語言模型帶來的普惠 AI 能力時，它也給開發者們帶來了前所未有的挑戰。隨著模型不斷增大，計算量也達到了空前的高度，直接導致推理時間變長。為了解決大語言模型推理的延遲問題，業界已經提供了一些解決方案，比如 Tensorrt、FasterTransformer 和 vllm。為了幫助用戶解決云原生系統中的大語言模型推理加速問題，云原生 AI 套件引入了 FasterTransformer 推理加速方案。

本文將在 ACK 容器服務上，以 Bloom7B1 模型為例展示如何使用 FasterTransformer 進行推理加速。本例中會使用以下組件：

Arena

Arena 是基于 Kubernetes 的機器學習輕量級解決方案，支持數據準備、模型開發，模型訓練、模型預測的完整生命周期，提升數據科學家工作效率。同時和阿里云的基礎云服務深度集成，支持 GPU 共享、CPFS 等服務，可以運行阿里云優化的深度學習框架，最大化使用阿里云異構設備的性能和成本的效益。更多 arena 信息，可以參考云原生 AI 套件開發者使用指南 [ 1] 。

Triton Server

Triton Server為Nvidia 提供了機器學習推理引擎，可以支持 Tensorflow、Pytorch、Tensorrt 和 Fastertransformer 多種 backend。云原生 AI 套件已經將 Triton Server 加入到 Arena 中，用戶可以通過簡單的命令行或 SDK 來在云原生系統中完成 Triton Server 服務的拉起、運維和監控。更多 AI 套件中使用 Triton Server 信息，可以參考部署 PyTorch 模型推理服務 [ 2] 。

FasterTransformer

FasterTransformer 是真對于 Transofrmer 類型模型（也包括 encoder-only、decoder-only）的推理加速方案，其提供了 Kernel Fuse、Memory reuse、kv cache、量化等多種優化方案，同時也提供了 Tensor Parallel 和 Pipeline Parallel 兩種分布式推理方案。本文將介紹如何在云原生 AI 套件中使用 FasterTransformer 進行模型的推理加速。

02 環境準備

環境準備分為兩個部分，第一個部分是創建包含 GPU 的 Kubernetes 集群 [ 3] 和安裝云原生 AI 套件 [ 4] ，第二個部分是從 huggingface 官網下載 bloom-7b1 模型。

模型的下載命令如下：

git lfs install
git clone git@hf.co:bigscience/bloom-7b1

通過上面的命令，可以將 huggingface repo 中的文件下載到本地：

在這里插入圖片描述

下載完成后，我們將 bloom-71 文件夾上傳到 OSS 中，作為推理時的共享存儲，OSS 的使用可以參考開始使用 OSS [ 5] 。

上傳到 OSS 之后，分別創建名稱為 bloom7b1-pv 和 bloom7b1-pvc 的 PV 和 PVC，以用于推理服務的容器掛載。具體操作，請參見使用 OSS 靜態存儲卷 [ 6] 。

03 模型轉換

FasterTransformer 本質上是對模型的重寫，它通過 CUDA、cuDNN 和 cuBLAS 重寫了 Transformer 模型結構，因此其具有自己的模型結構和模型參數的描述方式。而我們的模型一般是通過 Pytorch、Tesorflow、Megatron 或 huggingface 這樣的訓練框架產出，其往往又具有自己單獨的一套模型結構和參數的表達，因此在使用FasterTransformer時，就需要將模型原有的 checkpoint 轉換為 FasterTransformer 的結構。

FasterTransformer 中已經支持了多種類型的轉換腳本，這里我們使用 FasterTransofrmer 提供的 examples/pytorch/gpt/utils/huggingface_bloom_convert.py。

云原生 AI 套件已經接入了上述的轉換邏輯，因此，通過如下腳本即可完成一次模型的轉換。

arena submit pytorchjob\--gpus=1\--image ai-studio-registry.cn-beijing.cr.aliyuncs.com/kube-ai/fastertransformer:torch-0.0.1\--name convert-bloom\--workers 1\--namespace default-group\--data bloom-pvc:/mnt\'python  /FasterTransformer/examples/pytorch/gpt/utils/huggingface_bloom_convert.py -i /mnt/model/bloom-7b1 -o /mnt/model/bloom-7b1-ft-fp16 -tp 2  -dt fp16 -p 64 -v'

通過 arena log 來觀察轉換的日志：

$arena logs -n default-group convert-bloom
======================= Arguments =======================- input_dir...........: /mnt/model/bloom-7b1- output_dir..........: /mnt/model/bloom-7b1-ft-fp16- tensor_para_size....: 2- data_type...........: fp16- processes...........: 64- verbose.............: True- by_shard............: False
=========================================================
loading from pytorch bin format
model file num: 2- model.pre_decoder_layernorm.bias................: shape (4096,)            | saved at /mnt/model/bloom-7b1-ft-fp16/2-gpu/model.pre_decoder_layernorm.bias.bin- model.layers.0.input_layernorm.weight...........: shape (4096,)            | saved at /mnt/model/bloom-7b1-ft-fp16/2-gpu/model.layers.0.input_layernorm.weight.bin- model.layers.0.attention.dense.bias.............: shape (4096,)            | saved at /mnt/model/bloom-7b1-ft-fp16/2-gpu/model.layers.0.attention.dense.bias.bin- model.layers.0.input_layernorm.bias.............: shape (4096,)            | saved at /mnt/model/bloom-7b1-ft-fp16/2-gpu/model.layers.0.input_layernorm.bias.bin- model.layers.0.attention.query_key_value.bias...: shape (3, 2048)        s | saved at /mnt/model/bloom-7b1-ft-fp16/2-gpu/model.layers.0.attention.query_key_value.bias.0.bin (0/2)- model.layers.0.post_attention_layernorm.weight..: shape (4096,)            | saved at /mnt/model/bloom-7b1-ft-fp16/2-gpu/model.layers.0.post_attention_layernorm.weight.bin- model.layers.0.post_attention_layernorm.bias....: shape (4096,)            | saved at /mnt/model/bloom-7b1-ft-fp16/2-gpu/model.layers.0.post_attention_layernorm.bias.bin- model.layers.0.mlp.dense_4h_to_h.bias...........: shape (4096,)            | saved at /mnt/model/bloom-7b1-ft-fp16/2-gpu/model.layers.0.mlp.dense_4h_to_h.bias.bin- model.layers.0.mlp.dense_h_to_4h.bias...........: shape (8192,)          s | saved at /mnt/model/bloom-7b1-ft-fp16/2-gpu/model.layers.0.mlp.dense_h_to_4h.bias.0.bin (0/2)- model.layers.0.attention.query_key_value.bias...: shape (3, 2048)        s | saved at /mnt/model/bloom-7b1-ft-fp16/2-gpu/model.layers.0.attention.query_key_value.bias.1.bin (1/2)

通過 arena list 命令查看轉換是否執行結束：

NAME           STATUS     TRAINER     DURATION  GPU(Requested)  GPU(Allocated)  NODE
convert-bloom  SUCCEEDED  PYTORCHJOB  3m        1               N/A             192.168.123.35

轉換完成后，會在 OSS 上創建一個 model/arena/bloom-7b1-ft-fp16 文件夾，文件中會存儲 FasterTransofrmer 所對應的 checkpoint。

04 性能對比

此時，我們的 OSS 上已經有兩份 bloom-7b1 checkpoint，一份是 bloom-7b 文件夾存儲了 huggingface 原生的 checkpoint，另一份是 bloom-7b-ft-fp16 文件夾存儲了轉換后的 FasterTransformer 的 checkpoint。我們將使用這兩份 checkpoint 進行性能對比，看一下來 FasterTransformer 是否能夠帶來性能的提升。

性能對比使用 Fastertransformer 提供的 examples/pytorch/gpt/bloom_lambada.py，我們也已經集成到了 AI 套件中。這里我們分別提交兩個性能評測命令。對 Huggingface Bloom-7b1 評測的命令：

arena submit pytorchjob\--gpus=2\--image ai-studio-registry.cn-beijing.cr.aliyuncs.com/kube-ai/fastertransformer:torch-0.0.1\--name perf-hf-bloom \--workers 1\--namespace default-group\--data bloom7b1-pvc:/mnt\'python /FasterTransformer/examples/pytorch/gpt/bloom_lambada.py \--tokenizer-path /mnt/model/bloom-7b1 \--dataset-path /mnt/data/lambada/lambada_test.jsonl \--batch-size 16 \--test-hf \--show-progress'

查看 HuggingFace 的結果：

$arena -n default-group logs -t 5 perf-hf-bloom
Accuracy: 57.5587% (2966/5153) (elapsed time: 173.2149 sec)

對 Fastertransformer Blooom-7b 評測的命令：

arena submit pytorchjob\--gpus=2\--image ai-studio-registry.cn-beijing.cr.aliyuncs.com/kube-ai/fastertransformer:torch-0.0.1\--name perf-ft-bloom \--workers 1\--namespace default-group\--data bloom7b1-pvc:/mnt\'mpirun --allow-run-as-root -n 2 python /FasterTransformer/examples/pytorch/gpt/bloom_lambada.py \--lib-path /FasterTransformer/build/lib/libth_transformer.so \--checkpoint-path /mnt/model/2-gpu \--batch-size 16 \--tokenizer-path /mnt/model/bloom-7b1 \--dataset-path /mnt/data/lambada/lambada_test.jsonl \--show-progress'

查看 FasterTransformer 的結果，可以看見帶來了 2.5 倍的性能提升。

$arena -n default-group logs -t 5 perf-ft-bloom
Accuracy: 57.6363% (2970/5153) (elapsed time: 68.7818 sec)

通過結果對比可以看見，Fastertransformer 與原生的 Huggingface 相比有比較明顯的性能提升。

05 模型部署

在這一小節，我們使用 Triton Server 對 FasterTransformer 進行部署，Triton Server 中原生并不支持 FasterTransformer 的 backend，需要我們配合 Nvidia 提供的 Fastertransformer backend 來使用。通過使用 FasterTransformer backend，Triton Server 不再進行 GPU 資源的分配，FasterTransformer backend 會根據 CUDA_VISIBLE_DEVICES 判斷當前可用 GPU 資源，并分配給對應的 RANK 來執行分布式的推理。

FasterTransformer 對應的模型 Repo 目錄如下所示：

├── model_repo
│   └── fastertransformer
│       ├── 1
│       │   └── config.ini
│       └── config.pbtxt

使用功能 Arena 的如下命令來啟動 FasterTransformer：

arena serve triton \--namespace=default-group \--version=1 \--data=bloom7b1-pvc:/mnt \--name=ft-triton-bloom \--allow-metrics \--gpus=2 \--replicas=1 \--image=ai-studio-registry.cn-beijing.cr.aliyuncs.com/kube-ai/triton_with_ft:22.03-main-2edb257e-transformers \--model-repository=/mnt/triton_repo

通過 kubectl logs，我們可以看到 triton server 的部署日志，通過日志可以看到，triton server 啟動了兩個 gpu 來進行分布式推理。

I0721 08:57:28.116291 1 pinned_memory_manager.cc:240] Pinned memory pool is created at '0x7fd264000000' with size 268435456
I0721 08:57:28.118393 1 cuda_memory_manager.cc:105] CUDA memory pool is created on device 0 with size 67108864
I0721 08:57:28.118403 1 cuda_memory_manager.cc:105] CUDA memory pool is created on device 1 with size 67108864
I0721 08:57:28.443529 1 model_lifecycle.cc:459] loading: fastertransformer:1
I0721 08:57:28.625253 1 libfastertransformer.cc:1828] TRITONBACKEND_Initialize: fastertransformer
I0721 08:57:28.625307 1 libfastertransformer.cc:1838] Triton TRITONBACKEND API version: 1.10
I0721 08:57:28.625315 1 libfastertransformer.cc:1844] 'fastertransformer' TRITONBACKEND API version: 1.10
I0721 08:57:28.627137 1 libfastertransformer.cc:1876] TRITONBACKEND_ModelInitialize: fastertransformer (version 1)
I0721 08:57:28.628304 1 libfastertransformer.cc:372] Instance group type: KIND_CPU count: 1
I0721 08:57:28.628326 1 libfastertransformer.cc:402] Sequence Batching: disabled
I0721 08:57:28.628334 1 libfastertransformer.cc:412] Dynamic Batching: disabled
I0721 08:57:28.661657 1 libfastertransformer.cc:438] Before Loading Weights:
+-------------------+-----------------------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------+
| Backend           | Path                                                                        | Config                                                                                                                                                        |
+-------------------+-----------------------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------+
| fastertransformer | /opt/tritonserver/backends/fastertransformer/libtriton_fastertransformer.so | {"cmdline":{"auto-complete-config":"true","min-compute-capability":"6.000000","backend-directory":"/opt/tritonserver/backends","default-max-batch-size":"4"}} |
+-------------------+-----------------------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------+I0721 09:01:19.653743 1 server.cc:633]
+-------------------+---------+--------+
| Model             | Version | Status |
after allocation    : free:  7.47 GB, total: 15.78 GB, used:  8.31 GB
+-------------------+---------+--------+
| fastertransformer | 1       | READY  |
+-------------------+---------+--------+I0721 09:01:19.668137 1 metrics.cc:864] Collecting metrics for GPU 0: Tesla V100-SXM2-16GB
I0721 09:01:19.668167 1 metrics.cc:864] Collecting metrics for GPU 1: Tesla V100-SXM2-16GB
I0721 09:01:19.669954 1 metrics.cc:757] Collecting CPU metrics
I0721 09:01:19.670150 1 tritonserver.cc:2264]
+----------------------------------+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| Option                           | Value                                                                                                                                                                                                |
+----------------------------------+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| server_id                        | triton                                                                                                                                                                                               |
| server_version                   | 2.29.0                                                                                                                                                                                               |
| server_extensions                | classification sequence model_repository model_repository(unload_dependents) schedule_policy model_configuration system_shared_memory cuda_shared_memory binary_tensor_data statistics trace logging |
| model_repository_path[0]         | /mnt/triton_repo                                                                                                                                                                                     |
| model_control_mode               | MODE_NONE                                                                                                                                                                                            |
| strict_model_config              | 0                                                                                                                                                                                                    |
| rate_limit                       | OFF                                                                                                                                                                                                  |
| pinned_memory_pool_byte_size     | 268435456                                                                                                                                                                                            |
| cuda_memory_pool_byte_size{0}    | 67108864                                                                                                                                                                                             |
| cuda_memory_pool_byte_size{1}    | 67108864                                                                                                                                                                                             |
| response_cache_byte_size         | 0                                                                                                                                                                                                    |
| min_supported_compute_capability | 6.0                                                                                                                                                                                                  |
| strict_readiness                 | 1                                                                                                                                                                                                    |
| exit_timeout                     | 30                                                                                                                                                                                                   |
+----------------------------------+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+I0721 09:01:19.672326 1 grpc_server.cc:4819] Started GRPCInferenceService at 0.0.0.0:8001
I0721 09:01:19.672597 1 http_server.cc:3477] Started HTTPService at 0.0.0.0:8000
I0721 09:01:19.714356 1 http_server.cc:184] Started Metrics Service at 0.0.0.0:8002

06 服務請求

啟動 forward 進行驗證：

# 使用 kubectl 啟動port-forward
kubectl  -n default-group port-forward svc/ft-triton-bloom-1-tritoninferenceserver 8001:8001

這里我們使用 Triton Server 提供的 python SDK 所編寫的腳本來向 Triton Server 發起請求。腳本中主要完成三件事情：

通過 huggingface 中 bloom-7b1 對應的分詞器對 query 進行分詞和 token 轉換
通過 triton server SDK 向 triton server 發起請求
通過分詞器對 output token 進行轉換，拿到最終的結果

import os, sys
#from tkinter import _Padding
import numpy as np
import json
import torch
#import tritongrpcclient
import argparse
import time
from transformers import AutoTokenizer
import tritonclient.grpc as grpcclient# create tokenizer
tokenizer = AutoTokenizer.from_pretrained('/mnt/model/bloom-7b1', padding_side='right')
tokenizer.pad_token_id = tokenizer.eos_token_iddef load_image(img_path: str):"""Loads an encoded image as an array of bytes."""return np.fromfile(img_path, dtype='uint8')def tokeninze(query):# encodeencoded_inputs = tokenizer(query, padding=True, return_tensors='pt')input_token_ids = encoded_inputs['input_ids'].int()input_lengths = encoded_inputs['attention_mask'].sum(dim=-1, dtype=torch.int32).view(-1, 1)return input_token_ids.numpy().astype('uint32'), input_lengths.numpy().astype('uint32')if __name__ == "__main__":parser = argparse.ArgumentParser()parser.add_argument("--model_name",type=str,required=False,default="fastertransformer",help="Model name")parser.add_argument("--url",type=str,required=False,default="localhost:8001",help="Inference server URL. Default is localhost:8001.")parser.add_argument('-v',"--verbose",action="store_true",required=False,default=False,help='Enable verbose output')args = parser.parse_args()# 1.創建clienttry:triton_client = grpcclient.InferenceServerClient(url=args.url, verbose=args.verbose)except Exception as e:print("channel creation failed: " + str(e))sys.exit(1)output_name = "OUTPUT"# 2) 設置inputinputs = []## 2.1) input_idsquery="deepspeed is"input_ids, input_lengths = tokeninze(query)inputs.append(grpcclient.InferInput("input_ids", input_ids.shape, "UINT32"))inputs[0].set_data_from_numpy(input_ids)## 2.2) input_lengthinputs.append(grpcclient.InferInput("input_lengths", input_lengths.shape, "UINT32"))inputs[1].set_data_from_numpy(input_lengths)## 2.3） output lengthoutput_len=32output_len_np = np.array([[output_len]], dtype=np.uintc)inputs.append(grpcclient.InferInput("request_output_len", output_len_np.shape, "UINT32"))inputs[2].set_data_from_numpy(output_len_np)# 3) 設置outputoutputs = []outputs.append(grpcclient.InferRequestedOutput("output_ids"))# 4) 發起請求start_time = time.time()results = triton_client.infer(model_name=args.model_name, inputs=inputs,  outputs=outputs)latency = time.time() - start_time# 5) 結果處理：轉化為numpy 類型，計算max，轉化labeloutput0_data = results.as_numpy("output_ids")print(output0_data.shape)result = tokenizer.batch_decode(output0_data[0])print(result)

發起 client 請求命令如下：

$python3 bloom_7b_client.py 
(1, 1, 36)
['deepspeed is the speed of the ship at the time of the collision, and the\ndeepspeed of the other ship is the speed of the other ship
at the time']

07 總結

本文我們通過 Bloom-7b1 模型展示了如何在云原生 AI 套件中使用 FasterTransformer 對大語言模型進行加速，通過與 HuggingFace 的版本對比可以帶來 2.5 倍的性能提升。后續我們會逐步推出更多大模型相關的推理加速方案，以滿足不同的業務需求，大家敬請期待。

如果您對 Bloom 模型的微調訓練感興趣，您也可以點擊閱讀原文，參與實驗場景，體驗一鍵訓練大模型及部署 GPU 共享推理服務。

另外，歡迎加入云原生 AI 套件客戶交流釘釘群同我們一起探討（群號：33214567）。

相關鏈接：

[1] 云原生 AI 套件開發者使用指南

https://help.aliyun.com/zh/ack/cloud-native-ai-suite/getting-started/cloud-native-ai-component-set-user-guide

[2] 部署 PyTorch 模型推理服務

https://help.aliyun.com/zh/ack/cloud-native-ai-suite/user-guide/deploy-a-pytorch-model-as-an-inference-service?spm=a2c4g.11186623.0.0.2267225carYzgA

[3] 創建包含 GPU 的 Kubernetes 集群

https://help.aliyun.com/zh/ack/ack-managed-and-ack-dedicated/user-guide/use-gpu-scheduling-in-ack-clusters#task-1664343

[4] 安裝云原生 AI 套件

https://help.aliyun.com/document_detail/212117.htm#task-1917487

[5] 開始使用 OSS

https://help.aliyun.com/zh/oss/getting-started/getting-started-with-oss

[6] 使用 OSS 靜態存儲卷

https://help.aliyun.com/zh/ack/ack-managed-and-ack-dedicated/user-guide/mount-statically-provisioned-oss-volumes

點擊此處，體驗一鍵訓練大模型及部署 GPU 共享推理服務。