作者｜River Riddle、Eric Johnson、Abdul Dakak翻譯｜胡燕君、楊婷

機(jī)器學(xué)習(xí)模型逐漸發(fā)展成人們口中的“龐然大物”。全球頂尖的科技公司紛紛踏上“軍備競賽”之路，立志訓(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)力，部署時往往要等上好幾分鐘甚至好幾小時，這已經(jīng)成為AI工程師的隱痛，不但浪費(fèi)工程師的時間，降低工作效率，還會拖慢迭代速度。

致力于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）編譯。?

1AI模型配套工具的易用性不足

機(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這樣針對特定領(lǐng)域的編譯器，它們?yōu)锳I模型創(chuàng)建中間表示（可能是一張“圖”），然后將其編譯成另一種格式——例如機(jī)器碼或特定領(lǐng)域的運(yùn)行時表示（如CUDA圖）。

AI圖的編譯與傳統(tǒng)的編譯很不一樣。AI圖包含兩部分：圖拓?fù)洌ǜ鲗又g如何連接）和模型權(quán)重（特定層的參數(shù)）。從大小來看，圖拓?fù)湟訩B為單位，權(quán)重則以MB甚至GB為單位。舉個例子，Meta公司發(fā)布的Open Pre-trained Transformers模型，其參數(shù)量從300億、660億到1750億不等，相當(dāng)于100+GB權(quán)重。Gopher和Megatron模型甚至更大。

圖源DeepMind論文

AI生態(tài)系統(tǒng)中現(xiàn)有的工具尚不能很好地處理大型模型。比如，Protobufs限制了傳輸數(shù)據(jù)大小不能超過2GB，因此模型如果使用Protobufs序列化格式，就會備受掣肘。最新版TensorRT的文檔中寫道，“對于BERT和GPT等基于Transformer的神經(jīng)網(wǎng)絡(luò)模型，TensorRT在編譯時可以消耗10倍于模型大小的CPU內(nèi)存”，可見TensorRT不適合大型模型。如果使用ONNX文件格式存儲大型模型，就必須將模型權(quán)重分成多個文件分別存儲。

以上種種不但給AI開發(fā)工作流增加了不必要的復(fù)雜環(huán)節(jié)，也使模型喪失“單一事實(shí)來源”（SSOT），還導(dǎo)致模型分發(fā)更加困難。

為了應(yīng)對模型權(quán)重太大的問題，大家可能會采取變通方法，最終卻可能導(dǎo)致整個AI開發(fā)工作流變得更復(fù)雜。比如，由于某些編譯器階段耗時長達(dá)2分多鐘，打斷開發(fā)人員的工作節(jié)奏，所以Modular構(gòu)建了一種緩存臨時文件的機(jī)制。

雖然這種緩存機(jī)制和其他變通方法一樣，只是治標(biāo)不治本：它既非100%可靠，也不能解決Cache Miss（緩存缺失）的問題，不過由于Modular十分注重提高開發(fā)人員的工作效率，所以還是決定采用這種機(jī)制。

2Modular編譯棧中的MLIR

Modular的技術(shù)棧中，MLIR編譯器架構(gòu)負(fù)責(zé)表示和轉(zhuǎn)換AI模型，包括AI算子圖（用于多種框架）、中級運(yùn)行時原語和低級機(jī)器碼生成。

多級中間表示 (MLIR)

MLIR是LLVM編譯器基礎(chǔ)設(shè)施項(xiàng)目的子項(xiàng)目，LLVM旨在提供現(xiàn)代工具包，用以構(gòu)建針對特定領(lǐng)域的編譯器。MLIR提供一套核心組件，用于硬件設(shè)計、量子計算、人工智能等多種計算領(lǐng)域的建模、分析和轉(zhuǎn)換。

MLIR能夠幫助構(gòu)建單個涵蓋全棧的完整系統(tǒng)，比常規(guī)的技術(shù)棧功能更強(qiáng)大、模塊化程度和可拓展性更高，也更易于維護(hù)。使用統(tǒng)一的基礎(chǔ)設(shè)施讓我們得以便捷地將每一項(xiàng)改進(jìn)遷移到自己的工具棧，使開發(fā)工作流實(shí)現(xiàn)更高的模塊化和可組裝性。

除了Modular以外，TensorFlow、XLA、PyTorch和ONNX等也在使用MLIR進(jìn)行模型表示和轉(zhuǎn)換。隨著MLIR的用戶生態(tài)不斷擴(kuò)大，在贊美MLIR優(yōu)點(diǎn)的同時，也必須繼續(xù)進(jìn)行改進(jìn)和完善。?

3

MLIR管理權(quán)重的方法還有待提高

MLIR的基本組成部分之一是屬性機(jī)制（Attribute），可以把它理解為被unique（或被memoize、intern）的常量數(shù)據(jù)。屬性是用戶可拓展的，也就是說，可以根據(jù)不同用例使用不同的屬性類型。很多類型的值都可以被賦予屬性，比如常量表達(dá)式值（如“5”、“10.0”等）、字符串字面量、枚舉值（如“小于”、“大于”、“等于”等），數(shù)據(jù)組等等。大多數(shù)基于MLIR的AI工具都使用屬性來保存AI模型的權(quán)重。

然而，問題出現(xiàn)了：模型權(quán)重有可能極其龐大，但MLIR存儲2 GB權(quán)重的方式和存儲4 B權(quán)重的方式并沒有區(qū)別——都使用同一屬性，該屬性包含一組被unique的元素。但對GB級的龐大數(shù)據(jù)使用unique方法顯然不合理。

這個方法的難點(diǎn)在于：在MLIR中，當(dāng)某個東西被unique，它就會被分配（allocated）、被hash 、然后被儲存到MLIRContext中。這些東西具有和MLIRContext相同的生命周期，只有當(dāng)MLIRContext被銷毀，它們才會同時被銷毀。對于小的數(shù)值而言，這種機(jī)制帶來很多好處，可以把數(shù)值傳入傳出，可以通過指針對unique后的值進(jìn)行比較，還可以共享屬性的內(nèi)存分配（十分常見）等等。

但對數(shù)量龐大的權(quán)重而言，上述種種好處就變成了劣勢：我們不希望對權(quán)重進(jìn)行重新分配、復(fù)制或使用unique方法，我們只需要它們短暫存在——當(dāng)計算不再需要引用這些權(quán)重時，就要允許釋放分配。例如，當(dāng)運(yùn)行模型量化工具時，需要對算子圖進(jìn)行轉(zhuǎn)換，并生成新的權(quán)重，最終這些權(quán)重可能會被復(fù)制多份，大量權(quán)重副本在編譯結(jié)束前都將一直占用內(nèi)存。?

ML工具的另一個問題是MLIR如何序列化至文件系統(tǒng)。一開始，MLIR沒有二進(jìn)制序列化格式，只有文本格式。對數(shù)量龐大的權(quán)重來說，這就造成問題，因?yàn)槊總€字節(jié)的二進(jìn)制數(shù)據(jù)都會被轉(zhuǎn)化為十六進(jìn)制，后者占用的空間為前者的2倍。這樣一來，我們不但耗費(fèi)了相當(dāng)長的時間進(jìn)行進(jìn)制轉(zhuǎn)換（一個中等的GB級模型大概需要20秒），而且轉(zhuǎn)換后的中間文件還是原來的2倍大——2倍可是一個不小的數(shù)字！

4內(nèi)存占用：比拖慢開發(fā)效率更嚴(yán)重的影響

這一設(shè)計機(jī)制本意雖好，但它有可能降低編譯效率，即便用最好的編譯器也無濟(jì)于事。最明顯的問題是它會導(dǎo)致編譯、監(jiān)控和轉(zhuǎn)換模型的時間變長。但凡你曾用過“我的代碼還在編譯”作為日常摸魚的借口，你就明白等待編譯是多么痛苦的事情。采用這一機(jī)制，就意味著處理器不得不對GB級數(shù)據(jù)持續(xù)進(jìn)行分配、復(fù)制和hash處理。

XKCD漫畫 – 《還在編譯》?

比編譯時長更嚴(yán)重的問題是內(nèi)存占用，它會影響Modular技術(shù)棧中的其他架構(gòu)功能的實(shí)現(xiàn)。例如，由于我們的編譯器和技術(shù)棧本身都高度并行，而且使用線上搜索等高級功能，內(nèi)存占用會直接導(dǎo)致一些工作不能并行展開，也導(dǎo)致不能取得最高質(zhì)量的結(jié)果。

Modular的價值核心是構(gòu)建用戶喜歡的工具。高級功能如果不好用，或者會影響效率，又或者附帶一些注意事項(xiàng)（比如，“該功能對某些情況不適用”），那么用戶就根本不會用。因此，Modular致力于解決龐大權(quán)重帶來的基礎(chǔ)性問題，簡化用戶的使用流程和開發(fā)人員的工作流程。

5MLIR的核心改進(jìn)

Modular團(tuán)隊(duì)是MLIR項(xiàng)目的重要貢獻(xiàn)者，Modular企業(yè)文化的一大要點(diǎn)是“做對的產(chǎn)品”，Modular參與的所有項(xiàng)目都遵循這一要義。在推動MLIR發(fā)展的同時，Modular竭力保證MLIR項(xiàng)目的每一步路都正確，也加強(qiáng)與MLIR社區(qū)的合作，為所采取的辦法爭取認(rèn)可。?

Modular團(tuán)隊(duì)列出了大型模型工具應(yīng)該具備的特點(diǎn)：

非必要不分配內(nèi)存：對大型數(shù)據(jù)（比如權(quán)重）而言，從磁盤中實(shí)行內(nèi)存映射比將數(shù)據(jù)復(fù)制到已分配內(nèi)存的block中更高效。

無需進(jìn)行hash或unique處理：我們不希望費(fèi)力氣去檢查2 GB Blob數(shù)據(jù)的相等性；要辨別權(quán)重，希望能夠通過名稱辨別，而不是看具體內(nèi)容有沒有被unique。

允許內(nèi)聯(lián)變更（Inline Mutation）：如果數(shù)據(jù)只需要在一處使用，應(yīng)當(dāng)允許在原位置量化、轉(zhuǎn)化和操作數(shù)據(jù)，而不是先復(fù)制數(shù)據(jù)。

允許釋放內(nèi)存（deallocation）：由于大模型的數(shù)據(jù)量十分龐大，因此當(dāng)對某一數(shù)據(jù)的所有引用都不存在時，應(yīng)當(dāng)允許釋放內(nèi)存。

快速序列化：無論是即時編譯，搜索優(yōu)化參數(shù)，還是本地迭代，都需要緩存IR，所以這一步必須快。

上述觀點(diǎn)并不新穎，但傳統(tǒng)編譯器（比如適用于典型CPU編程語言的編譯器）卻還沒有實(shí)現(xiàn)這些要求。?

6調(diào)整權(quán)重屬性

上述前四點(diǎn)要求解決了我們應(yīng)該如何使用MLIR這一基本問題：權(quán)重雖然是常量數(shù)據(jù)，但對它的管理應(yīng)該區(qū)別于其他MLIR屬性。一直以來，我們的權(quán)重管理方式都很不適宜，這就好比試圖將一塊方釘擠進(jìn)圓孔中，不僅浪費(fèi)了空間，降低了我們的開發(fā)速度，同時也增加了用戶成本。

所以Modular決定換一種方法來管理權(quán)重數(shù)據(jù)，這促成了MLIR的第一個基本擴(kuò)展機(jī)制——“Resource機(jī)制”，在計算中將數(shù)據(jù)和對數(shù)據(jù)的引用區(qū)分開來。

在Resource機(jī)制中，序列化MLIR的每個Blob都可能包含額外的信息段，稱為Resource。Resource要么是dialect（擴(kuò)展MLIR時使用的類似namespace的抽象），要么是用于特定工具鏈數(shù)據(jù)的“外部（external）”資源。Resource中的數(shù)據(jù)用簡單的鍵值對表示，創(chuàng)造出如下圖所示的類似json的結(jié)構(gòu)。

/// Here we have some MLIR operations.module { func.func @foo() { // Cool stuff here ... }}/// Here we have an `external_resources` section. The resource section"s syntax is designed to be unique as to not conflict with other MLIR syntax (which is user extensible!).{-# external_resources: { mlir_reproducer: { pipeline: "func.func(cse,canonicalize),inline", disable_threading: true } }#-}

上面例子展示了如何調(diào)整MLIR來用Resource進(jìn)行復(fù)現(xiàn)。MLIR再生器（Reproducer）實(shí)際上是一種配置，它包含轉(zhuǎn)換管道（Transformation Pipeline）等執(zhí)行信息，用于復(fù)現(xiàn)某種故障或失敗。在使用Resource之前，我們通過在MLIR文件頂部添加注釋來表示這些執(zhí)行信息?，F(xiàn)在可以利用Resource將這些執(zhí)行信息合并為第一類信息。

從前需要進(jìn)行unique處理導(dǎo)致長期占用內(nèi)存的大型權(quán)重數(shù)據(jù)，現(xiàn)在可以利用Resource機(jī)制進(jìn)行儲存。在IR中，我們對屬性采用輕量級引用而不再采用底層數(shù)據(jù)：

其他優(yōu)勢：

使用IR進(jìn)行調(diào)試時更不容易出錯，從而帶來更好的開發(fā)體驗(yàn)：Resource是專門的信息段；我們不必?fù)?dān)心在調(diào)試時會不小心轉(zhuǎn)儲整整4GB的數(shù)據(jù)。

我們可以在無需數(shù)據(jù)的情況下合理地處理IR：因?yàn)镮R只保存對數(shù)據(jù)的引用，不保存數(shù)據(jù)本身，如果需要，我們可以省略底層Resource數(shù)據(jù)。這樣做的好處包括可以極大地簡化再生器生成流程，再生器本來就不需要用到大型權(quán)重數(shù)據(jù)（設(shè)想一下，你以前需要向同事發(fā)送1.2GB的再現(xiàn)器文件，現(xiàn)在的再生器文件只有20MB大）。

通過引入Resource這個新概念，我們在程序和數(shù)據(jù)之間建立清晰的分離機(jī)制。現(xiàn)在，我們不再將權(quán)重數(shù)據(jù)直接傳遞給某一屬性。相反，我們向?qū)傩詡魅胍粋€弱引用，并將權(quán)重數(shù)據(jù)傳給一個專門的管理器。這樣，我們就能更好地控制權(quán)重分配、變更和銷毀的時間和方式。

7新增MLIR二進(jìn)制編碼方式

有了更好的權(quán)重表示方法之后，下一步我們只需找到更高效的權(quán)重儲存方法來完成MLIR表示的序列化。?

到此為止，MLIR只有文本序列化格式，這種格式使用ASCII十六進(jìn)制字符串來表示權(quán)重。然而，Modular的終極目標(biāo)是盡可能加快本地開發(fā)流程，因此需要摒棄文本序列化格式，為MLIR新增合適的二進(jìn)制格式（https://discourse.llvm.org/t/rfc-a-binary-serialization-format-for-mlir/63518）。

二進(jìn)制格式需要考慮很多因素，況且二進(jìn)制格式?jīng)Q定了編譯器的穩(wěn)定性。MLIR需要高度的靈活性才能高效應(yīng)對各種各樣的用例，需要實(shí)現(xiàn)高速度，而且MLIR/LLVM不能依賴第三方編碼庫。

不過，MLIR的一大好處是編碼難度極低。因?yàn)镸LIR中所有操作的結(jié)構(gòu)都相同，所有操作都可以使用相同的編碼方式。上述的種種復(fù)雜要求都是為了保證MLIR核心概念的緊湊和高效?？紤]到這些限制，我們決定為MLIR定制編碼方式（https://mlir.llvm.org/docs/BytecodeFormat/）。

8用戶收益

為MLIR增加Resource機(jī)制和二進(jìn)制編碼方式大大加速了工具鏈和開發(fā)流程，并大幅降低內(nèi)存占用，提高了性能和速度表現(xiàn)，也整體改善了MLIR。

為了驗(yàn)證上述改進(jìn)帶來的性能變化，可以測試不同規(guī)模的模型上基于MLIR的圖編譯器中“降級”和“優(yōu)化”步驟的實(shí)際速度（將TensorFlow序列化模型轉(zhuǎn)化為符合MLIR運(yùn)行時輸入格式的模型），以及該過程中的實(shí)際內(nèi)存占用。?

速度提升：編譯工作流

測試結(jié)果發(fā)現(xiàn)，MLIR的速度大幅提升。從TensorFlow序列化模型（TensorFlow 2.10模型）轉(zhuǎn)化為MLIR運(yùn)行時輸入格式的模型，這一過程涉及大量底層表示轉(zhuǎn)換，經(jīng)過改進(jìn)后，實(shí)際執(zhí)行時間縮短了1.8~2倍，執(zhí)行速度隨模型大小按比例縮放。

具體而言，處理TensorFlow序列化模型耗時極短——生成MLIR時將大量權(quán)重數(shù)據(jù)寫入磁盤這一步驟是主要的耗時來源。經(jīng)改進(jìn)后，代碼處理時間比原來快10倍，整體執(zhí)行時間的快慢主要取決于固態(tài)硬盤（SSD）將 >1 GB數(shù)據(jù)寫入磁盤的速度。

ML開發(fā)人員使用我們的工具，可以加快模型編譯速度，從而提升生產(chǎn)效率，減少迭代時間。我們的工具可以優(yōu)化生產(chǎn)環(huán)境以及模型的加載和編譯，包括基于流入數(shù)據(jù)的動態(tài)模型加載和卸載，以及各種個性化或經(jīng)過精細(xì)化調(diào)整的用戶模型。

速度提升：序列化

引入二進(jìn)制編碼方式不但可以加快編譯工作流，還能加快序列化速度。通過外部工具與MLIR進(jìn)行交互，包括運(yùn)行時類型檢查（Introspection）、緩存和再生器生成等，都需要對序列化MLIR進(jìn)行讀寫。

通過對不同規(guī)模的模型進(jìn)行了序列化測試，結(jié)果同樣發(fā)現(xiàn)峰值性能大幅提速，且SSD寫入步驟依然為主要耗時來源。具體而言，大型模型文本數(shù)據(jù)的讀取耗時約5秒，而二進(jìn)制數(shù)據(jù)的讀取耗時僅不到10毫秒；二進(jìn)制格式的寫入速度則約是文本格式數(shù)據(jù)的5倍。

對Modular而言，引入二進(jìn)制編碼方式可以加快以MLIR為中心的基礎(chǔ)設(shè)施和工具的開發(fā)速度，改善原本高成本、低速度的開發(fā)狀況。比如，調(diào)試器（Debugger）的效率很大程度取決于編譯工作流中緩存模型表示的效率，而引入二進(jìn)制編碼方式可以提高調(diào)試器的效率，從而提高底層編譯器的性能。?

內(nèi)存占用

二進(jìn)制序列化格式的mmap（一種內(nèi)存映射方法）性能以及通過Resource機(jī)制實(shí)現(xiàn)的IR和數(shù)據(jù)的相互獨(dú)立性可以大幅減少內(nèi)存占用。測試發(fā)現(xiàn)，各種規(guī)模的模型編譯流程中的內(nèi)存占用都大大降低——因?yàn)橐郧靶枰獮槟Ｐ蜋?quán)重分配內(nèi)存，現(xiàn)在不需要了。 ?

9升級AI生態(tài)

Modular的愿景不只是為了方便我們自己，而是升級整個AI行業(yè)的生態(tài)。前文提及的新型Resource表示和二進(jìn)制編碼方式都已提交至上游的LLVM/MLIR倉庫中。?

Modular起初的研發(fā)動機(jī)是為了解決Modular的客戶遇到的問題并提升自身內(nèi)部基礎(chǔ)設(shè)施，但這些改進(jìn)產(chǎn)生的積極影響并不限于自己的用例，還能改善其他以MLIR為基礎(chǔ)技術(shù)的產(chǎn)品。例如，由于二進(jìn)制編碼方式的引進(jìn)，MLIR社區(qū)如今正在討論如何保證MLIR的穩(wěn)定性（https://discourse.llvm.org/t/mlir-generic-ir-stability-and-upgradability/65371）。

這些基礎(chǔ)技術(shù)的改進(jìn)最終都會融入產(chǎn)品中為用戶服務(wù)。以上只是Modular致力提升的無數(shù)核心技術(shù)之一。Modular一方面竭力適應(yīng)大模型，一方面努力完善模型在設(shè)備上的部署，目標(biāo)都是大幅提升AI基礎(chǔ)設(shè)施的性能和易用性。Modular非常看好AI的未來以及LLVM和MLIR的發(fā)展。 （本文由OneFlow社區(qū)翻譯，譯文轉(zhuǎn)載請聯(lián)系OneFlow獲得授權(quán)。原文：1. https://www.modular.com/blog/increasing-development-velocity-of-giant-ai-models；2.https://www.modular.com/blog/increasing-development-velocity-of-giant-ai-models-part-2）

其他人都在看

LLVM之父：編譯器的黃金時代

李白：你的模型權(quán)重很不錯，可惜被我沒收了

更快的YOLOv5問世，附送全面中文解析教程

LLVM之父：為什么我們要重建AI基礎(chǔ)設(shè)施軟件

LLVM之父：模塊化設(shè)計決定AI前途，不服來辯

開源吞噬AI界？從Stable Diffusion的爆火說起

OneEmbedding:單卡訓(xùn)練TB級推薦模型不是夢

歡迎Star、試用OneFlow最新版本：GitHub - Oneflow-Inc/oneflow: OneFlow is a deep learning framework designed to be user-friendly, scalable and efficient.OneFlow is a deep learning framework designed to be user-friendly, scalable and efficient. - GitHub - Oneflow-Inc/oneflow: OneFlow is a deep learning framework designed to be user-friendly, scalable and efficient.https://github.com/Oneflow-Inc/oneflow

關(guān)鍵詞： 基礎(chǔ)設(shè)施