作者｜River Riddle、Eric Johnson、Abdul Dakak翻譯｜胡燕君、楊婷
機器學習模型逐漸發展成人們口中的“龐然大物”。全球頂尖的科技公司紛紛踏上“軍備競賽”之路，立志訓練出規模最大的模型（MUM、OPT、GPT-3、Megatron），而其他專注于生產系統的公司也相繼擴大其原有模型，并取得良好成果。
一切如火如荼，然而，鮮少有人提及，龐大的模型給現有的AI基礎設施和開發流程帶來了諸多實際性挑戰。
(資料圖片僅供參考)
大模型的權重可達100+GB，而目前的開發工具卻還沒跟上，使用起來十分費力，部署時往往要等上好幾分鐘甚至好幾小時，這已經成為AI工程師的隱痛，不但浪費工程師的時間，降低工作效率，還會拖慢迭代速度。
致力于AI基礎設施工具研發的Modular團隊認為，開發人員的工作效率是訓練和部署模型的最大成本之一。因此需要不斷優化工具鏈，提升早期用戶的體驗，也方便開發人員。本文探討編譯過程中管理海量數據的技術難點，以及Modular為解決這些難點在基礎設施（以及MLIR編譯器框架）方面所做的改進。由OneFlow社區（ID：OneFlowTechnology）編譯。?
1AI模型配套工具的易用性不足
機器學習中的圖轉換（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十分注重提高開發人員的工作效率，所以還是決定采用這種機制。
2Modular編譯棧中的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致力于解決龐大權重帶來的基礎性問題，簡化用戶的使用流程和開發人員的工作流程。
5MLIR的核心改進
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/rfc-a-binary-serialization-format-for-mlir/63518）。
二進制格式需要考慮很多因素，況且二進制格式決定了編譯器的穩定性。MLIR需要高度的靈活性才能高效應對各種各樣的用例，需要實現高速度，而且MLIR/LLVM不能依賴第三方編碼庫。
不過，MLIR的一大好處是編碼難度極低。因為MLIR中所有操作的結構都相同，所有操作都可以使用相同的編碼方式。上述的種種復雜要求都是為了保證MLIR核心概念的緊湊和高效。考慮到這些限制，我們決定為MLIR定制編碼方式（https://mlir.llvm.org/docs/BytecodeFormat/）。
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/mlir-generic-ir-stability-and-upgradability/65371）。
這些基礎技術的改進最終都會融入產品中為用戶服務。以上只是Modular致力提升的無數核心技術之一。Modular一方面竭力適應大模型，一方面努力完善模型在設備上的部署，目標都是大幅提升AI基礎設施的性能和易用性。Modular非常看好AI的未來以及LLVM和MLIR的發展。 （本文由OneFlow社區翻譯，譯文轉載請聯系OneFlow獲得授權。原文：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之父：編譯器的黃金時代
李白：你的模型權重很不錯，可惜被我沒收了
更快的YOLOv5問世，附送全面中文解析教程
LLVM之父：為什么我們要重建AI基礎設施軟件
LLVM之父：模塊化設計決定AI前途，不服來辯
開源吞噬AI界？從Stable Diffusion的爆火說起
OneEmbedding:單卡訓練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
關鍵詞： 基礎設施
營業執照公示信息