就是快！BetterTransformer：實(shí)現(xiàn)Transformer的快速推理
點(diǎn)藍(lán)色字關(guān)注“機(jī)器學(xué)習(xí)算法工程師”
設(shè)為星標(biāo)，干貨直達(dá)！
Transformers 在 NLP 方面實(shí)現(xiàn)了最先進(jìn)的性能，并且在無數(shù)其他任務(wù)中變得流行。它們的計算成本很高，這一直是它們廣泛生產(chǎn)的障礙。從 PyTorch 1.12 開始，BetterTransformer 為 Transformer Encoder Inference 實(shí)現(xiàn)了一個向后兼容的 torch.nn.TransformerEncoder 快速路徑，并且不需要模型作者修改他們的模型。對于許多常見的執(zhí)行場景，BetterTransformer 的改進(jìn)速度和吞吐量可以超過 2 倍。要使用 BetterTransformer，請安裝 PyTorch 1.12 并立即開始使用帶有 PyTorch API 的高質(zhì)量、高性能 Transformer 模型。
在這篇博文中，我們分享了以下主題——性能改進(jìn)、向后兼容性和利用 FastPath。在下面了解有關(guān)這些主題的更多信息。
性能改進(jìn)
BetterTransformer 在 CPU 和 GPU 上使用MultiHeadAttention 和 TransformerEncoderLayer的加速原生實(shí)現(xiàn)來啟動。這些快速路徑集成在標(biāo)準(zhǔn) PyTorch Transformer API 中，并將加速 TransformerEncoder、TransformerEncoderLayer 和 MultiHeadAttention這些nn模塊。這些新模塊實(shí)現(xiàn)了兩種類型的優(yōu)化：(1) 融合內(nèi)核結(jié)合了通常用于實(shí)現(xiàn) Transformer 的多個單獨(dú)的運(yùn)算符以提供更有效的實(shí)現(xiàn)，以及 (2) 利用輸入中的稀疏性來避免對填充令牌執(zhí)行不必要的操作。在用于自然語言處理的許多 Transformer 模型中，填充標(biāo)記通常占輸入批次的很大一部分。
向后兼容
有利的是，無需更改模型即可從 BetterTransformer 提供的性能提升中受益。為了從快速路徑執(zhí)行中受益，輸入和操作條件必須滿足一些訪問條件（見下文）。雖然 Transformer API 的內(nèi)部實(shí)現(xiàn)發(fā)生了變化，但 PyTorch 1.12 保持與先前版本中的 Transformer 模塊的嚴(yán)格兼容性，使 PyTorch 用戶能夠使用使用先前 PyTorch 版本創(chuàng)建和訓(xùn)練的模型，同時受益于 BetterTransformer 的改進(jìn)。
除了啟用 PyTorch nn.Modules 之外，BetterTransformer 還為 PyTorch 庫提供了改進(jìn)。性能優(yōu)勢將通過兩種不同的啟用路徑獲得：
透明加速：PyTorch nn.Modules 的當(dāng)前用戶（例如 MultiHeadAttention 以及更高級別的 Transformer 組件）將自動受益于新 nn.Modules 改進(jìn)的性能。torchvision 庫中使用的視覺轉(zhuǎn)換器 (ViT) 實(shí)現(xiàn)就是一個例子（代碼鏈接：https://github.com/pytorch/vision/blob/main/torchvision/models/vision_transformer.py#L103）。
Torchtext 庫加速：作為該項(xiàng)目的一部分，我們優(yōu)化了 Torchtext 以構(gòu)建在 PyTorch 核心 API 上，以從 BetterTransformer 增強(qiáng)功能中受益，同時保持與以前的庫版本和使用以前的 Torchtext 版本訓(xùn)練的模型的嚴(yán)格和透明的兼容性。在 Torchtext 中使用 PyTorch Transformer 還可以確保 Torchtext 將從 PyTorch Transformer 實(shí)現(xiàn)的預(yù)期未來增強(qiáng)中受益。
利用快速路徑
BetterTransformer 是 PyTorch Transformer API 的快速路徑?？焖俾窂绞沁m用于常見 Transformer 用例的 CPU 和 GPU 關(guān)鍵 Transformer 功能的本機(jī)專用實(shí)現(xiàn)。
要利用輸入稀疏性（即填充）來加速模型（參見圖 2），請在實(shí)例化 TransformerEncoder 時設(shè)置關(guān)鍵字參數(shù) enable_nested_tensor=True 并在推理期間傳入 src_key_padding_mask 參數(shù)（表示填充標(biāo)記）。這要求填充掩碼是連續(xù)的，這是典型的情況。
目前，BetterTransformer 加速僅適用于推理中使用的變壓器編碼器模型。要從快速路徑執(zhí)行中受益，模型必須由以下任何組件組成：TransformerEncoder、TransformerEncoderLayer 或 MultiheadAttention (MHA)?？焖俾窂綀?zhí)行也受制于一些標(biāo)準(zhǔn)。最重要的是，模型必須在推理模式下執(zhí)行，并在不收集梯度磁帶信息的輸入張量上運(yùn)行（例如，使用 torch.no_grad 運(yùn)行）。可以分別在 nn.MultiHeadAttention 和 nn.TransformerEncoder 的這些鏈接中找到完整的條件列表。如果不滿足標(biāo)準(zhǔn)，控制將流向具有相同 API 但缺乏快速路徑性能提升的舊版 PyTorch 1.11 Transformer 實(shí)現(xiàn)。
使用 PyTorch MultiheadAttention 模塊的其他轉(zhuǎn)換器模型（例如解碼器模型）將受益于 BetterTransformer 快速路徑。計劃的未來工作是將端到端 BetterTransformer 快速路徑擴(kuò)展到基于 TransformerDecoder 的模型，以支持流行的 seq2seq 和僅解碼器（例如 OPT）模型架構(gòu)，并用于訓(xùn)練。
加速效果
下圖顯示了具有小規(guī)模和大規(guī)模輸入的基于 BERT 的模型所取得的性能：
BetterTransformer 包括兩種類型的優(yōu)化：(1) 融合內(nèi)核在單個內(nèi)核中更有效地實(shí)現(xiàn)多個操作，以及 (2) 通過避免對填充令牌進(jìn)行不必要的處理來利用稀疏性。小輸入大小的增強(qiáng)性能主要受益于融合內(nèi)核實(shí)現(xiàn)，并且無論填充量如何都顯示出持續(xù)的性能改進(jìn)。雖然大輸入仍然受益于融合內(nèi)核，但計算繁重的處理限制了融合內(nèi)核可能獲得的好處，因?yàn)榛€性能已經(jīng)接近理論峰值。然而，隨著我們增加填充量，性能會顯著提高，因?yàn)榭梢酝ㄟ^利用 NLP 工作負(fù)載中填充引入的稀疏性來避免越來越多的計算量。
未來規(guī)劃
作為我們正在進(jìn)行的 PyTorch BetterTransformer 工作的一部分，我們正在努力將 BetterTransformer 改進(jìn)擴(kuò)展到 Transformer 解碼器。我們的目標(biāo)是從推理擴(kuò)展到訓(xùn)練。
我們正在合作在 FairSeq、MetaSeq 和 HuggingFace 等其他庫上啟用 BetterTransformer，以使所有基于 Transformer 的 PyTorch 模型受益。作為本博客系列的一部分，我們將在更大的 PyTorch 生態(tài)系統(tǒng)中提供有關(guān) BetterTransformer 加速進(jìn)展的未來更新。
本文翻譯自：https://pytorch.org/blog/a-better-transformer-for-fast-transformer-encoder-inference/
推薦閱讀
輔助模塊加速收斂，精度大幅提升！移動端實(shí)時的NanoDet-Plus來了！
SSD的torchvision版本實(shí)現(xiàn)詳解
機(jī)器學(xué)習(xí)算法工程師
一個用心的公眾號
