位元組跳動模型大規模部署實戰

一. 背景介紹

在位元組跳動，基於深度學習的應用遍地開花，工程師關注模型效果的同時也需要關注線上服務一致性和效能，早期這通常需要演算法專家和工程專家分工合作並緊密配合來完成，這種模式存在比較高的diff 排查驗證等成本。

隨著PyTorch/TensorFlow 框架的流行，深度學習模型訓練和線上推理完成了統一，開發者只需要專注於具體演算法邏輯，呼叫框架的Python API 完成訓練驗證過程即可，之後模型可以很方便的序列化導出，並由統一的高性能C 引擎完成推理工作。提升了開發者訓練到部署的體驗。

然而，完整的服務通常還存在大量的預處理/後處理等業務邏輯，這類邏輯通常是把各種輸入經過加工處理轉變為Tensor，再輸入到模型，之後模型的輸出Tensor 再加工成目標格式，一些典型的場景如下：

• Bert
• #Resnet

我們的目標就是為以上端到端的過程，提供自動化且統一的訓練、推理方案，減輕人工開發推理過程、對齊diff 等一系列問題，實現大規模的統一部署方案。

二. 核心問題

PyTorch/TensorFlow 等框架相對已經解決了模型的訓練/推理統一的問題，因此模型計算本身不存在訓推一體的問題了（算子性能優化不在本次討論範圍）。

核心要解決的問題就是：預處理和後處理需要提供高效能訓推一體的方案。

對於此類邏輯，TensorFlow 2.x 提供了 tf.function(還不完善)，PyTorch 提供了 TorchScript，其無一例外都是選擇了原生 Python 語法子集。但即使強大如此，仍然存在不可忽略的問題：

• 效能：此方案大多基於虛擬機實現，虛擬機方案靈活且非常可控，但深度學習框架中的虛擬機器大多通常性能不夠優良。補充說明一下，框架早期都是為 Tensor 運算設計，陣列計算每個算子成本很高，虛擬機器的派發和調度成本可以忽略。但是，移植到程式語言程式設計層面開銷難以忽略，程式碼寫多了就會成為效能瓶頸。根據測試，TorchScript 解譯器效能只有 Python 的 1/5 左右，tf.function 效能更差一些。
• 功能不全：事實上應用在真實場景中，我們仍然可以找出很多tf.function/TorchScript 不支援的重要功能，例如：自訂的資源不能打包，只能序列化內建類型；字串只能做bytes 處理，中文等unicode 會造成diff；容器必須同構，不支援自訂類型等等...

再者，還有很多非深度學習任務，例如在自然語言處理中仍然有很多非深度學習的應用或子任務，如序列標註，語言模型解碼，樹模型的人工特徵構造等任務，這些通常具有更靈活的特徵範式，但同時都沒有完整實現端到端的訓推一體方案，仍有大量的發展以及正確性校驗工作。

為了解決上述問題，我們開發了一套基於編譯的預處理方案：MATXScript！

三. MATXScript

在深度學習演算法開發中，開發者通常使用Python 進行快速迭代和實驗，同時使用C 開發高效能的線上服務，其中正確性校驗和服務開發都會成為較重負擔！

MatxScript（https://github.com/bytedance/matxscript） 是一個 Python 子語言的 AOT 編譯器，可以自動化將 Python 翻譯成 C ，並提供一鍵打包發布功能。使用 MATXScript 可以讓開發者快速進行模型迭代的同時以較低成本完成高效能服務的部署。

核心架構如下：

#

• 最底層是純 C /CUDA 的基礎函式庫，由高效能算子專家開發。
• 在基礎函式庫之上，準守約定封裝出來 Python 的 函式庫，可以用在 training 流程中。
• 需要 inferencing 時，利用 MATXScript 可以把 Python 程式碼，翻譯成對等的 C 程式碼，編譯成動態連結函式庫，加上模型及其他所依賴的資源，一起打包發布即可。

其中，編譯器作用非常關鍵，其核心流程如下：

#透過上述流程，使用者所寫的預處理程式碼，可以被編譯成Pipeline 中的一個JitOp，為了把前後處理和模型連動，我們也開發了tracing 系統（介面設計上參考了PyTorch），架構如下：

#基於MATXScript，我們可以訓練和推理使用同一套程式碼，大大降低了模型部署的成本。同時，架構和演算法得到了解耦，演算法同學完全使用 Python 工作即可，架構同學專注於編譯器開發及 Runtime 優化，在字節跳動，此方案得到了大規模部署驗證！

四.小試牛刀

此處以最簡單的英文文字預處理為例，展示一下 MATXScript 如何使用。

目標：把一段英文文本轉成indexes

1. 寫一個基本的查字典的邏輯

class Text2Ids:
def __init__(self) -> None:
self.table: Dict[str, int] = {
"hello": 0,
"world": 1,
"[UNK]": 2,
}

def lookup(self, word: str)
return self.table.get(word, 2)

def__call__ (self, words: List[str])
return [self.lookup(w) for w in words]

2. 寫Pipeline

import matx

class WorkFlow:
def __init__(self):
# 此处会进行代码编译，Python 代码自动编译封装为 Callable 对象
self.text2ids = matx.script(Text2Ids)()

def process(self, texts):
ids = self.text2ids(texts)
return ids

# test
handler = WorkFlow()
print(handler.process("hello world unknown"))
# output: [0, 1, 2]

3. Trace 匯出到磁碟

# dump
mod = matx.trace(handler.process, "hello world")
print(mod.run({"texts": "hello world"}))
mod.save('./my_dir')
# load
mod = matx.load('./my_dir', -1)
print(mod.run({"texts": "hello world"}))

4. #C 載入

#include <string>
#include <vector>
#include <map>
#include <iostream>
#include <matxscript/pipeline/tx_session.h>
using namespace ::matxscript::runtime;
int main()
{
// test case
std::unordered_map<std::string, RTValue> feed_dict;
feed_dict.emplace("texts", Unicode(U"hello world"));
std::vector<std::pair<std::string, RTValue>> result;
const char* module_path = "./my_dir";
const char* module_name = "model.spec.json";
{
// -1 mean cpu
auto sess = TXSession::Load(module_path, module_name, -1);
auto result = sess->Run(feed_dict);
for (auto& r : result) {
std::cout << "key: " << r.first << ", value: " << r.second << std::endl;
}
}
return 0;
}

完整的程式碼請見：https://github. com/bytedance/matxscript/tree/main/examples/text2ids

小結：以上是一個非常簡單的純Python 實作的預處理邏輯，且能被一段通用的C 程式碼載入運行，下面我們結合模型展示一個實際的多模態端對端案例！

五. 多模態案例

此處以圖文多模態(Bert Resnet)為例，模型使用 PyTorch 編寫，展示訓練和部署中實際的工作。

1. 配置環境
   a. 配置gcc/cuda 等基礎設施（通常是維運同學已經搞定）
   b. 安裝MATXScript 及基於此開發的基礎庫(text、vision等等)
2. 寫模型程式碼
   a. 這裡省略，大家可以參考論文或其他開源實作自行搞定
3. 編寫預處理程式碼

    a . text

from typing import List, Dict, Tuple
import libcut
import matx
class Vocabulary:
...
def utf8_decoder(s: List[bytes]):
return [x.decode() for x in s]
class TextNDArrayBuilder:
...
class TextPipeline:
def __init__(self, mode: str = "eval"):
self.mode = mode
self.cut_engine = libcut.Cutter('/path/to/cut_models', ...)
self.vocab = matx.script(Vocabulary)('/path/to/vocab.txt')
self.decoder = matx.script(utf8_decoder)
self.input_builder = matx.script(TextNDArrayBuilder)(self.vocab)
def process(self, text: List[bytes]):
# List[bytes] 是对齐 C++ 的 vector<string>
text: List[str] = self.decoder(text)
words: List[List[str]] = self.cut_engine(text)
batch_ids: List[List[int]] = self.vocab(words)
input_ids, segment_ids, mask_ids = self.input_builder(batch_ids, 32)
if self.mode == "train":
return input_ids.torch(), segment_ids.torch(), mask_ids.torch()
return input_ids, segment_ids, mask_ids

    b. vision

from typing import List, Dict, Tuple
import matx
from matx import vision
class VisionPipeline:
def __init__(self,
 device_id: int = 0,
 mode: str = "eval",
 image_size: int = 224,):
self.is_training = mode == 'train'
self.mode = mode
...
def process(self, image,):
if self.is_training:
decode_nds = self.random_crop_decode(image)
flip_nds = self.random_flip(decode_nds)
resize_nds = self.resize(flip_nds)
transpose_nd = self.transpose_norm(resize_nds, vision.SYNC)
else:
decode_nds = self.decode(image)
resize_nds = self.resize(decode_nds)
crop_nds = self.center_crop(resize_nds)
transpose_nd = self.transpose_norm(crop_nds, vision.SYNC)
if self.mode == "trace":
return transpose_nd
return transpose_nd.torch()

4. 存取DataLoader
   a. TextPipeline 可以當成一個正常的Python Class 存取Dataset 即可
   b.到GPU 預處理，更適合按batch 進行處理，需要自行單獨建構一個DataLoader（這裡埋個點，之後會開源位元組跳動內部基於多執行緒的DataLoader）
5. 加上模型程式碼，開始訓練吧
6. 導出端到端的Inference Model

class MultimodalEvalPipeline:
def __init__(self):
self.text_pipe = TextPipeline(mode="eval", ...)
self.vision_pipe = VisionPipeline(mode="eval", ...)
self.torch_model = torch.jit.load('/path/to/multimodal.jit', map_locatinotallow='cuda:0')
self.tx_model_op = matx.script(self.torch_model, device=0)

def eval(self, texts: List[bytes], images: List[bytes])
input_ids, segment_ids, mask_ids = self.text_pipe.process(texts)
images = self.vision_pipe.process(images)
scores = self.tx_model_op(input_ids, segment_ids, mask_ids, images)
return scores
# examples
example_batch_size = 8
text_examples = ['hello, world'.encode()] * example_batch_size
with open('/path/image.jpg', 'rb') as f:
image_example = f.read()
image_examples = [image_example] * example_batch_size
# pipeline instance
pipe = MultimodalEvalPipeline(...)
mod = matx.trace(pipe.eval, text_examples, image_examples)
# test
print(mod.run({"texts": text_examples, "images": image_examples}))
# save
mod.save('/path/to/my_multimodal')

小結：經過以上步驟，我們即可完成端到端的訓練&發布工作，且整個過程是純Python 程式碼完成的，可以完全由演算法同學自己控制。當然，如果模型計算本身還有效能問題，也是可以在背後透過自動改圖優化工作完成。

附註：完整程式碼範例請見https://github.com/bytedance/matxscript/tree/main/examples/e2e_multi_modal

六. 統一Server

在上個章節，我們得到了一個演算法同學發布的模型包，本章節論述如果用統一的服務進行載入和運行。

完整的Server 包括：IDL 協定、Batching 策略、進/執行緒調度和排布、模型推理...

這裡，我們只討論模型推理這塊，其他的都是可依約定開發即可。我們以一個main 函數來範例模型載入和運行的過程：

#include <string>
#include <vector>
#include <map>
#include <iostream>
#include <matxscript/pipeline/tx_session.h>
using namespace ::matxscript::runtime;
int main()
{
// test case
std::unordered_map<std::string, RTValue> feed_dict;
feed_dict.emplace("texts", List({String("hello world")}));
feed_dict.emplace("images", List({String("......")}));
std::vector<std::pair<std::string, RTValue>> result;
const char* module_path = "/path/to/my_multimodal";
const char* module_name = "model.spec.json";
{
// cuda:0
auto sess = TXSession::Load(module_path, module_name, 0);
auto result = sess->Run(feed_dict);
for (auto& r : result) {
std::cout << "key: " << r.first << ", value: " << r.second << std::endl;
}
}
return 0;
}

以上程式碼就是最簡單的一個C 載入多模態模型的案例，對Server 開發的同學來說，只需進行簡單的抽象和約定，即可把上述程式碼改造成一個統一的C 模型服務架構。

七. 更多資訊

我們是位元組跳動-AML-機器學習系統團隊，致力於為公司提供統一的高性能訓推一體化框架，同時也將透過火山引擎機器學習平台服務於合作企業，火山引擎機器學習平台預計2023 年起提供MATX 的相關支持，包括預置鏡像環境、常用場景的公開樣本、企業接入和使用過程中的技術保障等，可以達到訓練和推理場景低成本加速和一體化的效果。歡迎在 https://www.volcengine.com/product/ml-platform 詳細了解我們的產品。

以上是位元組跳動模型大規模部署實戰的詳細內容。更多資訊請關注PHP中文網其他相關文章！