一. 背景介紹

在字節(jié)跳動(dòng)，基于深度學(xué)習(xí)的應(yīng)用遍地開(kāi)花，工程師關(guān)注模型效果的同時(shí)也需要關(guān)注線上服務(wù)一致性和性能，早期這通常需要算法專家和工程專家分工合作并緊密配合來(lái)完成，這種模式存在比較高的 diff 排查驗(yàn)證等成本。

隨著 PyTorch/TensorFlow 框架的流行，深度學(xué)習(xí)模型訓(xùn)練和在線推理完成了統(tǒng)一，開(kāi)發(fā)者僅需要關(guān)注具體算法邏輯，調(diào)用框架的 Python API 完成訓(xùn)練驗(yàn)證過(guò)程即可，之后模型可以很方便的序列化導(dǎo)出，并由統(tǒng)一的高性能 C++ 引擎完成推理工作。提升了開(kāi)發(fā)者訓(xùn)練到部署的體驗(yàn)。

然而，完整的服務(wù)通常還存在大量的預(yù)處理/后處理等業(yè)務(wù)邏輯，這類邏輯通常是把各種輸入經(jīng)過(guò)加工處理轉(zhuǎn)變?yōu)?Tensor，再輸入到模型，之后模型的輸出 Tensor 再加工成目標(biāo)格式，一些典型的場(chǎng)景如下：

• Bert
• Resnet

我們的目標(biāo)就是為以上端到端的過(guò)程，提供自動(dòng)化且統(tǒng)一的訓(xùn)練、推理方案，減輕人工開(kāi)發(fā)推理過(guò)程、對(duì)齊 diff 等一系列問(wèn)題，實(shí)現(xiàn)大規(guī)模的統(tǒng)一部署方案。

二. 核心問(wèn)題

PyTorch/TensorFlow 等框架相對(duì)已經(jīng)解決了模型的訓(xùn)練/推理統(tǒng)一的問(wèn)題，因此模型計(jì)算本身不存在訓(xùn)推一體的問(wèn)題了（算子性能優(yōu)化不在本次討論范圍）。

核心要解決的問(wèn)題就是：預(yù)處理和后處理需要提供高性能訓(xùn)推一體的方案。

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

• 性能：此方案大多基于虛擬機(jī)實(shí)現(xiàn)，虛擬機(jī)方案靈活并且非?？煽?，但深度學(xué)習(xí)框架中的虛擬機(jī)大多通常性能不夠優(yōu)良。補(bǔ)充說(shuō)明一下，框架早期都是為 Tensor 計(jì)算設(shè)計(jì)，數(shù)組計(jì)算每個(gè)算子成本很高，虛擬機(jī)的派發(fā)和調(diào)度成本可以忽略。但是，移植到程序語(yǔ)言編程層面開(kāi)銷難以忽略，代碼寫(xiě)多了就會(huì)成為性能瓶頸。據(jù)測(cè)試，TorchScript 解釋器性能只有 Python 的 1/5 左右，tf.function 性能更差一些。
• 功能不全：事實(shí)上應(yīng)用到真實(shí)場(chǎng)景中，我們?nèi)匀豢梢哉页龊芏? tf.function/TorchScript 不支持的重要功能，比如：自定義的資源不能打包，只能序列化內(nèi)置類型；字符串只能做 bytes 處理，中文等 unicode 會(huì)造成 diff；容器必須同構(gòu)，不支持自定義類型等等...

再者，還有很多非深度學(xué)習(xí)任務(wù)，比如在自然語(yǔ)言處理中仍然有很多非深度學(xué)習(xí)的應(yīng)用或者子任務(wù)，如序列標(biāo)注，語(yǔ)言模型解碼，樹(shù)模型的人工特征構(gòu)造等任務(wù)，這些通常具有更靈活的特征范式，但同時(shí)都沒(méi)有完整實(shí)現(xiàn)端到端的訓(xùn)推一體方案，仍然有大量的開(kāi)發(fā)以及正確性校驗(yàn)工作。

為了解決上述問(wèn)題，我們開(kāi)發(fā)了一套基于編譯的預(yù)處理方案：MATXScript！

三. MATXScript

在深度學(xué)習(xí)算法開(kāi)發(fā)中，開(kāi)發(fā)者通常使用 Python 進(jìn)行快速迭代和實(shí)驗(yàn)，同時(shí)使用 C++ 開(kāi)發(fā)高性能的線上服務(wù)，其中正確性校驗(yàn)和服務(wù)開(kāi)發(fā)都會(huì)成為較重負(fù)擔(dān)！

MatxScript（https://github.com/bytedance/matxscript） 是一個(gè) Python 子語(yǔ)言的 AOT 編譯器，可以自動(dòng)化將 Python 翻譯成 C++，并提供一鍵打包發(fā)布功能。使用 MATXScript 可以讓開(kāi)發(fā)者快速進(jìn)行模型迭代的同時(shí)以較低成本完成高性能服務(wù)的部署。

核心架構(gòu)如下：

• 最底層是純 C++/CUDA 的基礎(chǔ)庫(kù)，由高性能算子專家開(kāi)發(fā)。
• 在基礎(chǔ)庫(kù)之上，準(zhǔn)守約定封裝出來(lái) Python 的 庫(kù)，可以用在 training 過(guò)程中。
• 需要 inferencing 時(shí)，利用 MATXScript 可以把 Python 代碼，翻譯成對(duì)等的 C++ 代碼，編譯成動(dòng)態(tài)鏈接庫(kù)，加上模型及其他依賴的資源，一起打包發(fā)布即可。

其中，編譯器作用非常關(guān)鍵，其核心流程如下：

通過(guò)以上流程，用戶所編寫(xiě)的預(yù)處理代碼，可以被編譯成 Pipeline 中的一個(gè) JitOp，為了把前后處理和模型聯(lián)動(dòng)，我們還開(kāi)發(fā)了 tracing 系統(tǒng)（接口設(shè)計(jì)上參考了 PyTorch），架構(gòu)如下：

基于 MATXScript，我們可以訓(xùn)練和推理使用同一套代碼，大大降低了模型部署的成本。同時(shí)，架構(gòu)和算法得到了解耦，算法同學(xué)完全使用 Python 工作即可，架構(gòu)同學(xué)專注于編譯器開(kāi)發(fā)及 Runtime 優(yōu)化，在字節(jié)跳動(dòng)，此方案得到了大規(guī)模部署驗(yàn)證！

四. 小試牛刀

此處以最簡(jiǎn)單的英文文本預(yù)處理為例，展示一下 MATXScript 如何使用。

目標(biāo)：把一段英文文本轉(zhuǎn)成 indexes

1. 編寫(xiě)一個(gè)基本的查字典的邏輯

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. 編寫(xiě) Pipeline

import matx

class WorkFlow:
    def __init__(self):
        # 此處會(huì)進(jìn)行代碼編譯，Python 代碼自動(dòng)編譯封裝為 Callable 對(duì)象
        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 導(dǎo)出到 磁盤(pán)

# 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;
}

完整的代碼見(jiàn)：https://github.com/bytedance/matxscript/tree/main/examples/text2ids

小結(jié)：以上是一個(gè)非常簡(jiǎn)單的純 Python 實(shí)現(xiàn)的預(yù)處理邏輯，且能被一段通用的 C++ 代碼加載運(yùn)行，下面我們結(jié)合模型展示一個(gè)實(shí)際的多模態(tài)端到端案例！

五. 多模態(tài)案例

此處以圖文多模態(tài)(Bert+Resnet)為例，模型使用 PyTorch 編寫(xiě)，展示訓(xùn)練和部署中實(shí)際的工作。

1. 配置環(huán)境
   a. 配置 gcc/cuda 等基礎(chǔ)設(shè)施（通常是運(yùn)維同學(xué)已經(jīng)搞定）
   b. 安裝 MATXScript 及基于此開(kāi)發(fā)的基礎(chǔ)庫(kù)(text、vision等)
2. 編寫(xiě)模型代碼
   a. 此處省略，大家可以參考論文或其他開(kāi)源實(shí)現(xiàn)自行搞定
3. 編寫(xiě)預(yù)處理代碼

    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] 是對(duì)齊 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 可以當(dāng)成一個(gè)正常的 Python Class 接入 Dataset 即可
   b. VisionPipeline 涉及到 GPU 預(yù)處理，更適合按 batch 進(jìn)行處理，需要自己?jiǎn)为?dú)構(gòu)造一個(gè) DataLoader（這里埋個(gè)點(diǎn)，之后會(huì)開(kāi)源字節(jié)跳動(dòng)內(nèi)部基于多線程的 DataLoader）
5. 加上模型代碼，開(kāi)始訓(xùn)練吧
6. 導(dǎo)出端到端的 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')

小結(jié)：經(jīng)過(guò)以上步驟，我們即可完成端到端的訓(xùn)練&發(fā)布工作，且整個(gè)過(guò)程是純 Python 代碼完成的，可以完全由算法同學(xué)自己控制。當(dāng)然，如果模型計(jì)算本身還有性能問(wèn)題，也是可以在背后通過(guò)自動(dòng)改圖優(yōu)化工作完成。

注：完整代碼示例見(jiàn) https://github.com/bytedance/matxscript/tree/main/examples/e2e_multi_modal

六. 統(tǒng)一Server

在上個(gè)章節(jié)，我們得到了一個(gè)算法同學(xué)發(fā)布的模型包，本章節(jié)論述如果用統(tǒng)一的服務(wù)進(jìn)行加載和運(yùn)行。

完整的 Server 包括：IDL 協(xié)議、Batching 策略、進(jìn)/線程調(diào)度和排布、模型推理...

這里，我們只討論模型推理這塊，其他的都是可以按約定開(kāi)發(fā)即可。我們以一個(gè) main 函數(shù)來(lái)示例模型加載和運(yùn)行的過(guò)程：

#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;
}

以上代碼就是最簡(jiǎn)單的一個(gè) C++ 加載多模態(tài)模型的案例，對(duì) Server 開(kāi)發(fā)的同學(xué)來(lái)說(shuō)，只需進(jìn)行簡(jiǎn)單的抽象和約定，即可把上述代碼改造成一個(gè)統(tǒng)一的 C++ 模型服務(wù)框架。

七. 更多信息

我們是字節(jié)跳動(dòng)-AML-機(jī)器學(xué)習(xí)系統(tǒng)團(tuán)隊(duì)，致力于為公司提供統(tǒng)一的高性能訓(xùn)推一體化框架，同時(shí)也會(huì)通過(guò)火山引擎機(jī)器學(xué)習(xí)平臺(tái)服務(wù)于合作企業(yè)，火山引擎機(jī)器學(xué)習(xí)平臺(tái)預(yù)計(jì) 2023 年起提供 MATX 的相關(guān)支持，包括預(yù)置鏡像環(huán)境、常用場(chǎng)景的公開(kāi)樣例、企業(yè)接入和使用過(guò)程中的技術(shù)保障等，可以達(dá)到訓(xùn)練和推理場(chǎng)景低成本加速和一體化的效果。歡迎在 https://www.volcengine.com/product/ml-platform 詳細(xì)了解我們的產(chǎn)品。