AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報導了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯絡報導。投稿郵箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com
在人工智慧領域，模型引數的增多往往意味著效能的提升。但隨著模型規模的擴大，其對終端裝置的算力與記憶體需求也日益增加。低位元量化技術，由於可以大幅降低儲存和計算成本並提升推理效率，已成為實現大模型在資源受限裝置上高效執行的關鍵技術之一。然而，如果硬體裝置不支援低位元量化後的資料模式，那麼低位元量化的優勢將無法發揮。
為了解決這一問題，微軟亞洲研究院推出了全新的資料編譯器 Ladder 和演算法 T-MAC，使當前只支援對稱精度計算的硬體能夠直接執行混合精度矩陣乘法。測試結果表明，Ladder 在支援 GPU 原本不支援的自定義資料型別方面，最高提速可達 14.6 倍；T-MAC 在搭載了最新高通 Snapdragon X Elite 晶片組的 Surface AI PC 上，使 CPU 上執行的大模型吞吐率比專用加速器 NPU 快兩倍。此外，研究員們還設計了 LUT Tensor Core 硬體架構，這種精簡設計使硬體能夠直接支援各種低位元混合精度計算，為人工智慧硬體設計提供了新思路。
大模型已經越來越多地被部署在智慧手機、膝上型電腦、機器人等端側裝置上，以提供先進的智慧及實時響應服務。但包含上億引數的大模型對終端裝置的記憶體和計算能力提出了極高的要求，也因此限制了它們的廣泛應用。低位元量化技術因其能顯著壓縮模型規模，降低對計算資源的需求，成為了大模型在端側部署和實現高效推理的有效手段。
隨著低位元量化技術的發展，資料型別日益多樣化，如 int4、int2、int1 等低位元資料，使得大模型在推理中越來越多地採用低位元權重和高位元權重計算的混合精度矩陣乘法（mixed-precision matrix multiplication，mpGEMM）。然而，現有的 CPU、GPU 等硬體計算單元通常只支援對稱計算模式，並不相容這種混合精度的矩陣乘法。
混合精度矩陣乘法與傳統的矩陣乘法有何不同？
在傳統的矩陣乘法中，參與運算的兩端數值是對稱的，例如 FP16*FP16、int8*int8。但大模型的低位元量化打破了這種對稱性，使乘法的一端是高位元，另一端是低位元，例如在 1-bit 的 BitNet 模型中實現的 int8*int1 或 int8*int2，以及浮點數與整數的混合乘法 FP16*int4。
為了充分發揮低位元量化的優勢，讓硬體裝置能夠直接支援混合精度矩陣乘法，確保大模型在端側裝置上的高速有效執行，微軟亞洲研究院的研究員們針對現有 CPU、GPU 計算運算元和硬體架構進行創新：
推出了資料型別編譯器 Ladder，支援各種低精度資料型別的表達和相互轉換，將硬體不支援的資料型別無損轉換為硬體支援的資料型別指令，在傳統計算模式下，使得硬體能夠支援混合精度的 DNN（深度神經網路） 計算；
研發了全新演算法 T-MAC，基於查詢表（Lookup Table，LUT）的方法，實現了硬體對混合精度矩陣乘法的直接支援，軟體層面，在 CPU 上的計算相比傳統計算模式取得了更好的加速；
提出了新的硬體架構 LUT Tensor Core，為下一代人工智慧硬體設計開啟了新思路。
Ladder：自定義資料型別無損轉換成硬體支援的資料型別
當前，前沿加速器正在將更低位元的計算單元，如 FP32、FP16，甚至 FP8 的運算整合到新一代的架構中。然而，受限於晶片面積和高昂的硬體成本，每個加速器只能為標準的資料型別提供有限型別的計算單元，比如 NVIDIA V100 TENSOR CORE GPU 僅支援 FP16，而 A100 雖然加入了對 int2、int4、int8 的支援，但並未涵蓋更新的 FP8 或 OCP-MXFP 等資料格式。此外，大模型的快速迭代與硬體升級的緩慢步伐之間存在差距，導致許多新資料型別無法得到硬體支援，進而影響大模型的加速和執行。
微軟亞洲研究院的研究員們發現，儘管硬體加速器缺乏針對自定義資料型別的計算指令，但其記憶體系統可以將它們轉換為固定位寬的不透明資料塊來儲存任意資料型別。同時，大多數自定義資料型別可以無損地轉換為現有硬體計算單元支援的更多位的標準資料型別。例如，NF4 張量可以轉換成 FP16 或 FP32 以執行浮點運算。
基於這些發現，研究員們提出了一種透過分離資料儲存和計算來支援所有自定義資料型別的方法，並研發了資料編譯器 Ladder，以彌合不斷出現的自定義資料型別與當前硬體支援的固有精度格式之間的差距。
Ladder 定義了一套資料型別系統，包括資料型別之間無損轉換的抽象，它能夠表示演算法和硬體支援的各種資料型別，並定義了資料型別之間的轉換規則。當處理低位元演算法應用時，Ladder 透過一系列最佳化，將低位元資料轉譯成當前硬體上最高效的執行格式，包括對計算和儲存的最佳化 —— 將演算法對映到匹配的計算指令，並將不同格式的資料儲存到不同級別的儲存單元中，以實現最高效的運算。
圖 1：Ladder 的系統架構
在 NVIDIA A100、NVIDIA V100、NVIDIA RTX A6000、NVIDIA RTX 4090 和 AMD Instinct MI250 GPU 上執行的 DNN 推理效能評估顯示，Ladder 在原生支援資料型別上超越了現有最先進的 DNN 編譯器，並且在支援 GPU 原本不支援的自定義資料型別方面表現出色，最高提速可達 14.6 倍。
Ladder 是首個在現代硬體加速器上執行 DNN 時，可以系統性地支援以自定義資料型別表示低位元精度資料的系統。這為模型研究者提供了更靈活的資料型別最佳化方法，同時也讓硬體架構開發者在不改變硬體的情況下，支援更廣泛的資料型別。
T-MAC：無需乘法的通用低位元混合精度矩陣乘計算
為了讓現有硬體裝置支援不同的資料模式和混合精度矩陣乘法，在端側部署大模型時，常見的做法是對低位元模型進行反量化。然而，這種方法存在兩大問題：首先，從效能角度來看，反量化過程中的轉換開銷可能會抵消低位元量化帶來的效能提升；其次，從開發角度來看，開發者需要針對不同的混合精度重新設計資料佈局和計算核心。微軟亞洲研究院的研究員們認為，在裝置上部署低位元量化的大模型，關鍵在於如何基於低位元的特點來突破傳統矩陣乘法的實現。
為此，研究員們從系統和演算法層面提出了一種基於查詢表（LUT，Look-Up Table）的方法 T-MAC，幫助低位元量化的大模型在 CPU 上實現高效推理。T-MAC 的核心思想在於利用混合精度矩陣乘法的一端為極低位元（如 1 位元或 2 位元）的特點。它們的輸出結果只有 2 的 1 次方和 2 的 2 次方種可能，這些較少的輸出結果完全可以提前計算並儲存在表中，在運算時，只需從表中讀取結果，避免了重複計算，大幅減少了乘法和加法的運算次數。
具體而言，T-MAC 將傳統的以資料型別為中心的乘法轉變為基於位的查詢表操作，實現了一種統一且可擴充套件的混合精度矩陣乘法解決方案，減小了表的大小並使其停留在最快的記憶體單元中，降低了隨機訪問表的成本。這一創新為在資源受限的邊緣裝置上部署低位元量化大模型鋪平了道路。
圖 2：T-MAC 示意圖
在針對低位元量化的 Llama 和 1 位元的 BitNet 大語言模型的測試中，T-MAC 展現出了顯著的效能優勢。在搭載了最新高通 Snapdragon X Elite 晶片組的 Surface Laptop 7 上，T-MAC 讓 3B BitNet-b1.58 模型的生成速率達到每秒 48 個 token，2bit 7B Llama 模型的生成速率達到每秒 30 個 token，4bit 7B Llama 模型的生成速率可達每秒 20 個 token，這些速率均遠超人類的平均閱讀速度。與原始的 Llama.cpp 框架相比，其提升了 4 至 5 倍，甚至比專用的 NPU 加速器還快兩倍。
即使是在效能較低的裝置上，如 Raspberry Pi（樹莓派）5，T-MAC 也能使 3B BitNet-b1.58 模型達到每秒 11 個 token 的生成速率。T-MAC 還具有顯著的功耗優勢，在資源受限的裝置上可以達到相同的生成速率，而它所需的核心數僅為原始 Llama.cpp 的 1/4 至 1/6。
這些結果表明，T-MAC 提供了一種實用的解決方案，使得在使用通用 CPU 的邊緣裝置上部署大語言模型更為高效，且無需依賴 GPU，讓大模型在資源受限的裝置上也能高效執行，從而推動大模型在更廣泛的場景中的應用。
LUT Tensor Core：推動下一代硬體加速器原生支援混合精度矩陣乘法
T-MAC 和 Ladder 都是在現有 CPU 和 GPU 架構上，實現對混合精度矩陣乘法的最佳化支援。儘管這些軟體層面的創新顯著提升了計算效率，但它們在效率上仍無法與能夠直接實現一個專門查詢表的硬體加速器相比。研究員們認為，最理想的方法是重新設計硬體加速器，讓 CPU、GPU 等能夠原生支援混合精度矩陣乘法，但這一目標面臨三大挑戰：
效率：設計和實現方式必須具有成本效益，透過最佳化晶片的利用面積，最大限度地提高低位元資料的計算效益。
靈活性：由於不同的模型和場景需要不同的權重和啟用精度，因此硬體中的混合精度矩陣乘法設計必須能夠處理各種權重精度 (如 int4/2/1) 和啟用精度 (如 FP16/8、int8) 及其組合。
相容性：新設計必須與現有的 GPU 架構和軟體生態系統無縫整合，以加速新技術的應用。
為了應對這些挑戰，微軟亞洲研究院的研究員們設計了 LUT Tensor Core，這是一種利用查詢表直接執行混合精度矩陣乘法的 GPU Tensor Core 微架構。一方面，基於查詢表的設計將乘法運算簡化為表預計算操作，可直接在表中查詢結果，提高計算效率。另一方面，這種方法也簡化了對硬體的需求，它只需用於表儲存的暫存器和用於查詢的多路選擇器，無需乘法器和加法器。同時，LUT Tensor Core 透過位元序列設計實現了權重精度的靈活性，並利用表量化實現了啟用精度的靈活性。
此外，為了與現有 GPU 微架構和軟體堆疊整合，研究員們擴充套件了 GPU 中現有的 MMA 指令集，加入了一組 LMMA 指令，並設計了一個類似於 cuBLAS 的軟體堆疊，用於整合到現有的 DNN 框架中。研究員們還設計了一個編譯器，用於在具有 LUT Tensor Core 的 GPU 上進行端到端的執行計劃。這些創新方法可以讓 LUT Tensor Core 被無縫、快速地採用。
圖 3：LUT Tensor Core 微架構概述
在 Llama 和 BitNet 模型上的測試顯示，LUT Tensor Core 可以提供高達 6.93 倍的推理速度，且只佔傳統 Tensor Core 面積的 38.7%。在幾乎相同的模型精度下，這相當於 20.7 倍的計算密度和 19.1 倍的能效提升。隨著人工智慧大模型規模和複雜性的不斷增長，LUT Tensor Core 有助於進一步釋放低位元大語言模型的潛力，推動人工智慧在新場景中的應用。
“查詢表方法引領了計算正規化的轉變。在過去，我們依賴於矩陣乘法和累加運算，而在大模型時代，得益於低位元量化技術，查詢表方法將成為主流。相較於傳統的浮點運算或矩陣乘法，查詢表方法在計算上更輕便高效，而且在硬體層面上更易於擴充套件，能夠實現更高的電晶體密度，在單位晶片面積上提供更大的吞吐量，從而推動硬體架構的革新。” 微軟亞洲研究院首席研究員曹婷表示。
低位元量化的長尾效應：為具身智慧帶來新可能
低位元量化技術不僅最佳化了大模型在端側裝置上的執行效率，還透過減少單個引數的 “體積”，為模型引數的擴充套件（Scale up）提供了新的空間。這種引數擴充套件能力，使模型擁有了更強的靈活性和表達能力，正如 BitNet 模型所展示的，從低位元模型出發，逐步擴充套件至更大規模的訓練。
微軟亞洲研究院的 T-MAC、Ladder 和 LUT Tensor Core 等創新技術，為各種低位元量化大模型提供了高效能的執行方案，使得這些模型能夠在各種裝置上高效執行，並推動科研人員從低位元角度設計和最佳化大模型。其中部分技術已經在微軟必應（Bing）搜尋及其廣告業務等搜尋大模型中發揮作用。隨著對記憶體和計算資源的降低，低位元大模型在機器人等具身智慧系統上的部署也將成為可能，可以使這些裝置更好地實現與環境的動態感知和實時互動。
目前，T-MAC 和 Ladder 已經在 GitHub 上開源，歡迎相關研發人員測試應用，與微軟亞洲研究院共同探索人工智慧技術的更多可能。
相關連結：
Ladder 論文連結：https://www.usenix.org/conference/osdi24/presentation/wang-lei
BitBLAS/Ladder GitHub 連結：https://github.com/microsoft/BitBLAS
T-MAC 論文連結：https://arxiv.org/abs/2407.00088
T-MAC GitHub 連結：https://github.com/microsoft/T-MAC
LUT Tensor Core 論文連結：https://arxiv.org/abs/2408.06003
BitDistiller 論文連結: https://arxiv.org/abs/2402.10631
BitDistiller GitHub 連結：https://github.com/DD-DuDa/BitDistiller