在去年 10 月的 Google Cloud Next 2022 活動中，OpenXLA 項目正式浮出水面，谷歌與包括阿里巴巴、AMD、Arm、亞馬遜、英特爾、英偉達等科技公司推動的開源 AI 框架合作，致力于匯集不同機器學習框架，讓機器學習開發(fā)人員獲得能主動選擇框架、硬件的能力。

本周三，谷歌宣布 OpenXLA 項目正式開源。

項目鏈接：https://github.com/openxla/xla

通過創(chuàng)建與多種不同機器學習框架、硬件平臺共同工作的統(tǒng)一機器學習編譯器，OpenXLA 可以加速機器學習應用的交付并提供更大的代碼可移植性。對于 AI 研究和應用來說，這是一個意義重大的項目，Jeff Dean 也在社交網(wǎng)絡上進行了宣傳。

如今，機器學習開發(fā)和部署受到碎片化的基礎設施的影響，這些基礎設施可能因框架、硬件和用例而異。這種相互隔絕限制了開發(fā)人員的工作速度，并對模型的可移植性、效率和生產(chǎn)化造成了障礙。

3 月 8 日，谷歌等機構通過 OpenXLA 項目（其中包括 XLA、StableHLO 和 IREE 存儲庫）的開放，朝著消除這些障礙邁出了重要一步。

OpenXLA 是由 AI / 機器學習行業(yè)領導者共同開發(fā)的開源 ML 編譯器生態(tài)系統(tǒng)，貢獻者包括阿里巴巴、AWS、AMD、蘋果、Arm、Cerebras、谷歌、Graphcore、Hugging Face、英特爾、Meta 和英偉達。它使得開發(fā)人員能夠編譯和優(yōu)化來自所有領先機器學習框架的模型，以便在各種硬件上進行高效訓練和服務。使用 OpenXLA 的開發(fā)人員可以觀察到訓練時間、吞吐量、服務延遲以及最終發(fā)布和計算成本方面的明顯提升。

機器學習技術設施面臨的挑戰(zhàn)

隨著 AI 技術進入實用階段，許多行業(yè)的開發(fā)團隊都在使用機器學習來應對現(xiàn)實世界的挑戰(zhàn)，例如進行疾病的預測和預防、個性化學習體驗和黑洞物理學探索。

隨著模型參數(shù)數(shù)量呈指數(shù)級增長，深度學習模型所需的計算量每六個月翻一番，開發(fā)人員正在尋求基礎架構的最大性能和利用率。大量團隊正在利用多型號種類的硬件，從數(shù)據(jù)中心中的節(jié)能機器學習專用 ASIC 到可以提供更快響應速度的 AI 邊緣處理器。相應的，為了提高效率，這些硬件設備使用定制化的獨特算法和軟件庫。

但另一方面，如果沒有通用的編譯器將不同硬件設備橋接到當今使用的多種框架（例如 TensorFlow、PyTorch）上，人們就需要付出大量努力才能有效地運行機器學習。在實際工作中，開發(fā)人員必須手動優(yōu)化每個硬件目標的模型操作。這意味著使用定制軟件庫或編寫特定于設備的代碼需要領域專業(yè)知識。

這是一個矛盾的結果，為了提高效率使用專用技術，結果卻是跨框架和硬件的孤立、不可概括的路徑導致維護成本高，進而導致供應商鎖定，減緩了機器學習開發(fā)的進度。

解決方法和目標

OpenXLA 項目提供了最先進的 ML 編譯器，可以在 ML 基礎設施的復雜性中進行擴展。它的核心支柱是性能、可擴展性、可移植性、靈活性和易用性。借助 OpenXLA，我們渴望通過加速人工智能的開發(fā)和交付來實現(xiàn) AI 在現(xiàn)實世界中的更大潛力。

OpenXLA 的目標在于：

• 通過適用于任何框架，接入專用設備后端和優(yōu)化的統(tǒng)一編譯器 API，使開發(fā)人員可以輕松地在他們的首選框架中針對各種硬件編譯和優(yōu)化任何模型。
• 為當前和新興模型提供行業(yè)領先的性能，也可擴展至多個主機和加速器滿足邊緣部署的限制，并推廣到未來的新型模型架構上。
• 構建一個分層和可擴展的機器學習編譯器平臺，為開發(fā)人員提供基于 MLIR 的組件，這些組件可針對其獨特的用例進行重新配置，用于硬件定制化編譯流程。

AI/ML 領導者社區(qū)

我們今天在機器學習基礎架構中面臨的挑戰(zhàn)是巨大的，沒有任何一個組織可以單獨有效地解決這些挑戰(zhàn)。OpenXLA 社區(qū)匯集了在 AI 堆棧的不同級別（從框架到編譯器、runtime 和芯片）上運行的開發(fā)人員和行業(yè)領導者，因此非常適合解決我們在 ML 領域看到的碎片化問題。

作為一個開源項目，OpenXLA 遵循以下原則：

• 平等地位：個人無論從屬關系如何，都平等地做出貢獻。技術領導者是那些貢獻最多時間和精力的人。
• 尊重文化：所有成員都應維護項目價值觀和行為準則，無論他們在社區(qū)中的職位如何。
• 可擴展、高效的治理：小團隊做出基于共識的決策，具有清晰但很少使用的升級路徑。
• 透明度：所有決定和理由都應該對公眾清晰可見。

OpenXLA 生態(tài)系統(tǒng)：性能、規(guī)模和可移植能力

OpenXLA 通過模塊化工具鏈消除了機器學習開發(fā)人員的障礙，它通過通用編譯器接口得到所有領先框架的支持，利用可移植的標準化模型表示，并提供具有強大的目標向和特定硬件優(yōu)化的特定領域編譯器。該工具鏈包括 XLA、StableHLO 和 IREE，所有這些工具都利用 MLIR：一種編譯器基礎架構，使機器學習模型能夠在硬件上一致地表示、優(yōu)化和執(zhí)行。

OpenXLA 主要亮點

機器學習用例的范圍

OpenXLA 當前的使用涵蓋了 ML 用例的范圍，包括在阿里云上對 DeepMind 的 AlphaFold、GPT2 和 Swin Transformer 等模型進行全面訓練，以及在 Amazon.com 上進行多模態(tài) LLM 訓練。Waymo 等客戶利用了 OpenXLA 進行車載實時推理。此外，OpenXLA 還用于優(yōu)化配備 AMD RDNA? 3 的本地機器上的 Stable Diffusion 服務。

最佳性能，開箱即用

OpenXLA 使開發(fā)人員無需編寫特定于設備的代碼，即可輕松加快模型性能。它具有整體模型優(yōu)化功能，包括簡化代數(shù)表達式、優(yōu)化內(nèi)存數(shù)據(jù)布局以及改進調度以減少峰值內(nèi)存使用和通信開銷。高級算子融合和內(nèi)核生成有助于提高設備利用率并降低內(nèi)存帶寬要求。

輕松擴展工作負載

開發(fā)高效的并行化算法非常耗時并且需要專業(yè)知識。借助 GSPMD 等功能，開發(fā)人員只需注釋關鍵張量的一個子集，然后編譯器就可以使用這些子集自動生成并行計算。這消除了跨多個硬件主機和加速器對模型進行分區(qū)和高效并行化所需的大量工作。

便攜性和可選性

OpenXLA 為多種硬件設備提供開箱即用的支持，包括 AMD 和 NVIDIA GPU、x86 CPU 和 Arm 架構以及 ML 加速器，如 Google TPU、AWS Trainium 和 Inferentia、Graphcore IPU、Cerebras Wafer-Scale Engine 等等。OpenXLA 還通過 StableHLO 支持 TensorFlow、PyTorch 和 JAX，StableHLO 是一個用作 OpenXLA 輸入格式的可移植層。

靈活性

OpenXLA 為用戶提供了手動調整模型熱點的靈活性。自定義調用等擴展機制使用戶能夠用 CUDA、HIP、SYCL、Triton 和其他內(nèi)核語言編寫深度學習原語，從而能夠充分利用硬件特性。

StableHLO

StableHLO 是 ML 框架和 ML 編譯器之間的一個可移植層，是一個支持動態(tài)、量化和稀疏性的高級運算（HLO）的運算集。此外，它可以被序列化為 MLIR 字節(jié)碼以提供兼容性保證。所有主要的 ML 框架（JAX、PyTorch、TensorFlow）都可以產(chǎn)生 StableHLO。2023 年，谷歌計劃與 PyTorch 團隊緊密合作，實現(xiàn)與 PyTorch 2.0 版本的整合。