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

本周三，谷歌宣布 OpenXLA 項(xiàng)目正式開源。

項(xiàng)目鏈接：https://github.com/openxla/xla

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

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

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

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

機(jī)器學(xué)習(xí)技術(shù)設(shè)施面臨的挑戰(zhàn)

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

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

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

這是一個(gè)矛盾的結(jié)果，為了提高效率使用專用技術(shù)，結(jié)果卻是跨框架和硬件的孤立、不可概括的路徑導(dǎo)致維護(hù)成本高，進(jìn)而導(dǎo)致供應(yīng)商鎖定，減緩了機(jī)器學(xué)習(xí)開發(fā)的進(jìn)度。

解決方法和目標(biāo)

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

OpenXLA 的目標(biāo)在于：

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

AI/ML 領(lǐng)導(dǎo)者社區(qū)

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

作為一個(gè)開源項(xiàng)目，OpenXLA 遵循以下原則：

• 平等地位：個(gè)人無論從屬關(guān)系如何，都平等地做出貢獻(xiàn)。技術(shù)領(lǐng)導(dǎo)者是那些貢獻(xiàn)最多時(shí)間和精力的人。
• 尊重文化：所有成員都應(yīng)維護(hù)項(xiàng)目價(jià)值觀和行為準(zhǔn)則，無論他們在社區(qū)中的職位如何。
• 可擴(kuò)展、高效的治理：小團(tuán)隊(duì)做出基于共識(shí)的決策，具有清晰但很少使用的升級路徑。
• 透明度：所有決定和理由都應(yīng)該對公眾清晰可見。

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

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

OpenXLA 主要亮點(diǎn)

機(jī)器學(xué)習(xí)用例的范圍

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

最佳性能，開箱即用

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

輕松擴(kuò)展工作負(fù)載

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

便攜性和可選性

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

靈活性

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

StableHLO

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