使用推理服務(wù)器加速大型Transformer模型的推理

這是討論 NVIDIA FasterTransformer 庫(kù)的兩部分系列的第一部分，該庫(kù)是用于對(duì)任意大小（多達(dá)數(shù)萬(wàn)億個(gè)參數(shù)）的Transformer進(jìn)行分布式推理的最快庫(kù)之一。它提供了 FasterTransformer 的概述，包括使用該庫(kù)的好處。

使用 FasterTransformer 和 Triton 推理服務(wù)器部署 GPT-J 和 T5（第 2 部分）是一個(gè)指南，說(shuō)明了使用 FasterTransformer 庫(kù)和 Triton 推理服務(wù)器以具有張量并行性的最佳方式為 T5-3B 和 GPT-J 6B 模型提供服務(wù)。

Transformers 是當(dāng)今最具影響力的 AI 模型架構(gòu)之一，正在塑造未來(lái) AI 研發(fā)的方向。它們最初是作為自然語(yǔ)言處理 （NLP） 的工具而發(fā)明的，現(xiàn)在幾乎用于任何 AI 任務(wù)，包括計(jì)算機(jī)視覺(jué)、自動(dòng)語(yǔ)音識(shí)別、分子結(jié)構(gòu)分類(lèi)和金融數(shù)據(jù)處理。考慮到如此廣泛使用的是注意力機(jī)制，它顯著提高了模型的計(jì)算效率、質(zhì)量和準(zhǔn)確性。

具有數(shù)千億參數(shù)的大型基于 Transformer 的模型的行為就像一個(gè)巨大的百科全書(shū)和大腦，其中包含有關(guān)它所學(xué)到的一切的信息。他們以獨(dú)特的方式對(duì)所有這些知識(shí)進(jìn)行結(jié)構(gòu)化、表示和總結(jié)。擁有具有大量先驗(yàn)知識(shí)的此類(lèi)模型使我們能夠使用新的強(qiáng)大的一次性或少量學(xué)習(xí)技術(shù)來(lái)解決許多 NLP 任務(wù)。

由于它們的計(jì)算效率，Transformer可以很好地?cái)U(kuò)展——通過(guò)增加網(wǎng)絡(luò)的規(guī)模和訓(xùn)練數(shù)據(jù)的數(shù)量，研究人員可以改善觀察并提高準(zhǔn)確性。

然而，訓(xùn)練如此大的模型并非易事。這些模型可能需要比一個(gè) GPU 供應(yīng)更多的內(nèi)存——甚至是數(shù)百個(gè) GPU。值得慶幸的是，NVIDIA 研究人員已經(jīng)創(chuàng)建了強(qiáng)大的開(kāi)源工具，例如 NeMo Megatron，可以優(yōu)化訓(xùn)練過(guò)程。

快速和優(yōu)化的推理使企業(yè)能夠充分發(fā)揮這些大型模型的潛力。最新研究表明，增加模型和數(shù)據(jù)集的大小可以提高這種模型在不同領(lǐng)域（NLP、CV 等）下游任務(wù)上的質(zhì)量。

同時(shí)，數(shù)據(jù)表明這種技術(shù)也適用于多域任務(wù)。 （例如，參見(jiàn) OpenAI 的 DALLE-2 和 Google 的 Imagen 等關(guān)于文本到圖像生成的研究論文。）依賴于大型模型的“凍結(jié)”拷貝的 p-tuning 等研究方向甚至增加了擁有穩(wěn)定且優(yōu)化的推理流程。此類(lèi)大型模型的優(yōu)化推理需要分布式多 GPU 多節(jié)點(diǎn)解決方案。

用于加速推斷大型Transformer的庫(kù)

NVIDIA FasterTransformer （FT） 是一個(gè)庫(kù)，用于實(shí)現(xiàn)基于Transformer的神經(jīng)網(wǎng)絡(luò)推理的加速引擎，特別強(qiáng)調(diào)大型模型，以分布式方式跨越許多 GPU 和節(jié)點(diǎn)。

FasterTransformer 包含Transformer塊的高度優(yōu)化版本的實(shí)現(xiàn)，其中包含編碼器和解碼器部分。

使用此模塊，您可以運(yùn)行完整的編碼器-解碼器架構(gòu)（如 T5）以及僅編碼器模型（如 BERT）或僅解碼器模型（如 GPT）的推理。 它是用 C++/CUDA 編寫(xiě)的，依賴于高度優(yōu)化的 cuBLAS、cuBLASLt 和 cuSPARSELt 庫(kù)。 這使您可以在 GPU 上構(gòu)建最快的Transformer推理流程。

圖 1.使用張量并行（張量 MP 分區(qū)）和管道并行（管道 MP 分區(qū)），在四個(gè) GPU 之間分布了幾個(gè) transformer / attention 塊

與 NVIDIA TensorRT 等其他編譯器相比，F(xiàn)T 的顯著特點(diǎn)是它支持以分布式方式推斷大型Transformer模型。

上圖顯示了如何使用張量并行 （TP） 和流水線并行 （PP） 技術(shù)將具有多個(gè)經(jīng)典transformer/attention的神經(jīng)網(wǎng)絡(luò)拆分到多個(gè) GPU 和節(jié)點(diǎn)上。

當(dāng)每個(gè)張量被分成多個(gè)塊時(shí)，就會(huì)發(fā)生張量并行性，并且張量的每個(gè)塊都可以放置在單獨(dú)的 GPU 上。在計(jì)算過(guò)程中，每個(gè)塊在不同的 GPU 上單獨(dú)并行處理，并且可以通過(guò)組合來(lái)自多個(gè) GPU 的結(jié)果來(lái)計(jì)算結(jié)果（最終張量）。

當(dāng)模型被深度拆分并將不同的完整層放置到不同的 GPU/節(jié)點(diǎn)上時(shí)，就會(huì)發(fā)生流水線并行。

在底層，啟用節(jié)點(diǎn)間/節(jié)點(diǎn)內(nèi)通信依賴于 MPI 和 NVIDIA NCCL。使用此軟件堆棧，您可以在多個(gè) GPU 上以張量并行模式運(yùn)行大型Transformer，以減少計(jì)算延遲。

同時(shí)，TP 和 PP 可以結(jié)合在一起，在多 GPU 和多節(jié)點(diǎn)環(huán)境中運(yùn)行具有數(shù)十億和數(shù)萬(wàn)億個(gè)參數(shù)（相當(dāng)于 TB 級(jí)權(quán)重）的大型 Transformer 模型。

除了 C 中的源代碼，F(xiàn)asterTransformer 還提供 TensorFlow 集成（使用 TensorFlow 操作）、PyTorch 集成（使用 PyTorch 操作）和 Triton 集成作為后端。

目前，TensorFlow op 僅支持單 GPU，而 PyTorch op 和 Triton 后端都支持多 GPU 和多節(jié)點(diǎn)。

為了避免為模型并行性而拆分模型的額外工作，F(xiàn)asterTransformer 還提供了一個(gè)工具，用于將模型從不同格式拆分和轉(zhuǎn)換為 FasterTransformer 二進(jìn)制文件格式。然后 FasterTransformer 可以直接以二進(jìn)制格式加載模型。

目前，F(xiàn)T 支持 Megatron-LM GPT-3、GPT-J、BERT、ViT、Swin Transformer、Longformer、T5 和 XLNet 等模型。您可以在 GitHub 上的 FasterTransformer 存儲(chǔ)庫(kù)中查看最新的支持矩陣。

FT 適用于計(jì)算能力 》= 7.0 的 GPU，例如 V100、A10、A100 等。

圖 2.GPT-J 6B 模型推斷和加速比較

FasterTransformer 中的優(yōu)化

與深度學(xué)習(xí)訓(xùn)練的通用框架相比，F(xiàn)T 使您能夠獲得更快的推理管道，并且基于 Transformer 的 NN 具有更低的延遲和更高的吞吐量。

允許 FT 對(duì) GPT-3 和其他大型Transformer模型進(jìn)行最快推理的一些優(yōu)化技術(shù)包括：

層融合——預(yù)處理階段的一組技術(shù)，將多層神經(jīng)網(wǎng)絡(luò)組合成一個(gè)單一的神經(jīng)網(wǎng)絡(luò)，將使用一個(gè)單一的內(nèi)核進(jìn)行計(jì)算。 這種技術(shù)減少了數(shù)據(jù)傳輸并增加了數(shù)學(xué)密度，從而加速了推理階段的計(jì)算。 例如， multi-head attention 塊中的所有操作都可以合并到一個(gè)內(nèi)核中。

圖 3. NVIDIA Faster transformer 庫(kù)中緩存機(jī)制的演示

自回歸模型/激活緩存的推理優(yōu)化

為了防止通過(guò)Transformer重新計(jì)算每個(gè)新token 生成器的先前鍵和值，F(xiàn)T 分配一個(gè)緩沖區(qū)來(lái)在每一步存儲(chǔ)它們。

雖然需要一些額外的內(nèi)存使用，但 FT 可以節(jié)省重新計(jì)算的成本、在每一步分配緩沖區(qū)以及連接的成本。該過(guò)程的方案上圖 所示。相同的緩存機(jī)制用于 NN 的多個(gè)部分。

內(nèi)存優(yōu)化

與 BERT 等傳統(tǒng)模型不同，大型 Transformer 模型具有多達(dá)數(shù)萬(wàn)億個(gè)參數(shù)，占用數(shù)百 GB 存儲(chǔ)空間。即使我們以半精度存儲(chǔ)模型，GPT-3 175b 也需要 350 GB。因此有必要減少其他部分的內(nèi)存使用。

例如，在 FasterTransformer 中，我們?cè)诓煌慕獯a器層重用了激活/輸出的內(nèi)存緩沖區(qū)。由于 GPT-3 中的層數(shù)為 96，因此我們只需要 1/96 的內(nèi)存量用于激活。

使用 MPI 和 NCCL 實(shí)現(xiàn)節(jié)點(diǎn)間/節(jié)點(diǎn)內(nèi)通信并支持模型并行性

在 GPT 模型中，F(xiàn)asterTransormer 同時(shí)提供張量并行和流水線并行。 對(duì)于張量并行性，F(xiàn)asterTransformer 遵循了 Megatron 的思想。 對(duì)于自注意力塊和前饋網(wǎng)絡(luò)塊，F(xiàn)T 按行拆分第一個(gè)矩陣的權(quán)重，并按列拆分第二個(gè)矩陣的權(quán)重。 通過(guò)優(yōu)化，F(xiàn)T 可以將每個(gè) Transformer 塊的歸約操作減少到兩倍。

對(duì)于流程并行性，F(xiàn)asterTransformer 將整批請(qǐng)求拆分為多個(gè)微批，隱藏了通信的泡沫。 FasterTransformer 會(huì)針對(duì)不同情況自動(dòng)調(diào)整微批量大小。

MatMul 內(nèi)核自動(dòng)調(diào)整（GEMM 自動(dòng)調(diào)整）

矩陣乘法是基于Transformer的神經(jīng)網(wǎng)絡(luò)中主要和最繁重的操作。 FT 使用來(lái)自 CuBLAS 和 CuTLASS 庫(kù)的功能來(lái)執(zhí)行這些類(lèi)型的操作。 重要的是要知道 MatMul 操作可以在“硬件”級(jí)別使用不同的低級(jí)算法以數(shù)十種不同的方式執(zhí)行。

GemmBatchedEx 函數(shù)實(shí)現(xiàn) MatMul 操作，并以“cublasGemmAlgo_t”作為輸入?yún)?shù)。 使用此參數(shù)，您可以選擇不同的底層算法進(jìn)行操作。

FasterTransformer 庫(kù)使用此參數(shù)對(duì)所有底層算法進(jìn)行實(shí)時(shí)基準(zhǔn)測(cè)試，并為模型的參數(shù)和您的輸入數(shù)據(jù)（注意層的大小、注意頭的數(shù)量、隱藏層的大?。┻x擇最佳的一個(gè)。 此外，F(xiàn)T 對(duì)網(wǎng)絡(luò)的某些部分使用硬件加速的底層函數(shù)，例如 __expf、__shfl_xor_sync。

精度較低的推理

FT 的內(nèi)核支持使用 fp16 和 int8 中的低精度輸入數(shù)據(jù)進(jìn)行推理。 由于較少的數(shù)據(jù)傳輸量和所需的內(nèi)存，這兩種機(jī)制都允許加速。 同時(shí)，int8 和 fp16 計(jì)算可以在特殊硬件上執(zhí)行，例如張Tensor Core（適用于從 Volta 開(kāi)始的所有 GPU 架構(gòu)），以及即將推出的 Hopper GPU 中的Transformer引擎。

更多

快速的 C++ BeamSearch 實(shí)現(xiàn)

針對(duì) TensorParallelism 8 模式優(yōu)化 all-reduce 當(dāng)模型的權(quán)重部分在 8 個(gè) GPU 之間拆分時(shí)

具有 FasterTransformer 后端的 NVIDIA Triton 推理服務(wù)器

NVIDIA Triton 推理服務(wù)器是一款開(kāi)源推理服務(wù)軟件，有助于標(biāo)準(zhǔn)化模型部署和執(zhí)行，在生產(chǎn)中提供快速且可擴(kuò)展的 AI。 Triton 穩(wěn)定且快速，允許您使用準(zhǔn)備好的 Docker 容器以簡(jiǎn)單的方式運(yùn)行 ML/DL 模型的推理，該容器僅使用一行代碼和簡(jiǎn)單的類(lèi)似 JSON 的配置。

Triton 支持使用多個(gè)后端的模型，例如 PyTorch、TorchScript、Tensorflow、ONNXRuntime 和 OpenVINO。 Triton 采用您在其中一個(gè)框架中訓(xùn)練的導(dǎo)出模型，并使用相應(yīng)的后端為您透明地運(yùn)行該模型進(jìn)行推理。 它也可以使用自定義后端進(jìn)行擴(kuò)展。 Triton 使用 HTTP/gRPC API 包裝您的模型，并為多種語(yǔ)言提供客戶端庫(kù)。

圖 4. Triton 推理服務(wù)器，具有多個(gè)后端，用于對(duì)使用不同框架訓(xùn)練的模型進(jìn)行推理

Triton 包含 FasterTransformer 庫(kù)作為后端（圖 4），該庫(kù)支持使用 TP 和 PP 運(yùn)行大型Transformer模型的分布式多 GPU、多節(jié)點(diǎn)推理。 今天，帶有 FasterTransformer 后端的 Triton 支持 GPT-J、GPT-Megatron 和 T5 模型。

關(guān)于作者

Denis Timonin 是一名深度學(xué)習(xí)解決方案架構(gòu)師和工程師，目前在 NVIDIA 從事世界上最大的神經(jīng)網(wǎng)絡(luò)的訓(xùn)練和推理，并在計(jì)算機(jī)視覺(jué)、自然語(yǔ)言處理和自動(dòng)語(yǔ)音識(shí)別領(lǐng)域構(gòu)建人工智能解決方案。在此之前，丹尼斯在華為研究了移動(dòng)設(shè)備的小型精確神經(jīng)網(wǎng)絡(luò)，并構(gòu)建了復(fù)雜的人工智能管道，用于金融數(shù)據(jù)分類(lèi)、醫(yī)學(xué)圖像分割、目標(biāo)檢測(cè)和跟蹤。

Bo Yang Hsueh 是 FasterTransformer 的領(lǐng)導(dǎo)者和主要開(kāi)發(fā)人員。三年前他參加了 transformer 加速賽。最近，他專(zhuān)注于大型 NLP 模型加速，包括 T5 和 GPT-J 等公共模型。楊波獲得國(guó)立交通大學(xué)計(jì)算機(jī)科學(xué)碩士學(xué)位。

審核編輯：郭婷