云廠(chǎng)商布局AIGC，火山引擎如何做到訓(xùn)練提升40%、推理提升3.47倍？
為AIGC客戶(hù)降本增效。
AIGC是當(dāng)前AI領(lǐng)域最熱門(mén)的話(huà)題之一。
尤其是伴隨著Stable Diffusion、Midjourney、NovelAI、DALL-E等為代表的文本生成圖像的跨模態(tài)應(yīng)用涌現(xiàn)，AIGC更是火爆出圈。
最常見(jiàn)的AIGC應(yīng)用莫過(guò)于AI繪畫(huà)了。抖音在11月底上線(xiàn)了AI繪畫(huà)特效，根據(jù)照片生成二次元頭像，目前已經(jīng)有近2600萬(wàn)人使用。
抖音AI繪畫(huà)特效背后的技術(shù)支撐，來(lái)自于字節(jié)跳動(dòng)的云服務(wù)平臺(tái)火山引擎，抖音AI繪畫(huà)的上線(xiàn)也意味著火山引擎正式進(jìn)入AIGC領(lǐng)域，為抖音以及全行業(yè)提供底層的技術(shù)支撐、工程優(yōu)化等能力。
AIGC由人工智能生成聲音、圖像、視頻等內(nèi)容，被認(rèn)為是繼PGC（專(zhuān)業(yè)生產(chǎn)內(nèi)容）、UGC（用戶(hù)生產(chǎn)內(nèi)容）之后，新一代的內(nèi)容生產(chǎn)方式。
但AIGC的大規(guī)模普及，仍然尚需時(shí)日，還需要克服高昂的成本等問(wèn)題。
火山引擎的機(jī)器學(xué)習(xí)平臺(tái)，正是為了解決這些問(wèn)題，以此來(lái)推動(dòng)AIGC產(chǎn)業(yè)的發(fā)展。「甲子光年」采訪(fǎng)了火山引擎機(jī)器學(xué)習(xí)平臺(tái)相關(guān)負(fù)責(zé)人，來(lái)探究火山引擎對(duì)AIGC的投入與創(chuàng)新。
火山引擎相對(duì)PyTorch原生基線(xiàn)訓(xùn)練提升40%、推理提升3.47倍。對(duì)此，火山引擎是如何做到的？
AIGC火了之后，也催生了新興的AI獨(dú)角獸，最具代表性的就是主打AI文字生成的Jasper.ai，以及主打AI生成繪畫(huà)的Stability AI。
科技巨頭也不甘落后，紛紛推出各自的AI工具，比如英偉達(dá)推出了全新的Magic3D，能夠基于文本描述自動(dòng)生成3D模型；Meta推出了高性能推理引擎AITemplate，大幅簡(jiǎn)化多GPU后端部署。
今年7月Stability AI這家公司將其AI文本-圖像生成模型Stable Diffusion開(kāi)源，后續(xù)開(kāi)發(fā)者只需要在此開(kāi)源模型的基礎(chǔ)上做一定的調(diào)整開(kāi)發(fā)工作，就可以得到屬于自己的AI繪畫(huà)創(chuàng)作。
但是，AIGC仍然存在一定的門(mén)檻，尤其是對(duì)于國(guó)內(nèi)的開(kāi)發(fā)者來(lái)說(shuō)。
現(xiàn)有Diffusion模型主要使用英文數(shù)據(jù)進(jìn)行訓(xùn)練。由于中英文在文化和表達(dá)上的的差異性，如果直接將英文數(shù)據(jù)翻譯成中文進(jìn)行模型訓(xùn)練，產(chǎn)出的模型通常無(wú)法建模中文特有的現(xiàn)象。此外，通用的Stable Diffusion模型由于數(shù)據(jù)源的限制，很難用于生成特定領(lǐng)域、特定場(chǎng)景下的高清圖片。
同時(shí)，高昂的成本也一定程度上限制了AIGC產(chǎn)業(yè)的大規(guī)模普及。
以Stable Diffusion為例，最初版本的Stable Diffusion動(dòng)用了256塊英偉達(dá)A100 GPU，花費(fèi)了15 萬(wàn)小時(shí)的總訓(xùn)練時(shí)長(zhǎng)，僅這一項(xiàng)成本就是60萬(wàn)美元。即使是Fine Tune（微調(diào)），也需要消耗幾天時(shí)間。換算成錢(qián)來(lái)看，一次完整的預(yù)訓(xùn)練可能需要上百萬(wàn)的成本，F(xiàn)ine Tune階段也需要二十多萬(wàn)/次。目前AIGC市場(chǎng)火爆，各種畫(huà)風(fēng)和產(chǎn)品形態(tài)頻頻出現(xiàn)且快速演進(jìn)，這樣的預(yù)訓(xùn)練和Fine Tune的次數(shù)肯定不會(huì)少。
對(duì)于創(chuàng)業(yè)公司來(lái)說(shuō)，開(kāi)發(fā)這樣體量的大模型困難重重，不僅僅是寫(xiě)一套算法那么簡(jiǎn)單，還包括數(shù)據(jù)準(zhǔn)備、模型開(kāi)發(fā)、訓(xùn)練調(diào)優(yōu)、模型部署等多個(gè)環(huán)節(jié)，甚至在模型的訓(xùn)練與模型的推理階段還需使用不同的芯片，每一個(gè)環(huán)節(jié)工作量都不小。
其次，在模型推理方面，隨著AIGC火爆全網(wǎng)，各個(gè)產(chǎn)品的QPS（Queries Per Second，每秒查詢(xún)率）壓力都非常大，原先預(yù)計(jì)的幾百上千資源，可能需要短時(shí)間擴(kuò)容至上千甚至是萬(wàn)規(guī)模，才能保證前端的服務(wù)質(zhì)量。而目前的開(kāi)源加速方案，僅適用于標(biāo)準(zhǔn)的Stable Diffusion加速，需要用戶(hù)有較強(qiáng)的工程適配能力。
火山引擎就觀(guān)察到，火山上的很多用戶(hù)，由于自己的模型有些許改動(dòng)、且時(shí)間有限，都是用原始框架直接部署。
因此，不少AI初創(chuàng)公司選擇把模型開(kāi)發(fā)這項(xiàng)“AI煉丹”的工作通過(guò)云來(lái)完成，以期相對(duì)低成本、又能無(wú)限算力的方式獲取云端資源。
Stability AI選擇了AWS的云服務(wù)，目前擁有一個(gè)在 AWS 中運(yùn)行的由 4000 多個(gè)英偉達(dá)A100 GPU 組成的集群，用于訓(xùn)練包括 Stable Diffusion在內(nèi)的 AI 系統(tǒng)。在今年的亞馬遜re：Invent大會(huì)上宣布雙方達(dá)成戰(zhàn)略合作。
云廠(chǎng)商基本也是基于這個(gè)邏輯布局AIGC，即為AIGC產(chǎn)業(yè)提供基礎(chǔ)設(shè)施。
作為最年輕的一朵云，火山引擎在AIGC領(lǐng)域的表現(xiàn)卻一點(diǎn)也不像個(gè)新手，其模型推理速度達(dá)到了66it/s，是目前已公開(kāi)的模型優(yōu)化速度最快的平臺(tái)。
火山引擎憑什么做到？
火山引擎在性能、顯存消耗以及接入方式這3個(gè)方面，都做到了極致。
首先，是訓(xùn)練和推理的GPU性能。火山引擎計(jì)算加速團(tuán)隊(duì)對(duì)Stable Diffusion的整個(gè)模型做了詳盡的分析，然后按照如下思路逐步優(yōu)化訓(xùn)練和推理的性能。
第一步，確定Kernel Fusion優(yōu)化范圍。
眾所周知，Kernel Fusion是GPU優(yōu)化最常用有效的技巧，通過(guò)將多個(gè)連續(xù)的或者可以并行執(zhí)行的操作，融合為盡可能少的GPU kernels，可以大量減少重復(fù)的顯存讀寫(xiě)、降低Kernel launch的開(kāi)銷(xiāo)，從而實(shí)現(xiàn)更優(yōu)的性能。
然而，由于Kernel Fusion后的算子粒度往往較原始實(shí)現(xiàn)大很多，有時(shí)候大量的Kernel Fusion會(huì)極大地限制優(yōu)化適用的靈活性。
對(duì)此，火山引擎優(yōu)化團(tuán)隊(duì)充分考慮模型內(nèi)在的計(jì)算邏輯、不同GPU硬件的硬件參數(shù)及特性、以及現(xiàn)有的GPU基礎(chǔ)庫(kù)提供的基礎(chǔ)能力，通過(guò)設(shè)計(jì)多個(gè)層級(jí)不同粒度的優(yōu)化算子，同時(shí)兼顧優(yōu)化性能以及模型未來(lái)迭代的發(fā)展趨勢(shì)。
比如，通過(guò)擴(kuò)展Cutlass，火山引擎在卷積操作前后可以額外添加豐富靈活的前后處理，比如add bias或者residual等，這就打破了以往以卷積為界限的Kernel Fusion斷點(diǎn)，在Stable Diffusion這樣卷積操作很重且Tensor很大的網(wǎng)絡(luò)中就可以取得很好的性能收益。再比如，依托Flash-Attention算法，火山引擎可以突破片上shared memory的限制，將整個(gè)Attention操作融合為一個(gè)操作。
經(jīng)過(guò)Kernel Fusion后，Stable Diffusion v1 訓(xùn)練和推理的CUDA Kernels數(shù)量相比原始的PyTorch實(shí)現(xiàn)，分別減少了50%和60%。
第二步，并行算法重設(shè)計(jì)。
對(duì)于每個(gè)算子，均會(huì)從問(wèn)題的本質(zhì)出發(fā)，并結(jié)合現(xiàn)有的GPU硬件參數(shù)，設(shè)計(jì)更優(yōu)的并行算法實(shí)現(xiàn)。
比如，LayerNorm是深度學(xué)習(xí)里非常典型常用的一個(gè)操作，其反向操作均涉及到大量的按列求和操作，PyTorch使用比較基礎(chǔ)的實(shí)現(xiàn)來(lái)組裝，而火山引擎通過(guò)更好的并行算法實(shí)現(xiàn)，僅用2個(gè)kernel就完成了全部的反向操作。再比如，GroupNorm是StableDiffusion中耗時(shí)較高的一個(gè)操作，其性能非常受Tensor Layout排布的影響。火山引擎對(duì)兩種Layout下的GroupNorm操作，做了不同的并行算法設(shè)計(jì)，最終該操作前后向分別加速50%和30%。
第三步，SOL評(píng)估。對(duì)每個(gè)操作，都會(huì)仔細(xì)評(píng)估其需要的最少數(shù)據(jù)讀寫(xiě)、以及浮點(diǎn)操作數(shù)量。過(guò)程中會(huì)借助一系列Profiling工具，深入SASS Code層，爭(zhēng)取每個(gè)CUDA Kernel都能達(dá)到SOL的80-90%的水平。
其次，降低顯存消耗。
性能優(yōu)化的過(guò)程中，通過(guò)大量的Kernel Fusion，就可以大幅減少中間Tensor的讀寫(xiě)。除此之外，火山引擎進(jìn)一步，在訓(xùn)練優(yōu)化過(guò)程中，將部分輕量級(jí)操作以re-compute的方式來(lái)進(jìn)行反向推導(dǎo)，可以進(jìn)一步減少顯存消耗。最終，可以使單張上可以同時(shí)計(jì)算更大的batch_size，進(jìn)一步提高整體吞吐。
最后，火山引擎在接入方式上也做了極致簡(jiǎn)化。
因?yàn)榛鹕揭嫣峁┑募铀偎阕樱瑫r(shí)支持高效訓(xùn)練和推理，從而真正實(shí)現(xiàn)了訓(xùn)推一體，因此可以從模型訓(xùn)練到推理部署做到無(wú)縫銜接。接入方式上提供由Pytorch Custom op搭建出的相同模型結(jié)構(gòu)，完全兼容原始Stable Diffusion模型的Pipeline, 接入方便且無(wú)需改動(dòng)框架。
在推理部署時(shí)，高性能算子庫(kù)可以支持不同GPU硬件，包括英偉達(dá)V100、A100、A10、A30等不同系列的GPU卡，而且不需要做額外的模型轉(zhuǎn)換工作。
未經(jīng)優(yōu)化時(shí)，絕大多數(shù)模型只能在最高端的A100 80GB訓(xùn)練卡上跑起來(lái)。而火山引擎提供的訓(xùn)練加速，能夠擺脫特定計(jì)算卡的限制，讓客戶(hù)在更廣泛的范圍內(nèi)尋找適合自己需求的解決方案。
火山引擎目前提供40%的訓(xùn)練加速，這樣在同類(lèi)型的云產(chǎn)品上訓(xùn)練，無(wú)論是迭代速度，還是單次的訓(xùn)練成本，都有了明顯的提升和節(jié)省。
今年7月，Stable Diffusion的V1版本開(kāi)源時(shí)，火山引擎已經(jīng)做了初步的嘗試。11月25日Stable Diffusion的V2版本發(fā)布，火山引擎在此基礎(chǔ)上又花了不到2周的時(shí)間，做了基礎(chǔ)工程的優(yōu)化。這是抖音AI繪畫(huà)特效上線(xiàn)的技術(shù)準(zhǔn)備工作。
除了做好技術(shù)的基本功，還需要有一個(gè)好的練兵場(chǎng)。對(duì)于火山引擎而言，抖音就在字節(jié)跳動(dòng)內(nèi)部為火山引擎提供了一個(gè)絕佳的內(nèi)部實(shí)踐。
火山引擎機(jī)器學(xué)習(xí)總監(jiān)吳迪告訴「甲子光年」：“要把包含了AIGC以及機(jī)器翻譯、人機(jī)交互、語(yǔ)音識(shí)別等多媒體模型的訓(xùn)練和推理做好，是有一定門(mén)檻的，需要長(zhǎng)時(shí)間的技術(shù)和經(jīng)驗(yàn)積累。而抖音的產(chǎn)品形態(tài)中包含了豐富的多模態(tài)內(nèi)容，抖音的高標(biāo)準(zhǔn)要求和多樣化的產(chǎn)品形態(tài)在不停地打磨火山引擎機(jī)器學(xué)習(xí)平臺(tái)的工程團(tuán)隊(duì)。在這樣嚴(yán)苛的長(zhǎng)期錘煉下， 火山機(jī)器學(xué)習(xí)團(tuán)隊(duì)有信心，把最有競(jìng)爭(zhēng)力的機(jī)器學(xué)習(xí)優(yōu)化能力帶給客戶(hù)。幫助客戶(hù)用更低的成本去敏捷迭代。”
在內(nèi)部與抖音的合作過(guò)程中，火山引擎的機(jī)器學(xué)習(xí)平臺(tái)團(tuán)隊(duì)主要做工程優(yōu)化方面的工作。有時(shí)候如果僅僅依靠工程團(tuán)隊(duì)仍然沒(méi)有辦法滿(mǎn)足上線(xiàn)需求，就會(huì)與算法團(tuán)隊(duì)一起做更透徹的優(yōu)化，比如進(jìn)一步的量化，或者是對(duì)模型進(jìn)行一些蒸餾、裁剪等。
抖音對(duì)于火山引擎所提出的訴求，往往也代表了所有市場(chǎng)用戶(hù)的基本訴求，首當(dāng)其沖的就是“降低成本”。
作為一款國(guó)民級(jí)短視頻應(yīng)用，抖音所持有的機(jī)器總量很大，因此即使是細(xì)微的性能節(jié)省，都將帶來(lái)巨大的資金回報(bào)。因此，當(dāng)火山引擎的AIGC模型，優(yōu)化速度提升40%，硬件成本降低百分之二三十，訓(xùn)練成本降低80%，相比目前最快的TensorRT部署方案，推理還快10%，將會(huì)帶來(lái)巨大的成本優(yōu)勢(shì)。
除了更省錢(qián)之外，還必須要穩(wěn)定，不能出錯(cuò)。
近期OpenAI推出的聊天機(jī)器人模型chatGPT，就因?yàn)樘^(guò)火爆而一度服務(wù)器崩潰。作為一款現(xiàn)象級(jí)產(chǎn)品，這并沒(méi)有影響到用戶(hù)對(duì)chatGPT的熱情，但用戶(hù)體驗(yàn)總歸不夠友好，同樣的事情如果發(fā)生在抖音，可能就是一次產(chǎn)品事故。
為了避免這種情況發(fā)生，火山引擎必須做到足夠的穩(wěn)定性，能夠抗住激增流量的壓力。為此，火山引擎不僅要做模型的優(yōu)化，還要在服務(wù)框架層面做一些穩(wěn)定性的優(yōu)化。
第三是要保證充足的資源供應(yīng)。
火山引擎管理了大量計(jì)算資源，來(lái)支撐字節(jié)跳動(dòng)的整體業(yè)務(wù)。但這些卡片不會(huì)是永遠(yuǎn)忙碌的，一定有的處在業(yè)務(wù)峰值，有的處于相對(duì)閑置的狀態(tài)。此時(shí)，就需要有一個(gè)自動(dòng)化的資源調(diào)動(dòng)系統(tǒng)，能夠“集中力量干大事”。比如，當(dāng)AIGC的應(yīng)用在短時(shí)間內(nèi)需要大量資源的時(shí)候，火山引擎可以在小時(shí)級(jí)或者分鐘級(jí)的時(shí)間內(nèi)，從其他相對(duì)閑置的場(chǎng)景資源去調(diào)取，形成一種靈活、彈性的資源供應(yīng)。
最后就是滿(mǎn)足差異化的場(chǎng)景。
AIGC的火爆也帶火了機(jī)器學(xué)習(xí)在多媒體方面的應(yīng)用。其實(shí)在抖音內(nèi)部，多媒體的機(jī)器學(xué)習(xí)應(yīng)用早就層出不窮了，火山引擎的模型優(yōu)化要去理解每一幀圖片、每一個(gè)視頻、文字與語(yǔ)音之間的關(guān)聯(lián)等等，其豐富程度通常遠(yuǎn)大于市面上常見(jiàn)的一些形態(tài)。
同時(shí)，因?yàn)槎兑粢灿袕?qiáng)大的算法團(tuán)隊(duì)，意味著模型的變化幅度與速度都是非常快的，這也是市場(chǎng)的的一些優(yōu)化方案僅僅支持好一些基礎(chǔ)版本的模型就可以，而火山引擎一定要考慮算法整體設(shè)計(jì)方案的靈活性，滿(mǎn)足“變來(lái)變?nèi)?rdquo;的需求。
如今，抖音的AI繪畫(huà)特效已經(jīng)有接近2600萬(wàn)人使用。
AIGC目前仍然處在大規(guī)模爆發(fā)的前夜，而諸如AI繪畫(huà)這樣的早期應(yīng)用，正是打開(kāi)AIGC市場(chǎng)的一把鑰匙。而隨著火山引擎這樣的云平臺(tái)不斷把降本、增效這件事做到極致，AIGC產(chǎn)業(yè)大規(guī)模爆發(fā)的那一刻，有望在不久的將來(lái)快速實(shí)現(xiàn)。
END.
親愛(ài)的朋友們，不知不覺(jué)2022即將過(guò)去，感謝大家又一年的陪伴，在2022的最后幾周，我們邀請(qǐng)所有的朋友們來(lái)分享你在2022難忘的回憶，我們將挑選其中的一些故事放入甲子光年2022年終回顧文章中，掃碼下方二維碼，分享屬于你的2022故事吧！
本文來(lái)自微信公眾號(hào)“甲子光年”（ID:jazzyear），作者：趙健，36氪經(jīng)授權(quán)發(fā)布。
