近年來(lái)，Transformer 已經(jīng)成為了 NLP 和 CV 等領(lǐng)域的主流模型，但龐大的模型參數(shù)限制了它的高效訓(xùn)練和推理。于是字節(jié)跳動(dòng)在 2019 年 12 月和 2021 年 6 月分別推出了高效推理和訓(xùn)練引擎 LightSeq，大大加速了 Transformer 系列模型的訓(xùn)練和推理，也打通了 Transformer 從訓(xùn)練到推理的整個(gè)流程，極大優(yōu)化了用戶使用體驗(yàn)。最近，LightSeq 訓(xùn)練引擎相關(guān)論文[1]，被錄用難度極高的超算領(lǐng)域國(guó)際頂會(huì) SC22 接收，得到了學(xué)術(shù)界的廣泛認(rèn)可！

SC22 接收論文：https://sc22.supercomputing.org/presentation/?id=pap211&sess=sess154

代碼地址：https://github.com/bytedance/lightseq

如何繼續(xù)提升速度？降低計(jì)算精度是比較直接的方法。2017 年以來(lái)，fp16 混合精度技術(shù) [2] 獲得了廣泛應(yīng)用。在對(duì)模型效果無(wú)損的前提下，將模型訓(xùn)練和推理的速度提升了 50% 以上。而為了維持模型效果，更低精度的方法（例如 int8）通常需要使用如下傳統(tǒng)方案：

首先使用 fp16 混合精度將模型訓(xùn)練至收斂；

然后在模型計(jì)算密集型算子的權(quán)重、輸入和輸出位置處，插入偽量化結(jié)點(diǎn)，進(jìn)行量化感知訓(xùn)練；

最后將帶有偽量化結(jié)點(diǎn)的模型計(jì)算圖轉(zhuǎn)換到專用的 int8 推理引擎中，進(jìn)行服務(wù)部署和模型推理。

雖然在多數(shù)任務(wù)上，上述方案可以實(shí)現(xiàn)模型效果無(wú)損，但還是存在以下問(wèn)題：

使用方法復(fù)雜。例如要多一次量化感知訓(xùn)練 [4] 的過(guò)程，并且?guī)в袀瘟炕?jié)點(diǎn)的計(jì)算圖轉(zhuǎn)換復(fù)雜。

訓(xùn)練速度慢。由于目前流行的深度學(xué)習(xí)框架不支持 int8 精度，所以量化感知訓(xùn)練需要插入 fp16 的偽量化結(jié)點(diǎn)來(lái)模擬 int8 量化，導(dǎo)致量化感知訓(xùn)練反而比 fp16 混合精度訓(xùn)練慢 2-3 倍。

推理部署難且加速比低。對(duì)比 fp32、fp16 等類型，int8 硬件和底層軟件庫(kù)優(yōu)化相對(duì)滯后。例如在 NVIDIA GPU 上，int8 矩陣乘法加速受限于硬件架構(gòu)和特定 shape，實(shí)際加速比遠(yuǎn)遠(yuǎn)低于理論值。

在下文中，如無(wú)特殊說(shuō)明，量化都是指的 int8 精度的量化。

針對(duì)這些問(wèn)題，字節(jié)跳動(dòng)推出了全新版本的 LightSeq GPU 量化訓(xùn)練與推理引擎。支持 Transformer 系列模型的量化訓(xùn)練與推理，并做到了開(kāi)箱即用，用戶友好。LightSeq 快準(zhǔn)狠地實(shí)現(xiàn)了 int8 精度的量化訓(xùn)練和推理：

快：A100 多卡訓(xùn)練最高加速 5.2 倍，T4 單卡推理最高加速 8.9 倍。

準(zhǔn)：訓(xùn)練和推理效果基本無(wú)損。

狠：相同數(shù)據(jù)量下，顯存占用最高減少 68%，模型存儲(chǔ)空間減少 75%。

總體來(lái)說(shuō)，LightSeq 新版量化訓(xùn)練與推理引擎具有如下幾個(gè)優(yōu)點(diǎn)：

1.?豐富的支持

支持完整的 Transformer 模塊和多種解碼算法，支持 Transformer、BERT、GPT、BART、ViT 等多種模型結(jié)構(gòu)，支持 Fairseq、Hugging Face、NeurST 等多種訓(xùn)練框架接入量化訓(xùn)練、導(dǎo)出模型以及量化推理，提供了豐富的樣例供用戶參考。

2.?卓越的性能

相比于 fp16 精度的 LightSeq 推理引擎，int8 量化還可以進(jìn)一步加速最高 70%，相比于 PyTorch 推理更是達(dá)到了最高 8.9 倍的加速比。同時(shí)顯存占用相比 fp16 推理引擎降低了 30% 左右，模型存儲(chǔ)空間只需要原來(lái)的四分之一。最后經(jīng)過(guò)多個(gè)任務(wù)的驗(yàn)證，推理效果幾乎無(wú)損。

3. 便捷的使用

LightSeq 已經(jīng)針對(duì)多個(gè)訓(xùn)練庫(kù)進(jìn)行了量化支持，可以一鍵開(kāi)啟量化訓(xùn)練，然后輕松導(dǎo)出為 LightSeq 支持的模型格式，最后實(shí)現(xiàn)量化推理。除此之外，LightSeq 還支持訓(xùn)練后量化，無(wú)需額外訓(xùn)練即可體驗(yàn)量化推理。

使用方法

如上圖所示，為了最大程度減小量化帶來(lái)的損失，首先需要用 fp16 精度訓(xùn)練一個(gè)浮點(diǎn)數(shù)模型，將模型效果訓(xùn)到最好。然后開(kāi)啟量化進(jìn)行 finetune，得到微調(diào)過(guò)的量化模型，此時(shí)模型效果已經(jīng)基本恢復(fù)到浮點(diǎn)數(shù)模型的水平。接著將量化模型轉(zhuǎn)換為 LightSeq 支持的 PB 或者 HDF5 模型格式，最后用 LightSeq 進(jìn)行量化推理。

安裝方法

LightSeq 安裝非常簡(jiǎn)單，只需要一行命令即可：

pip install lightseq

量化訓(xùn)練

LightSeq 支持 Fairseq、Hugging Face、NeurST 等訓(xùn)練框架的量化接入，同時(shí)也可以自定義模型并開(kāi)啟量化訓(xùn)練。以 encoder 層為例，只需要先定義浮點(diǎn)數(shù)模型，然后開(kāi)啟量化即可：

from lightseq.training import LSTransformerEncoderLayer
from lightseq.training.ops.pytorch.quantization import enable_quant

config = LSTransformerEncoderLayer.get_config(
    model="bert-base",
    max_batch_tokens=4096,
    max_seq_len=512,
    fp16=True,
    local_rank=0,
)
layer = LSTransformerEncoderLayer(config)
# 開(kāi)啟量化
layer.apply(enable_quant)

量化推理

LightSeq 提供了便捷的 python 推理接口，只需要三行代碼即可實(shí)現(xiàn)快速的量化推理：

import lightseq.inference as lsi

model = lsi.QuantTransformer(pb_path, batch_size)
result = model.infer(input)

此外 LightSeq 還提供了 BERT、GPT、ViT 等模型的 python 接口，分別調(diào)用 QuantBert、QuantGpt 和 QuanVit 即可體驗(yàn)。

梯度通信量化

LightSeq 支持 Transformer 模型的梯度通信量化[5]，使用 Fairseq 或者 Hugging Face 即可輕松開(kāi)啟分布式量化訓(xùn)練，并同時(shí)支持浮點(diǎn)數(shù)模型和量化模型。在構(gòu)建模型后，只需要為模型注冊(cè)一個(gè) communication hook 即可開(kāi)啟梯度通信量化，再開(kāi)始訓(xùn)練過(guò)程。

from lightseq.training.gradient_comm_quantization import encode_and_decode, GCQState
from torch.nn.parallel import DistributedDataParallel 


# model could be from Fairseq or Hugging Face, wrapped by DDP
model = DistributedDataParallel(model)
state =  GCQState(process_group)
# register hook
model.register_comm_hook(state=state, hook=encode_and_decode)

性能測(cè)試

LightSeq 在多個(gè)任務(wù)上測(cè)試了量化訓(xùn)練、量化推理和梯度通信量化的速度，并且分析了顯存占用情況和量化模型的效果。

量化訓(xùn)練速度

LightSeq 在 8 張 A100 顯卡上進(jìn)行了訓(xùn)練實(shí)驗(yàn)，主要對(duì)比對(duì)象是 Fairseq 的 Transformer、Hugging Face 的 BERT、GPT2 和 ViT。

可以看出，四種模型結(jié)構(gòu)加速趨勢(shì)都是類似的，加速比都會(huì)隨著數(shù)據(jù)量的增大而減小，原因有三點(diǎn)：

隨著數(shù)據(jù)量的增大，矩陣乘法 GEMM 的占比會(huì)明顯增加，因此 PyTorch QAT 增加的額外的偽量化結(jié)點(diǎn)時(shí)間占比會(huì)逐漸減小，最后速度會(huì)和 PyTorch fp16 無(wú)限接近。

與此同時(shí)，隨著 GEMM 占比升高，LightSeq fp16 自定義算子的提速效果也逐漸減小，因此時(shí)間上也會(huì)和 PyTorch fp16 無(wú)限接近。

由于 Ampere 架構(gòu)顯卡上 int8 GEMM 在 shape 較小時(shí)甚至不如 fp16 GEMM 快，在大 shape 下才能稍快一點(diǎn)，因此隨著數(shù)據(jù)量增大，LightSeq int8 也會(huì)無(wú)限接近 LightSeq fp16 的速度。

量化推理速度

LightSeq 在單張 T4 顯卡上進(jìn)行了推理實(shí)驗(yàn)，主要對(duì)比對(duì)象是 Hugging Face 的 Transformer、BERT、GPT2 和 ViT。

可以看出，隨著輸入數(shù)據(jù)量的增大，LightSeq 與 PyTorch 的差距會(huì)逐漸減小，這也是 GEMM 占比升高造成的。比較 LightSeq fp16 和 LightSeq int8，可以看出隨著數(shù)據(jù)量的增大，LightSeq int8 越來(lái)越快。這是因?yàn)樵?T4 顯卡上，int8 GEMM 的加速會(huì)隨著 shape 的增大而有明顯增加。因此在 T4 顯卡上進(jìn)行量化推理時(shí)，輸入數(shù)據(jù)量越大，加速效果越好。

LightSeq 還針對(duì)機(jī)器翻譯多個(gè)語(yǔ)向和多個(gè)測(cè)試集，測(cè)試了不同 batch size 下，LightSeq int8 推理相對(duì)于 LightSeq fp16 推理的加速比，實(shí)驗(yàn)同樣是在單張 T4 顯卡上進(jìn)行的，采用的模型都是標(biāo)準(zhǔn)的 Transformer-Big。

可以得到和上文中相同的結(jié)論，隨著 batch size 的增大，量化推理的加速比會(huì)逐漸升高。相比于 LightSeq fp16，最高還可以再加速近 70%，這極大地縮短了線上翻譯模型的推理延時(shí)。

最后如上圖所示，為了展示自動(dòng) GEMM 調(diào)優(yōu)技術(shù)的效果，LightSeq 測(cè)試對(duì)比了 A100 顯卡上 Transformer 和 BERT 模型 fp16、int8 調(diào)優(yōu)前和 int8 調(diào)優(yōu)后的延時(shí)?？梢钥闯稣{(diào)優(yōu)前某些 shape 的 int8 GEMM 速度甚至比 fp16 還要慢，而調(diào)優(yōu)后全面超越了 fp16。

顯存占用

LightSeq 分析了不同 batch size 下，量化模型相對(duì)于浮點(diǎn)數(shù)模型顯存占用的加速比?？梢钥闯鲭S著 batch size 的增大，量化模型的顯存占用優(yōu)勢(shì)更明顯，最高可以減少 30% 左右。而 LightSeq fp16 引擎相對(duì)于 PyTorch 模型也極大程度減少了顯存占用，因此 LightSeq int8 引擎最終能夠減少最多 68% 左右的顯存。

量化模型效果

針對(duì)機(jī)器翻譯多個(gè)語(yǔ)向和多個(gè)測(cè)試集，LightSeq 測(cè)試了量化模型推理相對(duì)于浮點(diǎn)數(shù)模型 BLEU 的損失，采用的模型都是標(biāo)準(zhǔn)的 Transformer-Big。

在數(shù)據(jù)量較大的語(yǔ)向 en2zh 上，LightSeq int8 相對(duì) BLEU 損失較大些，最大達(dá)到了 - 0.4。而在數(shù)據(jù)量較小的語(yǔ)向 en2es 上，LightSeq int8 不僅沒(méi)有任何效果損失，反而比浮點(diǎn)數(shù)模型更好?？傮w而言，int8 量化模型的平均 BLEU 相比浮點(diǎn)數(shù)模型基本無(wú)損。在 GLUE 和 SQuAD 等多個(gè)任務(wù)上，LightSeq 也驗(yàn)證了量化模型的效果。

梯度通信量化

由于在多機(jī)多卡場(chǎng)景下通信瓶頸更加明顯，所以梯度通信量化主要應(yīng)用在分布式訓(xùn)練場(chǎng)景。因此 LightSeq 在 2 機(jī) 8 卡的 A100 上進(jìn)行了分布式訓(xùn)練的速度測(cè)試。

可以看出，梯度通信量化的訓(xùn)練加速效果整體上隨著輸入數(shù)據(jù)的增大而減弱。這主要是因?yàn)殡S著輸入數(shù)據(jù)的增大，計(jì)算時(shí)間占比升高，梯度通信時(shí)間占比減少，梯度量化的收益也隨之減小。

LightSeq 還額外增加了不同數(shù)量網(wǎng)卡（NIC）下的訓(xùn)練速度測(cè)試?？梢钥吹绞褂锰荻韧ㄐ帕炕姆植际接?xùn)練速度相比原始的 LightSeq fp16 有大幅度提升。

量化技術(shù)

int8 量化的加速收益主要來(lái)自如下幾個(gè)方面：

GEMM 精度從 fp16 降低到 int8 后，計(jì)算時(shí)間縮短；

自定義算子采用 int8 輸入輸出后，數(shù)據(jù)讀寫(xiě)時(shí)間縮短；

梯度采用 int8 存儲(chǔ)后，多機(jī)之間通信時(shí)間縮短。

以 Transformer 模型為例，經(jīng)過(guò) LightSeq fp16 引擎加速后，自定義算子時(shí)間大大縮短，而 GEMM 時(shí)間占比提升到了 90% 左右，因此優(yōu)化的重點(diǎn)轉(zhuǎn)移到了 GEMM 提速。將 fp16 GEMM 替換為 int8 GEMM 不僅可以縮短 GEMM 時(shí)間，還可以減小前后算子的輸入輸出位寬，從而減小讀寫(xiě)數(shù)據(jù)的時(shí)間。最后多機(jī)訓(xùn)練的瓶頸主要在梯度的通信，將梯度量化為 int8 精度可以大大加快分布式訓(xùn)練的速度。

量化原理

為了彌補(bǔ)量化帶來(lái)的精度損失，通常需要用量化感知訓(xùn)練來(lái)模擬量化過(guò)程。如上圖所示，量化感知訓(xùn)練就是將 float GEMM 的兩個(gè) float 輸入分別做一遍量化和反量化（稱之為偽量化結(jié)點(diǎn)），離散化成分段的浮點(diǎn)數(shù)輸入，然后進(jìn)行 float GEMM 運(yùn)算。得到結(jié)果后再次進(jìn)行量化與反量化，得到最終的浮點(diǎn)數(shù)結(jié)果。而量化的過(guò)程是不可導(dǎo)的，因此需要用 STE 方法來(lái)估計(jì)量化參數(shù)的梯度。之所以量化感知訓(xùn)練中需要插入偽量化結(jié)點(diǎn)，然后用 float GEMM 去模擬量化過(guò)程，是因?yàn)?TensorFlow 和 PyTorch 等訓(xùn)練框架不支持 int8 GEMM。

而 LightSeq 量化訓(xùn)練直接采用 int8 GEMM 來(lái)真實(shí)還原量化過(guò)程，因此相比傳統(tǒng)的實(shí)現(xiàn)要更快，且更加節(jié)省顯存。在推理的時(shí)候，同樣采用離散化后的整數(shù)進(jìn)行 int8 GEMM 運(yùn)算，最后再反量化回浮點(diǎn)數(shù)結(jié)果。量化推理過(guò)程和量化訓(xùn)練完全一致，并且和傳統(tǒng)的量化感知訓(xùn)練是完全等價(jià)的。

量化位置

整個(gè)量化 Transformer 的網(wǎng)絡(luò)結(jié)構(gòu)如上圖所示，紅色箭頭表示需要加上量化和反量化結(jié)點(diǎn)的位置。

首先所有 int8 GEMM 的輸入和輸出都需要進(jìn)行量化。由于 int8 GEMM 的 shape 限制，部分 GEMM（例如注意力分?jǐn)?shù)的計(jì)算）仍然采用 float GEMM。此外第二層 FFN 的 GEMM 采用的是 int32 的輸出，因?yàn)樗?GEMM 輸入是 ReLU 激活函數(shù)的輸出結(jié)果，只包含正數(shù)，非對(duì)稱，因此如果采用 int8 輸出的 GEMM，將無(wú)法反量化為正確的浮點(diǎn)數(shù)結(jié)果。

然后所有的模型權(quán)重 weight 都需要存儲(chǔ)為 int8 類型，因此需要對(duì) weight 做量化。而權(quán)重 bias 參數(shù)量較小，無(wú)需量化，保留 float 精度反而可以提升模型效果。

最后需要對(duì) decoder 端的 cache 進(jìn)行量化。因?yàn)樵谕评頃r(shí)，decoder 端的 cache 需要頻繁進(jìn)行讀寫(xiě)，因此將 cache 量化為 int8 可以大大加快解碼的速度。

量化策略

將一個(gè)浮點(diǎn)數(shù)矩陣量化為 int8 整數(shù)矩陣有很多方法，LightSeq 采用的是對(duì)稱量化，即將正負(fù)數(shù)范圍對(duì)稱的浮點(diǎn)數(shù)區(qū)間等比例地映射到整數(shù)區(qū)間 [-127, 127] 上。

而實(shí)際上浮點(diǎn)數(shù)矩陣的數(shù)值范圍通常并不對(duì)稱，存在極少的離群值。如果直接按照離群值的范圍來(lái)量化矩陣，會(huì)影響到量化后的精度，所以需要先對(duì)矩陣進(jìn)行數(shù)值截?cái)唷?/p>

LightSeq 采用 PACT 方法進(jìn)行截?cái)郲6]，將截?cái)嗟姆秶?dāng)作模型可學(xué)習(xí)的參數(shù)，然后利用 STE 算法去估計(jì)參數(shù)的梯度，并進(jìn)行反向傳播優(yōu)化。根據(jù)實(shí)踐經(jīng)驗(yàn)，權(quán)重 weight 的初始截?cái)喾秶O(shè)為[-1, 1]，中間結(jié)果的初始截?cái)喾秶O(shè)為[-16, 16]，可以在大部分任務(wù)上達(dá)到最好的效果。最后經(jīng)過(guò)截?cái)喾秶推渌Ｐ蛥?shù)的聯(lián)合優(yōu)化，量化模型的效果可以達(dá)到基本無(wú)損。

梯度通信量化

針對(duì)分布式訓(xùn)練場(chǎng)景，LightSeq 推出了梯度量化壓縮技術(shù)。即對(duì)浮點(diǎn)精度的梯度進(jìn)行 int8 量化，以減少梯度通信的時(shí)間消耗，從而加速訓(xùn)練，這就是梯度通信量化（GCQ）。

如上圖所示，梯度通信量化的主要流程如下：

計(jì)算每張卡上各自梯度的截?cái)喾秶?/p>

對(duì)截?cái)喾秶鷪?zhí)行 all-reduce max 操作；

每張卡使用統(tǒng)一的截?cái)喾秶鷮?duì)各自梯度進(jìn)行 int8 量化；

對(duì) int8 梯度執(zhí)行 all-reduce sum 操作；

每張卡對(duì) all-reduce 后的梯度進(jìn)行反量化，還原為浮點(diǎn)數(shù)梯度，并進(jìn)行參數(shù)更新。

為了解決 int8 梯度在 all-reduce 過(guò)程中溢出的問(wèn)題，LightSeq 首先將每張卡上的浮點(diǎn)數(shù)梯度除以卡數(shù)，再使用除之前的截?cái)喾秶M(jìn)行量化，最后進(jìn)行 all-reduce 操作。這樣每張卡上量化后的 int8 整數(shù) all-reduce 完就不會(huì)溢出，但是單卡實(shí)際用于量化的比特?cái)?shù)也因此而減少，所以目前方案在 2 機(jī) 8 卡效果幾乎無(wú)損，但隨著卡數(shù)的上漲，訓(xùn)練效果會(huì)有所下降。以 en2de 和 en2fr 翻譯任務(wù)為例，在 4 機(jī) 8 卡上進(jìn)行分布式量化訓(xùn)練，BLEU 值分別會(huì)下降 0.4 和 1.5 左右。未來(lái) LightSeq 將會(huì)持續(xù)探索更好的方法來(lái)解決這一問(wèn)題。

通用技術(shù)

除了上一章節(jié)中提到的量化技術(shù)以外，此次更新 LightSeq 還提出了幾種通用的優(yōu)化技術(shù)，不僅可以應(yīng)用在量化模型中，也適用于其它所有精度模型的訓(xùn)練與推理。

算子融合

上圖是 encoder 模塊量化訓(xùn)練的計(jì)算圖，LightSeq 將兩次 GEMM 運(yùn)算之間的所有操作融合成一個(gè)算子[7]，減少了 kernel 調(diào)用的次數(shù)，因此減少了總的計(jì)算時(shí)間。

圖中黃色矩形表示 int8 GEMM，綠色矩形表示 float GEMM。這里采用 float GEMM 是由于 shape 的限制，不適合使用 int8 GEMM 加速。紅色箭頭表示流動(dòng)數(shù)據(jù)的類型是 int8，綠色箭頭表示第二層 FFN 的 GEMM 輸出是 int32 數(shù)據(jù)類型。int8 GEMM 輸入輸出的量化與反量化操作都被融合到了前后 kernel 里，這不僅可以減少數(shù)據(jù)搬運(yùn)，還可以減小顯存占用。

在推理時(shí)，LightSeq 還針對(duì) decoder 做了優(yōu)化。如上圖所示，在計(jì)算 self-attention 時(shí)，注意力得分的維度是(batch size, 1, sequence length)。因此在計(jì)算 value 乘積時(shí)，可以不采用 GEMM 運(yùn)算，而直接手寫(xiě)加權(quán)求和的算子，從而將圖中虛線框中的計(jì)算融合成一個(gè) kernel。

自動(dòng)顯存管理

模型量化引入了更復(fù)雜的張量類型和張量依賴關(guān)系，這給顯存管理帶來(lái)新的挑戰(zhàn)。為此，LightSeq 設(shè)計(jì)了新的顯存管理機(jī)制。如上圖所示，主要包括以下過(guò)程：

訓(xùn)練啟動(dòng)前，根據(jù)每個(gè)算子的拓?fù)湟蕾囮P(guān)系，自動(dòng)計(jì)算每個(gè)張量的生命周期及顯存空間大小。其中，包含動(dòng)態(tài)維度的張量按照此維度的最大量進(jìn)行計(jì)算，例如機(jī)器翻譯任務(wù)中的最大句長(zhǎng)和最大 batch 句子數(shù)量。這些最大量在訓(xùn)練前已被指定；

張量確定生命周期和大小后，分析顯存復(fù)用關(guān)系。其中，無(wú)生命周期重合的張量可以共用一片顯存空間，所有顯存空間都是無(wú)數(shù)據(jù)類型的，可以被分配到任意數(shù)據(jù)類型的張量上；

根據(jù)張量顯存復(fù)用關(guān)系，申請(qǐng)多段顯存空間，為每個(gè)張量分配實(shí)際的顯存起止地址。

張量顯存復(fù)用的分析，LightSeq 借鑒了論文 [3] 中提出的 Greedy by Size for Offset Calculation 方法，做了三個(gè)改進(jìn)：

支持了整個(gè)訓(xùn)練過(guò)程的顯存復(fù)用（forward/backward）；

不同數(shù)據(jù)類型能做到顯存復(fù)用（int8/fp16/fp32）；

在多段顯存空間上容納所有張量，而非一段非常大的顯存空間，這樣能有效提升顯存利用率。

自動(dòng) GEMM 調(diào)優(yōu)

LightSeq 的 int8 GEMM 采用了 NVIDIA 的 cuBLASLt 庫(kù)，這也是目前 NVIDIA 顯卡上最為高效的矩陣運(yùn)算庫(kù)。但是輸入數(shù)據(jù)的 shape 或者顯卡不同的話，GEMM 所采用的最優(yōu)配置（例如數(shù)據(jù)排布、GEMM 算法等等）也可能不同，因此需要進(jìn)行自動(dòng)選取。LightSeq 采取的自動(dòng)調(diào)優(yōu)方案如下：

在多種型號(hào)顯卡上（例如 T4 和 A100）進(jìn)行不同 shape 的 GEMM 最優(yōu)配置搜索，并將結(jié)果保存到配置文件中，用戶只需要下載即可；

模型初始化時(shí)，加載對(duì)應(yīng)型號(hào)顯卡的配置文件，解析并保存到鍵值對(duì)為 (shape, 最優(yōu)配置) 的字典中。如果沒(méi)有對(duì)應(yīng)型號(hào)顯卡的配置文件，或者沒(méi)有需要的 GEMM shape，那么用戶可以選擇自己搜索并保存，或者直接使用默認(rèn)配置；

模型前向或后向計(jì)算時(shí)，根據(jù)輸入的 shape 在字典中尋找最優(yōu)配置，然后進(jìn)行 GEMM 計(jì)算。如果沒(méi)有找到對(duì)應(yīng)的 shape，那么直接采用默認(rèn)的配置。

未來(lái)工作

未來(lái) LightSeq 還將繼續(xù)探索移動(dòng)端的低精度量化、反向傳播中梯度的量化、大模型量化等方向。

引用

[1] Wang, Xiaohui, et al. "LightSeq2: Accelerated training for transformer-based models on gpus." arXiv preprint arXiv:2110.05722 (2021).

[2] Micikevicius, Paulius, et al. "Mixed precision training." arXiv preprint arXiv:1710.03740 (2017).

[3] Pisarchyk, Yury, and Juhyun Lee. "Efficient memory management for deep neural net inference." arXiv preprint arXiv:2001.03288 (2020).

[4] Jacob, Benoit, et al. "Quantization and training of neural networks for efficient integer-arithmetic-only inference." Proceedings of the IEEE conference on computer vision and pattern recognition. 2018.

[5] Alistarh, Dan, et al. "QSGD: Communication-efficient SGD via gradient quantization and encoding." Advances in neural information processing systems 30 (2017).

[6] Choi, Jungwook, et al. "Pact: Parameterized clipping activation for quantized neural networks." arXiv preprint arXiv:1805.06085 (2018).

[7] Wang, Xiaohui, et al. "LightSeq: A high performance inference library for transformers." arXiv preprint arXiv:2010.13887 (2020).

編輯：黃飛

?