大模型狂潮背后：AI基礎(chǔ)設(shè)施的“老化”與改造工程
2023-02-25
機(jī)器學(xué)習(xí)模型逐漸發(fā)展成人們口中的“龐然大物”。全球頂尖的科技公司紛紛踏上“軍備競(jìng)賽”之路，立志訓(xùn)練出規(guī)模最大的模型（MUM、OPT、GPT-3、Megatron），而其他專注于生產(chǎn)系統(tǒng)的公司也相繼擴(kuò)大其原有模型，并取得良好成果。
一切如火如荼，然而，鮮少有人提及，龐大的模型給現(xiàn)有的AI基礎(chǔ)設(shè)施和開發(fā)流程帶來了諸多實(shí)際性挑戰(zhàn)。
大模型的權(quán)重可達(dá)100+GB，而目前的開發(fā)工具卻還沒跟上，使用起來十分費(fèi)力，部署時(shí)往往要等上好幾分鐘甚至好幾小時(shí)，這已經(jīng)成為AI工程師的隱痛，不但浪費(fèi)工程師的時(shí)間，降低工作效率，還會(huì)拖慢迭代速度。
致力于AI基礎(chǔ)設(shè)施工具研發(fā)的Modular團(tuán)隊(duì)認(rèn)為，開發(fā)人員的工作效率是訓(xùn)練和部署模型的最大成本之一。因此需要不斷優(yōu)化工具鏈，提升早期用戶的體驗(yàn)，也方便開發(fā)人員。本文探討編譯過程中管理海量數(shù)據(jù)的技術(shù)難點(diǎn)，以及Modular為解決這些難點(diǎn)在基礎(chǔ)設(shè)施（以及MLIR編譯器框架）方面所做的改進(jìn)。由OneFlow社區(qū)（ID：OneFlowTechnology）編譯。
AI模型配套工具的易用性不足
機(jī)器學(xué)習(xí)中的圖轉(zhuǎn)換（Graph Transformations）、優(yōu)化和編譯器等技術(shù)的作用是提升AI模型的性能和便攜性，讓模型可以部署在某些目標(biāo)硬件上。
編譯器中，有TensorFlow Lite Converter這樣的高層次“編譯器”，它可以將TensorFlow SavedModel模型轉(zhuǎn)換為高度優(yōu)化的程序格式（如FlatBuffer格式），讓模型可以在邊緣設(shè)備上執(zhí)行；也有XLA和TorchScript JIT Compiler這樣針對(duì)特定領(lǐng)域的編譯器，它們?yōu)锳I模型創(chuàng)建中間表示（可能是一張“圖”），然后將其編譯成另一種格式——例如機(jī)器碼或特定領(lǐng)域的運(yùn)行時(shí)表示（如CUDA圖）。
AI圖的編譯與傳統(tǒng)的編譯很不一樣。AI圖包含兩部分：圖拓?fù)洌ǜ鲗又g如何連接）和模型權(quán)重（特定層的參數(shù)）。從大小來看，圖拓?fù)湟訩B為單位，權(quán)重則以MB甚至GB為單位。舉個(gè)例子，Meta公司發(fā)布的Open Pre-trained Transformers模型，其參數(shù)量從300億、660億到1750億不等，相當(dāng)于100+GB權(quán)重。Gopher和Megatron模型甚至更大。
AI生態(tài)系統(tǒng)中現(xiàn)有的工具尚不能很好地處理大型模型。比如，Protobufs限制了傳輸數(shù)據(jù)大小不能超過2GB，因此模型如果使用Protobufs序列化格式，就會(huì)備受掣肘。最新版TensorRT的文檔中寫道，“對(duì)于BERT和GPT等基于Transformer的神經(jīng)網(wǎng)絡(luò)模型，TensorRT在編譯時(shí)可以消耗10倍于模型大小的CPU內(nèi)存”，可見TensorRT不適合大型模型。如果使用ONNX文件格式存儲(chǔ)大型模型，就必須將模型權(quán)重分成多個(gè)文件分別存儲(chǔ)。
以上種種不但給AI開發(fā)工作流增加了不必要的復(fù)雜環(huán)節(jié)，也使模型喪失“單一事實(shí)來源”（SSOT），還導(dǎo)致模型分發(fā)更加困難。
為了應(yīng)對(duì)模型權(quán)重太大的問題，大家可能會(huì)采取變通方法，最終卻可能導(dǎo)致整個(gè)AI開發(fā)工作流變得更復(fù)雜。比如，由于某些編譯器階段耗時(shí)長(zhǎng)達(dá)2分多鐘，打斷開發(fā)人員的工作節(jié)奏，所以Modular構(gòu)建了一種緩存臨時(shí)文件的機(jī)制。
雖然這種緩存機(jī)制和其他變通方法一樣，只是治標(biāo)不治本：它既非100%可靠，也不能解決Cache Miss（緩存缺失）的問題，不過由于Modular十分注重提高開發(fā)人員的工作效率，所以還是決定采用這種機(jī)制。