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

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

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

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

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

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

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

機器學習技術設施面臨的挑戰

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

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

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

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

解決方法和目標

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

OpenXLA 的目標在于：

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

AI/ML 領導者社區

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

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

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

OpenXLA 生態系統：性能、規模和可移植能力

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

OpenXLA 主要亮點

機器學習用例的范圍

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

最佳性能，開箱即用

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

輕松擴展工作負載

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

便攜性和可選性

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 和其他內核語言編寫深度學習原語，從而能夠充分利用硬件特性。

StableHLO

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