盡管AI的發(fā)展取得了巨大進(jìn)步，但編譯器LLVM之父Chris Lattner認(rèn)為，AI技術(shù)應(yīng)用并不深入，遠(yuǎn)遠(yuǎn)沒有發(fā)揮出已有機(jī)器學(xué)習(xí)研究的所有潛力。而AI系統(tǒng)和工具的單一化和碎片化正是造成這一問題的根源。 為了讓AI發(fā)揮其真正的潛力，計(jì)算碎片化是需要解決的重點(diǎn)問題之一，目標(biāo)是讓AI軟件開發(fā)人員能夠無縫地充分利用現(xiàn)有硬件和下一代創(chuàng)新硬件。但解決這一問題并不容易，硬件、模型和數(shù)據(jù)的多樣性使得當(dāng)前市場(chǎng)上的現(xiàn)有解決方案都只是單點(diǎn)性質(zhì)的，Chris Lattner創(chuàng)立的Modular團(tuán)隊(duì)從矩陣算法的角度對(duì)此進(jìn)行了深入分析。

(資料圖片)

（以下內(nèi)容由OneFlow編譯發(fā)布，譯文轉(zhuǎn)載請(qǐng)聯(lián)系OneFlow獲得授權(quán)。https://www.modular.com/blog/ais-compute-fragmentation-what-matrix-multiplication-teaches-us）

作者｜Eric Johnson、Abdul Dakkak、Chad Jarvis

OneFlow編譯翻譯｜徐佳渝、楊婷

1

算力碎片化正在阻礙AI的發(fā)展

AI由數(shù)據(jù)、算法（即模型）和算力驅(qū)動(dòng)，三者之間形成了良性循環(huán)。其中任意一方的發(fā)展會(huì)推動(dòng)其他方面需求的增長(zhǎng)，從而嚴(yán)重影響開發(fā)者在可用性和性能等方面的體驗(yàn)。如今，我們擁有更多的數(shù)據(jù)，做了更多的AI模型研究，但算力的擴(kuò)展速度卻沒有跟上，這主要是由于物理限制。

如果你一直在關(guān)注AI和硬件的發(fā)展，可能聽說過摩爾定律時(shí)代即將結(jié)束。過去60年，單核處理器每18個(gè)月翻一倍性能提升速度的情況已然改變。除了繼續(xù)制造越來越小的晶體管的物理限制之外（例如，電流泄漏會(huì)導(dǎo)致功耗過高，從而引起發(fā)熱），性能也越來越多地受到內(nèi)存延遲的限制，而這種限制的增長(zhǎng)速度比處理速度要緩慢得多。

（Hennessy和Patterson的圖靈演講：不同時(shí)期CPU處理器性能提升的分析（性能提升的速度保持穩(wěn)定））

然而，隨著模型不斷擴(kuò)大，在邊緣計(jì)算中創(chuàng)建和處理的企業(yè)數(shù)據(jù)更多，對(duì)AI計(jì)算的需求也在不斷增加。因此，盡可能地利用硬件設(shè)備的性能已成為業(yè)界的關(guān)注焦點(diǎn)。

（?機(jī)器學(xué)習(xí)三個(gè)時(shí)代的算力走向（https://arxiv.org/pdf/2202.05924.pdf），Sevilla：計(jì)算需求隨時(shí)間呈對(duì)數(shù)方式變化的分析。其中，2010年左右深度學(xué)習(xí)開始受到關(guān)注，2016年左右迎來了大型模型時(shí)代，這促使了算力需求急劇增長(zhǎng)。）

那么，算力的碎片化是如何阻礙AI發(fā)展的呢？由于傳統(tǒng)CPU無法擴(kuò)展以滿足更多的算力需求，因此，唯一的解決之道是創(chuàng)建并行的、用于特定領(lǐng)域的硬件平臺(tái)，盡管這些硬件平臺(tái)的通用性不強(qiáng)，但在特定的AI領(lǐng)域卻表現(xiàn)良好——例如圖形處理單元(GPU)、張量處理單元(TPU)和其他專用集成電路(ASIC)。 雖然這些創(chuàng)新推動(dòng)了AI行業(yè)的發(fā)展，讓邊緣設(shè)備能夠使用規(guī)模更大、效率更高的處理器，但硬件的多樣性使得整個(gè)AI行業(yè)變得碎片化，AI開發(fā)者們需要努力解決以下問題：

1. 開發(fā)出能夠充分利用硬件能力的軟件，且能與其他軟件協(xié)同工作。

2. 在任意一款設(shè)備上實(shí)現(xiàn)并行軟件算法。

3. 將軟件擴(kuò)展到多設(shè)備的生態(tài)系統(tǒng)，甚至擴(kuò)展到異構(gòu)系統(tǒng)。

Modular公司致力于從零開始重建全球的AI基礎(chǔ)設(shè)施。在本系列博客中，我們將探討如何采用全新的方法解決AI行業(yè)的算力碎片化問題。我們會(huì)專注于單個(gè)運(yùn)算符——矩陣乘法（matrix multiplication，matmul），這是機(jī)器學(xué)習(xí)算法中的關(guān)鍵計(jì)算。 通過這種方式，我們會(huì)看到構(gòu)建真正統(tǒng)一的解決方案所面臨的底層挑戰(zhàn)。我們將深入研究矩陣乘法的內(nèi)部運(yùn)作，探討其工作原理的一些細(xì)節(jié)，以了解矩陣乘法為什么如此困難。

2

矩陣乘法為何如此困難

矩陣對(duì)機(jī)器學(xué)習(xí)系統(tǒng)來說至關(guān)重要，因?yàn)樗峁┝艘环N簡(jiǎn)單而高效的方式來表示數(shù)據(jù)。例如，輸入數(shù)據(jù)（圖像中的像素集合）或者模型內(nèi)部不同層之間的運(yùn)作機(jī)制都可以用矩陣來表示。因此，矩陣相乘的運(yùn)算在深度學(xué)習(xí)模型總計(jì)算量中占據(jù)很大比例。 實(shí)際上，在許多當(dāng)前流行的Transformer模型BERT、CLIP以及ChatGPT中，矩陣乘法的運(yùn)行時(shí)長(zhǎng)約占其總運(yùn)行時(shí)長(zhǎng)的45-60%。矩陣乘法在計(jì)算卷積運(yùn)算中扮演著重要角色，該運(yùn)算是大多數(shù)計(jì)算機(jī)視覺模型的基礎(chǔ)，也是許多高性能計(jì)算應(yīng)用的核心。

考慮到矩陣乘法在機(jī)器學(xué)習(xí)和計(jì)算機(jī)視覺等領(lǐng)域的重要性，研究人員對(duì)其進(jìn)行了廣泛的算法研究，以編寫出高效的矩陣乘法算法。自60年代（https://ieeexplore.ieee.org/abstract/document/1687427）、70年代（https://apps.dtic.mil/sti/pdfs/AD0705509.pdf）、80年代

（https://dl.acm.org/doi/10.1145/356012.356020）,?90年代（https://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.54.9411&rep=rep1&type=pdf）、21世紀(jì)初至今（https://netlib.org/lapack/lawnspdf/lawn147.pdf），一直有論文在嘗試使用當(dāng)時(shí)的硬件來解決該問題。

簡(jiǎn)單的O(n^3)矩陣乘法算法。

但是概念上的矩陣乘法并不是難點(diǎn)。相反，挑戰(zhàn)在于編寫一個(gè)足夠高效的矩陣乘法，以在AI行業(yè)的所有硬件、模型和數(shù)據(jù)多樣性上實(shí)現(xiàn)SOTA性能。此外，與其他AI算子的協(xié)同工作更具挑戰(zhàn)性。

硬件

每種用于運(yùn)行AI模型的設(shè)備都有其獨(dú)特的特征，例如內(nèi)存層次結(jié)構(gòu)不同、乘法和累加單元（MAC）不同等等。

例如，CPU采用了這種內(nèi)存層次結(jié)構(gòu)——從慢速的RAM到越來越快的緩存，包括Level-3、Level-2、Level-1和CPU寄存器。內(nèi)存大小與速度成反比，例如，L1緩存的訪問速度通常為1納秒量級(jí)，而RAM的訪問速度為100納秒量級(jí)。 要獲得最高性能的矩陣運(yùn)算，必須讓算法有效地處理不同的內(nèi)存級(jí)別和大小。因?yàn)樵季仃囂?，所以無法將其一次性放入寄存器或最快的內(nèi)存緩存中。因此，挑戰(zhàn)在于將原始矩陣分解為大小合適的block或tiles，以最大限度地利用最快內(nèi)存。

此外，處理核心矩陣功能（core matrix functionality）單元的實(shí)際形狀因硬件而異。就傳統(tǒng)意義而言，CPU是一種標(biāo)量（scalar）機(jī)器，這意味著它需要逐個(gè)處理指令。然而，在過去二十年中，所有CPU供應(yīng)商都增加了向量單元（如SIMD），GPU則采用了SIMT（單指令多線程，Single Instruction, Multiple Threads）的方式，以最大限度地提高高度并行、重復(fù)運(yùn)算的效率。

此外，更專業(yè)的硬件通過對(duì)二維矩陣進(jìn)行運(yùn)算來進(jìn)一步實(shí)現(xiàn)這一點(diǎn)。最著名的是Google的TPU，不過蘋果和英特爾已經(jīng)增加了自己的矩陣乘法功能AMX。雖然更先進(jìn)的MAC單元提高了性能，但也產(chǎn)生了相應(yīng)的需求，即能夠在標(biāo)量、向量和矩陣處理器上工作的靈活算法。

（深入了解谷歌的第一個(gè)張量處理單元：不同乘法和累加（MAC）單元形狀。）

模型

AI模型具有多樣性。雖然矩陣乘法是許多模型的基礎(chǔ)，但它們的矩陣大小可能差異巨大。例如，模型具有不同的輸入形狀（如不同的序列長(zhǎng)度）、不同的內(nèi)部形狀（即作為模型隱藏層一部分的相乘矩陣），以及不同的batch size大小（對(duì)于訓(xùn)練和推理效率至關(guān)重要）。因此，矩陣乘法在生產(chǎn)中有數(shù)百種不同的形狀，使得難以將它們分解為不同的block，以實(shí)現(xiàn)內(nèi)存效率最大化。

（Transformer圖解，Jay Alammar：多頭注意力塊中涉及的各種矩陣大小，Transformer模型（如BERT, GPT2和CLIP）的關(guān)鍵構(gòu)建塊。）

數(shù)據(jù)

數(shù)據(jù)也存在差異。大多數(shù)讀者可能比較熟悉結(jié)構(gòu)化和非結(jié)構(gòu)化數(shù)據(jù)方面的差異，但在本文，我們重點(diǎn)關(guān)注數(shù)據(jù)類型（“dtype”）。在AI模型中，數(shù)據(jù)通常使用dtype FP32，但為了減少模型大小，提高模型性能，業(yè)界也接受較低精度的數(shù)據(jù)類型，如Bfloat16，Int8以及更特殊的FP4和Int4。用例不同，Matmul算法需要運(yùn)行的數(shù)據(jù)精度也不同。

量化在AI中的重要性：將FP32量化為Int8。

3

當(dāng)前的SOTA矩陣乘法算法

那么，目前的SOTA矩陣乘法算法是如何實(shí)現(xiàn)的呢？鑒于其重要性，matmul通常是硬件供應(yīng)商最先優(yōu)化的算法之一，一般來說，供應(yīng)商會(huì)利用他們的庫(kù)來實(shí)現(xiàn)優(yōu)化，比如英特爾的MKL和OneDNN

庫(kù)、AMD的AOCL和RocBLAS、ARM的performance庫(kù)、Apple的Accelerate、Nvidia的CUBLAS。

就效率而言，在上述提及的硬件庫(kù)中，當(dāng)前的SOTA是有效地編寫匯編代碼：在最低層級(jí)向硬件提供直接指令來微觀管理硬件，而無需抽象。

面向AVX512，用x86匯編編寫的矩陣算法

這樣做的主要原因是什么？編寫匯編可以實(shí)現(xiàn)任意特定用例的最佳性能，通過匯編編寫，開發(fā)人員可以避免編譯器的不可預(yù)測(cè)性，編譯器將Python和C++等高級(jí)語(yǔ)言轉(zhuǎn)換為匯編語(yǔ)言，可以執(zhí)行編譯器難以做到的優(yōu)化，因?yàn)榫幾g器必須泛化。更重要的是，它們可以利用編譯器未意識(shí)到的指令和模式，因?yàn)閿U(kuò)展編譯器以支持新的硬件功能需要時(shí)間。

4

手寫匯編kernel無法擴(kuò)展

那么，匯編能否真正解決用戶的碎片問題呢？雖然匯編編寫可以最大限度地提高個(gè)人用例性能，但它不可移植、組合、擴(kuò)展且對(duì)用戶不夠友好。少數(shù)手工編寫和微調(diào)匯編代碼的專家怎么可能將他們的工作擴(kuò)展到不同配置，同時(shí)整合到所有AI框架中呢？這簡(jiǎn)直就是不可能完成的任務(wù)。

（機(jī)器學(xué)習(xí)編譯優(yōu)化器介紹，Chip Huyen：不斷增加的框架和硬件支持組合）

可移植性

匯編是使用特定于硬件的接口（指令集架構(gòu)ISA）編寫的，因此它不能跨硬件移植。事實(shí)上，匯編甚至無法在同一硬件供應(yīng)商的多代芯片上實(shí)現(xiàn)最佳性能！

此外，即使我們?cè)陂_發(fā)模型時(shí)考慮了目標(biāo)硬件，也仍然存在兩大實(shí)際問題：

1. 首先，我們無法在云上控制運(yùn)行匯編的特定硬件。你可能會(huì)說“我選擇了一個(gè)非常適合自己模型的實(shí)例”。但事實(shí)是，某些云提供商（如AWS）上的實(shí)例不能保證特定的CPU類型。比如，如果你選擇了c5.4xlarge實(shí)例，那么相應(yīng)的，你只能使用老一代英特爾SkyLake處理器或新一些的級(jí)聯(lián)湖處理器（Cascade Lake processor）。匯編無法適應(yīng)運(yùn)行代碼的特定芯片，更無法為其提供最佳性能。

2. 你的產(chǎn)品將繼續(xù)迅速迭代，你可能希望完全遷移到不同的硬件架構(gòu)。確定一種特定配置會(huì)限制你根據(jù)模型需求的變化或新一代硬件進(jìn)行適應(yīng)性調(diào)整的靈活性。

可擴(kuò)展性和可組合性

如前所述，AI模型的多樣性導(dǎo)致了矩陣形狀差異。使用基于匯編的庫(kù)意味著選擇對(duì)內(nèi)存tile大小等參數(shù)進(jìn)行硬編碼的特定處理器指令。這些硬編碼的匯編庫(kù)可以針對(duì)特定的張量形狀進(jìn)行優(yōu)化，但需要對(duì)其他張量形狀進(jìn)行不同的實(shí)現(xiàn)。 因此，許多現(xiàn)有的kernel庫(kù)膨脹到了千兆字節(jié)（例如，MKL為3.2GB，cuDNN最高可達(dá)2.5GB）。當(dāng)這些庫(kù)的大小影響到容器構(gòu)建時(shí)間時(shí)，如果你想部署到不切實(shí)際的邊緣端，或者想要部署供應(yīng)商尚未手動(dòng)優(yōu)化的新創(chuàng)新和研究時(shí)，這將成為一個(gè)麻煩。

從整體來看，高性能矩陣乘法對(duì)性能確實(shí)很重要。但為了獲得最佳結(jié)果，矩陣乘法運(yùn)算可以與其他運(yùn)算一起執(zhí)行，例如Elementwise、跨步訪問（strided accesses）、廣播（broadcasts）等。通過減少內(nèi)存數(shù)據(jù)流量，算子融合（Operator fusion）帶來了顯著的性能改進(jìn)，但問題是有成千上萬(wàn)的AI算子。 此外，模型使用了很多不同運(yùn)算的排列組合，手動(dòng)融合所有重要組合是不切實(shí)際的（雖然有些人已經(jīng)嘗試過了?。?，特別是在AI高速發(fā)展的情況下。

用戶友好性

最后，匯編對(duì)于用戶來說并不友好，不利于跨組織的生產(chǎn)力。匯編編程在現(xiàn)代編程語(yǔ)言（如參數(shù)化和面向?qū)ο缶幊蹋┲锌捎玫墓δ苡邢?，且不能在調(diào)試（debugging）、代碼覆蓋率（code coverage）、測(cè)試（testing）等方面提供有效助益。 雖然大多數(shù)想要編寫新運(yùn)算符的研究人員都對(duì)Python比較滿意，但還是有一些眾所周知的性能問題（https://en.wikipedia.org/wiki/Global_interpreter_lock）。為了應(yīng)對(duì)這些問題，組織不得不花高價(jià)高薪聘請(qǐng)專家來填補(bǔ)鴻溝。 ?

其他人都在看

“ChatGPT們”的淘金時(shí)代

機(jī)器學(xué)習(xí)編譯器的前世今生

手把手推導(dǎo)分布式矩陣乘的最優(yōu)并行策略

谷歌科學(xué)家：ChatGPT秘密武器的演進(jìn)與局限

LLVM之父Chris Lattner：編譯器的黃金時(shí)代

兩大圖靈獎(jiǎng)得主力作：計(jì)算機(jī)架構(gòu)的新黃金時(shí)代

GLM訓(xùn)練加速：性能最高提升3倍，顯存節(jié)省1/3

歡迎Star、試用OneFlow: github.com/Oneflow-Inc/oneflow/http://github.com/Oneflow-Inc/oneflow/

關(guān)鍵詞：