螞蟻鏈AIoT團(tuán)隊(duì)與NVIDIA合作加速AI推理

螞蟻鏈 AIoT 團(tuán)隊(duì)與 NVIDIA 合作，將量化感知訓(xùn)練(QAT)技術(shù)應(yīng)用于深度學(xué)習(xí)模型性能優(yōu)化中，并通過 NVIDIA TensorRT 高性能推理 SDK 進(jìn)行高效率部署， 通過 INT8 推理， 吞吐量提升了 3 倍， 助力螞蟻鏈版權(quán) AI 平臺(tái)中的模型推理服務(wù)大幅降本增效。

依托于螞蟻鏈自研的區(qū)塊鏈和 AI 技術(shù)，以及金融級的加密算法能力和云計(jì)算能力的加持，螞蟻鏈版權(quán) AI 平臺(tái)能夠?yàn)閿?shù)字媒體時(shí)代的內(nèi)容管理方、內(nèi)容平臺(tái)方以及創(chuàng)作者提供一站式的版權(quán)保護(hù)能力。作為螞蟻鏈版權(quán)平臺(tái)的核心能力，AI 技術(shù)可以快速提取音視頻及圖像特征，在短時(shí)間內(nèi)完成對相似內(nèi)容的自動(dòng)識別和判斷。得益于 NVIDIA 完善的軟硬件生態(tài)，螞蟻鏈團(tuán)隊(duì)除了可以方便的在云端 GPU 部署深度學(xué)習(xí)模型推理服務(wù)，在進(jìn)一步的合作中，雙方亦成功將其 TensorRT 支持的業(yè)界領(lǐng)先的 INT8-QAT 技術(shù)應(yīng)用到版權(quán)保護(hù)業(yè)務(wù)模型上，對 AI 模型的推理效率進(jìn)行大幅度優(yōu)化，實(shí)現(xiàn)了在幾乎精度無損的情況下，單 GPU 上的吞吐量提升約 300%。

利用深度學(xué)習(xí)模型高效且自動(dòng)化地進(jìn)行相似內(nèi)容識別是螞蟻鏈版權(quán) AI 平臺(tái)的一項(xiàng)關(guān)鍵技術(shù)能力，這對 AI 模型研發(fā)提出了較高的技術(shù)挑戰(zhàn)。一方面，多媒體版權(quán)保護(hù)業(yè)務(wù)場景中作品侵權(quán)的類型復(fù)雜而多變， 為了到達(dá)高召回率和低虛警率的效果精度，螞蟻鏈團(tuán)隊(duì)需采用較為復(fù)雜的算法模型方案;而另一方面，版權(quán)場景中高吞吐、低成本的需求又對模型提出了性能方面的挑戰(zhàn)。換言之，算法模型需要的存算資源必須有所限制、推理性能必須較高，而一般情況下，這與模型的效果要求是矛盾的，因?yàn)樯疃葘W(xué)習(xí)模型的效果往往與模型的尺寸和所需算力正相關(guān)。而這些復(fù)雜而又多維的業(yè)務(wù)需求，也確實(shí)給團(tuán)隊(duì)的算法研發(fā)和優(yōu)化工作，帶來了相當(dāng)大的挑戰(zhàn)。

因此，螞蟻鏈團(tuán)隊(duì)將算法研發(fā)拆解為不同的流程以解決不同維度的業(yè)務(wù)需求。簡而言之，就是優(yōu)先以滿足業(yè)務(wù)效果指標(biāo)為目的研發(fā)模型，隨后借助模型壓縮技術(shù)對模型存算需求進(jìn)行優(yōu)化，而第二部分工作實(shí)現(xiàn)，則需仰賴 NVIDIA 相關(guān)軟硬件生態(tài)對于高性能AI模型推理加速的優(yōu)秀支持。一方面，NVIDIA GPU 所提供的強(qiáng)大并行算力以及 INT8 Tensor Core 提供的整型計(jì)算能力，為深度學(xué)習(xí)模型推理的高效實(shí)現(xiàn)奠定了基礎(chǔ);另一方面，NVIDIA TensorRT 8 SDK 中對模型計(jì)算圖的高效融合，以及對于新型模型結(jié)構(gòu)的支持和優(yōu)化(例如 QAT 所采用的 QDQ 結(jié)構(gòu)以及 transformer-based 模型)讓我們可以在 PyTorch 模型的基礎(chǔ)上生成高效的量化推理模型。其中，QAT 作為此項(xiàng)目模型性能優(yōu)化的重要技術(shù)，最大的優(yōu)點(diǎn)莫過于可以在大幅提升推理速度的同時(shí)，做到幾乎完全消弭量化網(wǎng)絡(luò)整型計(jì)算相比浮點(diǎn)網(wǎng)絡(luò)數(shù)值精度差異帶來的算法效果損失。而近年來，隨著 NVIDIA 對相關(guān)軟硬件生態(tài)的對 QAT 技術(shù)的支持日臻完善，螞蟻鏈團(tuán)隊(duì)希望能藉由 NVIDIA 相關(guān)的軟硬件生態(tài)，將 QAT 技術(shù)應(yīng)用到螞蟻鏈版權(quán) AI 項(xiàng)目中，以滿足嚴(yán)苛的業(yè)務(wù)需求。

基于以上挑戰(zhàn)，螞蟻鏈 AIoT 技術(shù)部與 NVIDIA 技術(shù)專家合作，引入 QAT 技術(shù)對深度學(xué)習(xí)模型進(jìn)行性能優(yōu)化，并將模型轉(zhuǎn)換為 TensorRT 引擎部署至 NVIDIA T4 GPU 進(jìn)行在線推理。

首先，團(tuán)隊(duì)使用 QAT 技術(shù)對算法模型以“偽量化”模式進(jìn)行微調(diào)，使模型得以在 INT8 整型推理模式下在算法效果指標(biāo)上幾乎沒有損失。具體實(shí)現(xiàn)方案主要包括：使用 NVIDIA PyTorch Quantization 工具在模型中一些特定位置插入量化/反量化節(jié)點(diǎn)，在原模型的基礎(chǔ)上構(gòu)造一個(gè)“偽量化”模型。緊接著，使用該“偽量化”模型在原任務(wù)上進(jìn)行一定輪數(shù)的微調(diào)，使模型參數(shù)在受到量化誤差擾動(dòng)的同時(shí)依然可以收斂到一個(gè)局部最優(yōu)，最終最小化量化推理對模型算法指標(biāo)帶來的負(fù)面影響。實(shí)踐中，螞蟻鏈團(tuán)隊(duì)通過 QAT 技術(shù)可以將 INT8 模型與單精度模型的算法指標(biāo)的相對差距縮小到 0.2% 以內(nèi)，做到了算法效果幾乎無損的模型量化。

隨后，則需要將微調(diào)完成的 QAT 模型轉(zhuǎn)換為真正的 INT8 版本模型部署到線上生產(chǎn)環(huán)境中。這一流程主要依賴 NVIDIA 提供的 TensorRT 高性能推理框架，相較其他部署框架，TensorRT 在 NVIDIA GPU 上表現(xiàn)出巨大的性能優(yōu)勢：一方面通過 graph fusion，kernel tuning 等功能，可以自動(dòng)化精簡網(wǎng)絡(luò)結(jié)構(gòu)，為模型各層的不同 op 尋找最優(yōu) CUDA kernel 等優(yōu)化操作;更重要的是，TensorRT 8.0 開始，添加了針對“偽量化”節(jié)點(diǎn)的自動(dòng)化解析和融合功能，使 QAT 模型落地的門檻大大降低，讓 QAT 模型真正有可能在部署環(huán)境中展現(xiàn)出近似 INT8 后量化模型的極高推理性能。實(shí)踐中，采用 INT8-QAT 的模型做到了比單精度模型約 300% 的單位時(shí)間吞吐提升，極大提高了模型推理服務(wù)的效率;同時(shí)，量化模型更低的顯存占用也為模型部署帶來了更高的靈活性。

借助 NVIDIA 在高性能模型推理方面完善的軟硬件生態(tài)，螞蟻鏈團(tuán)隊(duì)得以使用 INT8-QAT 技術(shù)大幅優(yōu)化螞蟻鏈版權(quán) AI 平臺(tái)中模型推理服務(wù)的效率，幫助系統(tǒng)整體降本增效。300% 的推理速度提升，以及算法指標(biāo)幾乎無損的特性，讓此項(xiàng)目的 AI 模型可以做到兼顧效果與性能，幫助螞蟻鏈版權(quán) AI 平臺(tái)在業(yè)界樹立技術(shù)優(yōu)勢。

“版權(quán)保護(hù)是一個(gè)富有挑戰(zhàn)的技術(shù)領(lǐng)域，互聯(lián)網(wǎng)內(nèi)容類型多樣、隱匿性強(qiáng)、易復(fù)制與編輯等特點(diǎn)對我們提出了諸多挑戰(zhàn)。螞蟻鏈從 2019 年發(fā)布鵲鑿版權(quán)保護(hù)平臺(tái)以來，在音視圖文的侵權(quán)檢索、侵權(quán)比對與定位方面做了大量深入的研究工作。我們 AIoT 團(tuán)隊(duì)在 2022 年世界知識產(chǎn)權(quán)日向公眾發(fā)布了螞蟻鏈版權(quán) AI 計(jì)算引擎，可以對相似內(nèi)容的提取精確到幀，直接以秒為單位反饋比對結(jié)果。這種細(xì)顆粒度的識別能力極大提高了下游工作效率，同時(shí)也需要我們充分探索和利用 NVIDIA INT8-QAT 等加速技術(shù)，在效果和性能之間達(dá)到最佳平衡。接下去我們還會(huì)面向低成本高性能的版權(quán) AI 算法演進(jìn)，讓技術(shù)普惠更多的普通創(chuàng)作者?！蔽浵佹?AIoT 高級算法專家張曉博與錢烽表示。

審核編輯：湯梓紅