作者｜River Riddle、Eric Johnson、Abdul Dakak
翻譯｜胡燕君、楊婷
機器學習模型逐漸發展成人們口中的“龐然大物”。全球頂尖的科技公司紛紛踏上“軍備競賽”之路，立志訓練出規模最大的模型（MUM、OPT、GPT-3、Megatron），而其他專注於生產系統的公司也相繼擴大其原有模型，並取得良好成果。
一切如火如荼，然而，鮮少有人提及，龐大的模型給現有的AI基礎設施和開發流程帶來了諸多實際性挑戰。
大模型的權重可達100+GB，而目前的開發工具卻還沒跟上，使用起來十分費力，部署時往往要等上好幾分鐘甚至好幾小時，這已經成為AI工程師的隱痛，不但浪費工程師的時間，降低工作效率，還會拖慢迭代速度。
致力於AI基礎設施工具研發的Modular團隊認為，開發人員的工作效率是訓練和部署模型的最大成本之一。因此需要不斷最佳化工具鏈，提升早期使用者的體驗，也方便開發人員。本文探討編譯過程中管理海量資料的技術難點，以及Modular為解決這些難點在基礎設施（以及MLIR編譯器框架）方面所做的改進。由OneFlow社群（ID：OneFlowTechnology）編譯。
1
AI模型配套工具的易用性不足
機器學習中的圖轉換（Graph Transformations）、最佳化和編譯器等技術的作用是提升AI模型的效能和便攜性，讓模型可以部署在某些目標硬體上。
編譯器中，有TensorFlow Lite Converter這樣的高層次“編譯器”，它可以將TensorFlow SavedModel模型轉換為高度最佳化的程式格式（如FlatBuffer格式），讓模型可以在邊緣裝置上執行；也有XLA和TorchScript JIT Compiler這樣針對特定領域的編譯器，它們為AI模型建立中間表示（可能是一張“圖”），然後將其編譯成另一種格式——例如機器碼或特定領域的執行時表示（如CUDA圖）。
AI圖的編譯與傳統的編譯很不一樣。AI圖包含兩部分：圖拓撲（各層之間如何連線）和模型權重（特定層的引數）。從大小來看，圖拓撲以KB為單位，權重則以MB甚至GB為單位。舉個例子，Meta公司釋出的Open Pre-trained Transformers模型，其引數量從300億、660億到1750億不等，相當於100+GB權重。Gopher和Megatron模型甚至更大。
圖源DeepMind論文
AI生態系統中現有的工具尚不能很好地處理大型模型。比如，Protobufs限制了傳輸資料大小不能超過2GB，因此模型如果使用Protobufs序列化格式，就會備受掣肘。最新版TensorRT的文件中寫道，“對於BERT和GPT等基於Transformer的神經網路模型，TensorRT在編譯時可以消耗10倍於模型大小的CPU記憶體”，可見TensorRT不適合大型模型。如果使用ONNX檔案格式儲存大型模型，就必須將模型權重分成多個檔案分別儲存。
以上種種不但給AI開發工作流增加了不必要的複雜環節，也使模型喪失“單一事實來源”（SSOT），還導致模型分發更加困難。
為了應對模型權重太大的問題，大家可能會採取變通方法，最終卻可能導致整個AI開發工作流變得更復雜。比如，由於某些編譯器階段耗時長達2分多鐘，打斷開發人員的工作節奏，所以Modular構建了一種快取臨時檔案的機制。
雖然這種快取機制和其他變通方法一樣，只是治標不治本：它既非100%可靠，也不能解決Cache Miss（快取缺失）的問題，不過由於Modular十分注重提高開發人員的工作效率，所以還是決定採用這種機制。
2
Modular編譯棧中的MLIR
Modular的技術棧中，MLIR編譯器架構負責表示和轉換AI模型，包括AI運算元圖（用於多種框架）、中級執行時原語和低階機器碼生成。
多級中間表示 (MLIR)
MLIR是LLVM編譯器基礎設施專案的子專案，LLVM旨在提供現代工具包，用以構建針對特定領域的編譯器。MLIR提供一套核心元件，用於硬體設計、量子計算、人工智慧等多種計算領域的建模、分析和轉換。
MLIR能夠幫助構建單個涵蓋全棧的完整系統，比常規的技術棧功能更強大、模組化程度和可擴充性更高，也更易於維護。使用統一的基礎設施讓我們得以便捷地將每一項改進遷移到自己的工具棧，使開發工作流實現更高的模組化和可組裝性。
除了Modular以外，TensorFlow、XLA、PyTorch和ONNX等也在使用MLIR進行模型表示和轉換。隨著MLIR的使用者生態不斷擴大，在讚美MLIR優點的同時，也必須繼續進行改進和完善。
3
MLIR管理權重的方法還有待提高
MLIR的基本組成部分之一是屬性機制（Attribute），可以把它理解為被unique（或被memoize、intern）的常量資料。屬性是使用者可擴充的，也就是說，可以根據不同用例使用不同的屬性型別。很多型別的值都可以被賦予屬性，比如常量表示式值（如“5”、“10.0”等）、字串字面量、列舉值（如“小於”、“大於”、“等於”等），資料組等等。大多數基於MLIR的AI工具都使用屬性來儲存AI模型的權重。
然而，問題出現了：模型權重有可能極其龐大，但MLIR儲存2 GB權重的方式和儲存4 B權重的方式並沒有區別——都使用同一屬性，該屬性包含一組被unique的元素。但對GB級的龐大資料使用unique方法顯然不合理。
這個方法的難點在於：在MLIR中，當某個東西被unique，它就會被分配（allocated）、被hash 、然後被儲存到MLIRContext中。這些東西具有和MLIRContext相同的生命週期，只有當MLIRContext被銷燬，它們才會同時被銷燬。對於小的數值而言，這種機制帶來很多好處，可以把數值傳入傳出，可以透過指標對unique後的值進行比較，還可以共享屬性的記憶體分配（十分常見）等等。
但對數量龐大的權重而言，上述種種好處就變成了劣勢：我們不希望對權重進行重新分配、複製或使用unique方法，我們只需要它們短暫存在——當計算不再需要引用這些權重時，就要允許釋放分配。例如，當執行模型量化工具時，需要對運算元圖進行轉換，並生成新的權重，最終這些權重可能會被複制多份，大量權重副本在編譯結束前都將一直佔用記憶體。
ML工具的另一個問題是MLIR如何序列化至檔案系統。一開始，MLIR沒有二進位制序列化格式，只有文字格式。對數量龐大的權重來說，這就造成問題，因為每個位元組的二進位制資料都會被轉化為十六進位制，後者佔用的空間為前者的2倍。這樣一來，我們不但耗費了相當長的時間進行進位制轉換（一箇中等的GB級模型大概需要20秒），而且轉換後的中間檔案還是原來的2倍大——2倍可是一個不小的數字！
4
記憶體佔用：比拖慢開發效率更嚴重的影響
這一設計機制本意雖好，但它有可能降低編譯效率，即便用最好的編譯器也無濟於事。最明顯的問題是它會導致編譯、監控和轉換模型的時間變長。但凡你曾用過“我的程式碼還在編譯”作為日常摸魚的藉口，你就明白等待編譯是多麼痛苦的事情。採用這一機制，就意味著處理器不得不對GB級資料持續進行分配、複製和hash處理。
XKCD漫畫 – 《還在編譯》
比編譯時長更嚴重的問題是記憶體佔用，它會影響Modular技術棧中的其他架構功能的實現。例如，由於我們的編譯器和技術棧本身都高度並行，而且使用線上搜尋等高階功能，記憶體佔用會直接導致一些工作不能並行展開，也導致不能取得最高質量的結果。
Modular的價值核心是構建使用者喜歡的工具。高階功能如果不好用，或者會影響效率，又或者附帶一些注意事項（比如，“該功能對某些情況不適用”），那麼使用者就根本不會用。因此，Modular致力於解決龐大權重帶來的基礎性問題，簡化使用者的使用流程和開發人員的工作流程。
5
MLIR的核心改進
Modular團隊是MLIR專案的重要貢獻者，Modular企業文化的一大要點是“做對的產品”，Modular參與的所有專案都遵循這一要義。在推動MLIR發展的同時，Modular竭力保證MLIR專案的每一步路都正確，也加強與MLIR社群的合作，為所採取的辦法爭取認可。
Modular團佇列出了大型模型工具應該具備的特點：
- 非必要不分配記憶體： 對大型資料（比如權重）而言，從磁碟中實行記憶體對映比將資料複製到已分配記憶體的block中更高效。
- 無需進行hash或unique處理： 我們不希望費力氣去檢查2 GB Blob資料的相等性；要辨別權重，希望能夠透過名稱辨別，而不是看具體內容有沒有被unique。
- 允許內聯變更（Inline Mutation）： 如果資料只需要在一處使用，應當允許在原位置量化、轉化和運算元據，而不是先複製資料。
- 允許釋放記憶體（deallocation）： 由於大模型的資料量十分龐大，因此當對某一資料的所有引用都不存在時，應當允許釋放記憶體。
- 快速序列化： 無論是即時編譯，搜尋最佳化引數，還是本地迭代，都需要快取IR，所以這一步必須快。
上述觀點並不新穎，但傳統編譯器（比如適用於典型CPU程式語言的編譯器）卻還沒有實現這些要求。
6
調整權重屬性
上述前四點要求解決了我們應該如何使用MLIR這一基本問題：權重雖然是常量資料，但對它的管理應該區別於其他MLIR屬性。一直以來，我們的權重管理方式都很不適宜，這就好比試圖將一塊方釘擠進圓孔中，不僅浪費了空間，降低了我們的開發速度，同時也增加了使用者成本。
所以Modular決定換一種方法來管理權重資料，這促成了MLIR的第一個基本擴充套件機制——“Resource機制”，在計算中將資料和對資料的引用區分開來。
在Resource機制中，序列化MLIR的每個Blob都可能包含額外的資訊段，稱為Resource。Resource要麼是dialect（擴充套件MLIR時使用的類似namespace的抽象），要麼是用於特定工具鏈資料的“外部（external）”資源。Resource中的資料用簡單的鍵值對錶示，創造出如下圖所示的類似json的結構。
/// 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
}
}
#-}上面例子展示瞭如何調整MLIR來用Resource進行復現。MLIR再生器（Reproducer）實際上是一種配置，它包含轉換管道（Transformation Pipeline）等執行資訊，用於復現某種故障或失敗。在使用Resource之前，我們透過在MLIR檔案頂部新增註釋來表示這些執行資訊。現在可以利用Resource將這些執行資訊合併為第一類資訊。
從前需要進行unique處理導致長期佔用記憶體的大型權重資料，現在可以利用Resource機制進行儲存。在IR中，我們對屬性採用輕量級引用而不再採用底層資料：
其他優勢：
- 使用IR進行除錯時更不容易出錯，從而帶來更好的開發體驗： Resource是專門的資訊段；我們不必擔心在除錯時會不小心轉儲整整4GB的資料。
- 我們可以在無需資料的情況下合理地處理IR：因為IR只儲存對資料的引用，不儲存資料本身，如果需要，我們可以省略底層Resource資料。這樣做的好處包括可以極大地簡化再生器生成流程，再生器本來就不需要用到大型權重資料（設想一下，你以前需要向同事傳送1.2GB的再現器檔案，現在的再生器檔案只有20MB大）。
透過引入Resource這個新概念，我們在程式和資料之間建立清晰的分離機制。現在，我們不再將權重資料直接傳遞給某一屬性。相反，我們向屬性傳入一個弱引用，並將權重資料傳給一個專門的管理器。這樣，我們就能更好地控制權重分配、變更和銷燬的時間和方式。
7
新增MLIR二進位制編碼方式
有了更好的權重表示方法之後，下一步我們只需找到更高效的權重儲存方法來完成MLIR表示的序列化。
到此為止，MLIR只有文字序列化格式，這種格式使用ASCII十六進位制字串來表示權重。然而，Modular的終極目標是儘可能加快本地開發流程，因此需要摒棄文字序列化格式，為MLIR新增合適的二進位制格式（https://discourse.llvm.org/t/...）。
二進位制格式需要考慮很多因素，況且二進位制格式決定了編譯器的穩定性。MLIR需要高度的靈活性才能高效應對各種各樣的用例，需要實現高速度，而且MLIR/LLVM不能依賴第三方編碼庫。
不過，MLIR的一大好處是編碼難度極低。因為MLIR中所有操作的結構都相同，所有操作都可以使用相同的編碼方式。上述的種種複雜要求都是為了保證MLIR核心概念的緊湊和高效。考慮到這些限制，我們決定為MLIR定製編碼方式（https://mlir.llvm.org/docs/By...）。
8
使用者收益
為MLIR增加Resource機制和二進位制編碼方式大大加速了工具鏈和開發流程，並大幅降低記憶體佔用，提高了效能和速度表現，也整體改善了MLIR。
為了驗證上述改進帶來的效能變化，可以測試不同規模的模型上基於MLIR的圖編譯器中“降級”和“最佳化”步驟的實際速度（將TensorFlow序列化模型轉化為符合MLIR執行時輸入格式的模型），以及該過程中的實際記憶體佔用。
速度提升：編譯工作流
測試結果發現，MLIR的速度大幅提升。從TensorFlow序列化模型（TensorFlow 2.10模型）轉化為MLIR執行時輸入格式的模型，這一過程涉及大量底層表示轉換，經過改進後，實際執行時間縮短了1.8~2倍，執行速度隨模型大小按比例縮放。
具體而言，處理TensorFlow序列化模型耗時極短——生成MLIR時將大量權重資料寫入磁碟這一步驟是主要的耗時來源。經改進後，程式碼處理時間比原來快10倍，整體執行時間的快慢主要取決於固態硬碟（SSD）將 >1 GB資料寫入磁碟的速度。
ML開發人員使用我們的工具，可以加快模型編譯速度，從而提升生產效率，減少迭代時間。我們的工具可以最佳化生產環境以及模型的載入和編譯，包括基於流入資料的動態模型載入和解除安裝，以及各種個性化或經過精細化調整的使用者模型。
速度提升：序列化
引入二進位制編碼方式不但可以加快編譯工作流，還能加快序列化速度。透過外部工具與MLIR進行互動，包括執行時型別檢查（Introspection）、快取和再生器生成等，都需要對序列化MLIR進行讀寫。
透過對不同規模的模型進行了序列化測試，結果同樣發現峰值效能大幅提速，且SSD寫入步驟依然為主要耗時來源。具體而言，大型模型文字資料的讀取耗時約5秒，而二進位制資料的讀取耗時僅不到10毫秒；二進位制格式的寫入速度則約是文字格式資料的5倍。
對Modular而言，引入二進位制編碼方式可以加快以MLIR為中心的基礎設施和工具的開發速度，改善原本高成本、低速度的開發狀況。比如，偵錯程式（Debugger）的效率很大程度取決於編譯工作流中快取模型表示的效率，而引入二進位制編碼方式可以提高偵錯程式的效率，從而提高底層編譯器的效能。
記憶體佔用
二進位制序列化格式的mmap（一種記憶體對映方法）效能以及透過Resource機制實現的IR和資料的相互獨立性可以大幅減少記憶體佔用。測試發現，各種規模的模型編譯流程中的記憶體佔用都大大降低——因為以前需要為模型權重分配記憶體，現在不需要了。
9
升級AI生態
Modular的願景不只是為了方便我們自己，而是升級整個AI行業的生態。前文提及的新型Resource表示和二進位制編碼方式都已提交至上游的LLVM/MLIR倉庫中。
Modular起初的研發動機是為了解決Modular的客戶遇到的問題並提升自身內部基礎設施，但這些改進產生的積極影響並不限於自己的用例，還能改善其他以MLIR為基礎技術的產品。例如，由於二進位制編碼方式的引進，MLIR社群如今正在討論如何保證MLIR的穩定性（https://discourse.llvm.org/t/...）。
這些基礎技術的改進最終都會融入產品中為使用者服務。以上只是Modular致力提升的無數核心技術之一。Modular一方面竭力適應大模型，一方面努力完善模型在裝置上的部署，目標都是大幅提升AI基礎設施的效能和易用性。Modular非常看好AI的未來以及LLVM和MLIR的發展。
（本文由OneFlow社群翻譯，譯文轉載請聯絡OneFlow獲得授權。原文：1. https://www.modular.com/blog/...；2.https://www.modular.com/blog/...）
歡迎下載體驗 OneFlow v0.8.0 最新版本：https://github.com/Oneflow-In...