FasterTransformer GPT介紹

FasterTransformer GPT 介紹（翻譯）

GPT 是 Decooding 模型的一種變體，沒(méi)有 Encoder 模塊，沒(méi)有交叉多頭注意力模塊，使用 GeLU 作為激活函數(shù)。2020 年，OpenAI 在他們的論文中表明，使用非常龐大的模型和大量的訓(xùn)練數(shù)據(jù)可以顯著提高 GPT 模型的容量。但是，不可能將這樣的模型放入單個(gè) GPU 中。例如，最大的模型 GPT-3 有 1750 億個(gè)參數(shù)，half 數(shù)據(jù)類(lèi)型下大約需要 350 GB顯存。因此，多GPU，甚至多節(jié)點(diǎn)，是很有必要的。為了解決模型大小導(dǎo)致的延遲和內(nèi)存瓶頸，F(xiàn)asterTransformer 提供了高性能、低內(nèi)存占用的 kernel，并使用了模型并行技術(shù)。

支持的特性

Checkpoint converter

Huggingface

Megatron

Nemo Megatron

TensorFlow

Data type

限制:

Note:

權(quán)重被切分后，隱藏層的維度必須是 64 的倍數(shù)。

cuda kernel通常只為小的 batch（如32和64）和權(quán)重矩陣很大時(shí)提供性能優(yōu)勢(shì)。

權(quán)重的 PTQ 量化只支持 FP16/BF16。

僅支持 Volta 和更新的 GPU 架構(gòu)。

根據(jù)當(dāng)前 GPU 的情況，權(quán)重被提前離線預(yù)處理，以降低 TensorCore 做權(quán)重對(duì)齊的開(kāi)銷(xiāo)。目前，我們直接使用 FP32/BF16/FP16 權(quán)重并在推理前對(duì)其進(jìn)行量化。如果我們想存儲(chǔ)量化的權(quán)重，必須要在推理的 GPU 上來(lái)進(jìn)行預(yù)處理。

使用 torch API 時(shí)，int8 模式只能通過(guò) Parallel GPT Op 使用。Parallel GPT Op 也可以在單個(gè) GPU 上使用。

FP32

FP16

BF16

INT8 weight only PTQ.

INT8 with SmoothQuant

FP8 (Experimental)

Feature

Multi-GPU multi-node inference

Dynamic random seed

Stop tokens

Beam search and sampling are both supported

Loading FP32 or FP16 weights

Frameworks

TensorFlow

PyTorch

C++

Triton backend

FasterTransformer GPT 結(jié)構(gòu)（翻譯）

工作流

Fig 1. Workflow of GPT model

Fig 1展示了 FasterTransformer GPT 的工作流程。與 BERT 和編碼器-解碼器結(jié)構(gòu)不同，GPT 接收一些輸入 id 作為上下文，并生成相應(yīng)的輸出 id 作為響應(yīng)。在這個(gè)工作流中，主要的瓶頸是 GptDecoderLayer （Transformer塊），因?yàn)楫?dāng)我們?cè)黾訉訑?shù)的時(shí)候耗時(shí)也是線性增加的。在GPT-3中，GptDecoderLayer占用了大約95%的時(shí)間。

FasterTransformer把整個(gè)工作流分成了兩個(gè)部分。第一個(gè)部分是：“根據(jù)上下文context(也就是輸入ids)計(jì)算k/v cache”。第二個(gè)部分是：“自回歸的生成輸出ids”。這兩部分的操作類(lèi)似，但是selfAttention部分的輸入tensors的形狀是不一樣的。所以FasterTransformer提供了2種計(jì)算方式，如Fig2所示。

在DecoderSelfAttention里面，query的序列長(zhǎng)度總是1，所以我們使用自定義的fused masked multi-head attention kernel 進(jìn)行處理。另一方面，在ContextSelfAttention中，query的序列長(zhǎng)度最大時(shí)輸入的長(zhǎng)度，所以我們使用cuBLAS來(lái)利用TensorCore。

Fig 2. Comparison between different self attention.

這個(gè)地方?jīng)]有理解為什么要分成2個(gè)Attention，因?yàn)樽曰貧w的解碼也是需要把輸入的句子 padding 到最大的長(zhǎng)度吧。這里的seq_len為1的情況是什么時(shí)候發(fā)生呢？我看了一下hugging face的GPT，似乎沒(méi)有找到對(duì)應(yīng)的位置。然后在FasterTransformer的GPT C++實(shí)現(xiàn)中也沒(méi)有找到這個(gè)DecoderSelfAttention的實(shí)現(xiàn)：https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/models/multi_gpu_gpt/ParallelGptContextDecoder.cc 。不過(guò)本文主要是在后面介紹下 FasterTransformer 的優(yōu)化點(diǎn)以及優(yōu)缺點(diǎn)，這個(gè)暫時(shí)不影響解讀。

以下示例演示如何運(yùn)行多 GPU 和多節(jié)點(diǎn) GPT 模型。

examples/cpp/multi_gpu_gpt_example.cc: 它使用 MPI 來(lái)組織所有 GPU。

examples/cpp/multi_gpu_gpt_triton_example.cc: 它在節(jié)點(diǎn)內(nèi)使用多線程，節(jié)點(diǎn)間使用 MPI。此示例還演示了如何使用 FasterTransformer 的 Triton 后端 API 來(lái)運(yùn)行 GPT 模型。

examples/pytorch/gpt/multi_gpu_gpt_example.py: 這個(gè)例子和 examples/cpp/multi_gpu_gpt_example.cc 很類(lèi)似, 但是通過(guò) PyTorch OP 封裝了 FasterTransformer 的實(shí)例。

總之，運(yùn)行 GPT 模型的工作流程是:

通過(guò) MPI 或多線程初始化 NCCL 通信并設(shè)置張量并行和流水并行的等級(jí)。

按張量并行、流水并行的ranks和其它模型超參數(shù)加載權(quán)重。

按張量并行，流水并行的ranks和其它模型超參數(shù)創(chuàng)建ParalelGpt實(shí)例。

接收來(lái)自客戶(hù)端的請(qǐng)求并將請(qǐng)求轉(zhuǎn)換為ParallelGpt的輸入張量格式.

運(yùn)行 forward 函數(shù)

將 ParallelGpt 的輸出張量轉(zhuǎn)換為客戶(hù)端的響應(yīng)并返回響應(yīng)。

在c++示例代碼中，我們跳過(guò)第4步和第6步，通過(guò)examples/cpp/multi_gpu_gpt/start_ids.csv加載請(qǐng)求。在 PyTorch 示例代碼中，請(qǐng)求來(lái)自 PyTorch 端。在 Triton 示例代碼中，我們有從步驟 1 到步驟 6 的完整示例。

源代碼放在 src/fastertransformer/models/multi_gpu_gpt/ParallelGpt.cc 中。GPT的參數(shù)、輸入張量和輸出張量：

Constructor of GPT

Input of GPT

Output of GPT

beam_width 值直接由輸出形狀設(shè)置。當(dāng)output_ids的beam_width大于1時(shí)，F(xiàn)T會(huì)使用beam search來(lái)生成token；否則，F(xiàn)T 將使用 topk 或 topp 采樣。當(dāng) beam search 和 sampling 的輸入無(wú)效時(shí)，比如 beam width 1，top k 0，top p 0.0，F(xiàn)T 會(huì)自動(dòng)運(yùn)行 greedy search?！?/p>

優(yōu)化

kernel 優(yōu)化：很多 kernel 都是基于已經(jīng)高度優(yōu)化的解碼器和解碼模塊的 kernel。為了防止重新計(jì)算以前的key 和 value，我們將在每一步分配一個(gè)緩沖區(qū)來(lái)存儲(chǔ)它們。雖然它需要一些額外的內(nèi)存使用，但我們可以節(jié)省重新計(jì)算的成本，在每一步分配緩沖區(qū)以及串行的成本。

內(nèi)存優(yōu)化：與 BERT 等傳統(tǒng)模型不同，GPT-3 有 1750 億個(gè)參數(shù)，即使我們以半精度存儲(chǔ)模型也需要 350 GB。因此，我們必須減少其他部分的內(nèi)存使用。在 FasterTransformer 中，我們將重用不同解碼器層的內(nèi)存緩沖區(qū)。由于 GPT-3 的層數(shù)是 96，我們只需要 1/96 的內(nèi)存。

模型并行：在GPT模型中，F(xiàn)asterTransormer同時(shí)提供張量并行和流水線并行。對(duì)于張量并行，F(xiàn)asterTransformer 遵循了 Megatron 的思想。對(duì)于自注意力塊和前饋網(wǎng)絡(luò)塊，我們按行拆分第一個(gè)矩陣乘法的權(quán)重，按列拆分第二個(gè)矩陣乘法的權(quán)重。通過(guò)優(yōu)化，我們可以將每個(gè)Transformer塊的歸約操作減少到 2 次。工作流程如Fig 3 所示。對(duì)于流水并行，F(xiàn)asterTransformer 將整個(gè)Batch的請(qǐng)求拆分為多個(gè)Micro Batch并隱藏通信氣泡。FasterTransformer 會(huì)針對(duì)不同情況自動(dòng)調(diào)整微批量大小。用戶(hù)可以通過(guò)修改 gpt_config.ini 文件來(lái)調(diào)整模型并行度。我們建議在節(jié)點(diǎn)內(nèi)使用張量并行，在節(jié)點(diǎn)間使用流水并行，因?yàn)閺埩坎⑿行枰嗟?NCCL 通信。

多框架：FasterTransformer除了C上的源代碼，還提供了TensorFlow op、PyTorch op和Triton backend。目前TensorFlow op只支持單GPU，而PyTorch op和Triton backend支持多GPU和多節(jié)點(diǎn)。FasterTransformer 還提供了一個(gè)工具，可以將 Megatron 的模型拆分并轉(zhuǎn)換為二進(jìn)制文件，然后 FasterTransformer 可以直接加載二進(jìn)制文件，從而避免為模型并行而進(jìn)行額外的模型拆分工作。

Fig 3. Workflow of GPT with tensor parallelism.

筆記

is_context_qk_buf_float_（是否對(duì) GPT context QK GEMM 使用浮點(diǎn)累加）默認(rèn)設(shè)置為 false。如果您遇到與 GPT Context注意力模塊相關(guān)的準(zhǔn)確性問(wèn)題，請(qǐng)嘗試在 ParallelGpt.h 中啟用它。

CUDA相關(guān)優(yōu)化點(diǎn)解讀

TensorRT fused multi-head attention kernel: 和 BERT 一樣對(duì)于 GPT 的 ContextSelfAttention，F(xiàn)asterTransformer 使用 TensorRT 的 fused multi-head attention kernel 將 batch GEMM，softmax, GEMM，transpose 等操作都合并成一個(gè) cuda kernel，不僅可以減少數(shù)據(jù)搬提升帶寬利用率還可以減少 kernel launch 的開(kāi)銷(xiāo)。

AddBiasResidualLayerNorm：在 Decoder 中將 Attention 的最后一個(gè) Linear 的 bias_add，殘差連接（elementwise_add）以及 LayerNorm 合并成一個(gè) AddBiasResidualLayerNorm Kernel，降低 Kernel Launch 開(kāi)銷(xiāo)以及提升訪問(wèn)帶寬。

GeluFFN：從Fig1中的示意圖可以看到 GeluFFN 包含兩個(gè) Linear 層，中間夾了一個(gè) GeLU 的激活函數(shù)，這里做的優(yōu)化是把第一個(gè) Linear 層的 bias_add 和 GeLU 激活函數(shù) fuse 到一起，也就是 AddBiasGeLU Kernel。

AddBiasResidual：從Fig1的示意圖可以看到，Decoder的最后一層就是 AddBiasResidual，這個(gè)Kernel就是把 bias_add 和 殘差連接(element_wise add) 融合到一起。

GEMM 試跑：和BERT一樣仍然是在運(yùn)行模型之前先試跑一下 GPT 網(wǎng)絡(luò)中涉及到的GEMM的尺寸，并且保存 GEMM 性能最高的超參數(shù)配置，這個(gè)對(duì)于 cublas 和 cutlass 實(shí)現(xiàn)的卷積應(yīng)該都是成立的。

高效的 LayerNorm：在 TensorFlow 里面 LayerNorm 是零碎的 Kernel 拼接的，在 FasterTransformer 中實(shí)現(xiàn)了一個(gè) LayerNorm Kernel 來(lái)完成這個(gè)功能。實(shí)際上 PyTorch/OneFlow 等框架也有 LayerNorm Kernel，并且 OneFlow 的 LayerNorm 性能最強(qiáng)。

GEMM 的 FP16 累加：上面提到 is_context_qk_buf_float_ 參數(shù)，在 GPT 的 fp8 實(shí)現(xiàn)中，默認(rèn)使用 GEMM 的 FP16 累加而非 FP32 累加，進(jìn)一步提升性能，但是也可能帶來(lái)精度問(wèn)題。

最近我做一個(gè)大模型的推理工作時(shí)也發(fā)現(xiàn)如果基于 cutlass 的 gemm 使用 FP16 累加，最后生成的結(jié)果會(huì)部分亂碼，所以這個(gè)優(yōu)化必須用環(huán)境變量或者類(lèi)似于這里用一個(gè)單獨(dú)的參數(shù)來(lái)控制。

和通信相關(guān)的實(shí)現(xiàn)以及shared_context相關(guān)的優(yōu)化這里就不提了，代碼的可讀性比較差，我個(gè)人建議有需要的讀者學(xué)習(xí)下某些kernel的實(shí)現(xiàn)即可。

FasterTransformer 優(yōu)點(diǎn)

從之前對(duì) BERT 的優(yōu)化點(diǎn)介紹以及這里對(duì) GPT 的優(yōu)化點(diǎn)介紹，我們可以發(fā)現(xiàn)FasterTransformer集成了大量針對(duì)Transformer架構(gòu)的優(yōu)化，并且實(shí)現(xiàn)了各種Transformer架構(gòu)中常見(jiàn)的各種fuse pattern對(duì)應(yīng)的kernel。并且較為完整的支持了Transformer架構(gòu)的int8推理，整體的性能始終保持在一個(gè)SOTA水平。對(duì)于我這種入門(mén)CUDA優(yōu)化的學(xué)習(xí)者來(lái)說(shuō)有一定的學(xué)習(xí)意義。此外，F(xiàn)asterTransformer也將實(shí)現(xiàn)的這些組件注冊(cè)到TensorFlow，PyTorch等框架中使得讀者可以對(duì)常見(jiàn)的Transformer架構(gòu)進(jìn)行推理。

FasterTransformer 缺點(diǎn)

CUDA Kernel之外的代碼寫(xiě)得很抽象，特別對(duì)于多卡模式來(lái)說(shuō)需要用戶(hù)手動(dòng)管理通信和模型切分，這個(gè)門(mén)檻是很高的。如果用戶(hù)想基于FasterTreansformer這個(gè)框架實(shí)現(xiàn)新的Transformer架構(gòu)的網(wǎng)絡(luò)會(huì)非常困難，必須要非常了解FasterTransformer才可以。除了要手動(dòng)管理通信以及模型切分之外，如果用戶(hù)的新模型中出現(xiàn)了新的組件不僅要實(shí)現(xiàn)CUDA Kernel還需要手動(dòng)管理內(nèi)存的申請(qǐng)和釋放，比如GPT的內(nèi)存申請(qǐng)和釋放：

最近試跑了一個(gè)第三方模型的FasterTransformer實(shí)現(xiàn)，就出現(xiàn)了類(lèi)似的問(wèn)題。

個(gè)人認(rèn)為 FasterTransformer 的整體架構(gòu)實(shí)現(xiàn)的用戶(hù)體驗(yàn)類(lèi)似于 “九轉(zhuǎn)大腸”，易用性方面我還是比較看好 PyTorch ，OneFlow 等將內(nèi)存管理，通信集成到框架底層用戶(hù)新增模型只需要關(guān)心自定義 CUDA Kernel 的傳統(tǒng)深度學(xué)習(xí)框架。個(gè)人建議可以學(xué)習(xí)下 FasterTransformer 某些 CUDA Kernel 實(shí)現(xiàn)，但基于這個(gè)框架來(lái)搭建應(yīng)用要慎重。如果基于 PyTorch，OneFlow 等框架能將大量的 Transformer 架構(gòu)性能追平甚至超越 FasterTransformer 就完全沒(méi)必要折磨自己。

總結(jié)

這里總結(jié)了一下 FasterTransformer 里面和 CUDA Kernel相關(guān)的優(yōu)化技巧，并且給出了Kernel實(shí)現(xiàn)的位置，并從易用性，性能多方便對(duì)比了 FasterTransformer 和 PyTorch/OneFlow 等框架的優(yōu)缺點(diǎn)，供大家參考學(xué)習(xí)。

審核編輯：劉清