FP8 和更低的浮點數量化精度，不再是 H100 的 " 專利 " 了！
老黃想讓大家用 INT8/INT4，微軟 DeepSpeed 團隊在沒有英偉達官方支持的條件下，硬生生在 A100 上跑起 FP6。
測試結果表明，新方法 TC-FPx 在 A100 上的 FP6 量化，速度接近甚至偶爾超過 INT4，而且擁有比後者更高的精度。
在此基礎之上，還有端到端的大模型支持，目前已經開源并集成到了 DeepSpeed 等深度學習推理框架中。
這一成果對大模型的加速效果也是立竿見影——在這種框架下用單卡跑 Llama，吞吐量比雙卡還要高 2.65 倍。
一名機器學習研究人員看了後表示，微軟的這項研究簡直可以用 crazy 來形容。
表情包也第一時間上線，be like：
英偉達：隻有 H100 支持 FP8。
微軟：Fine，我自己搞定。
那麽，這個框架到底能實現什麽樣的效果，背後又采用了什麽樣的技術呢？
用 FP6 跑 Llama，單卡比雙卡還快
在 A100 上使用 FP6 精度，帶來的是内核級的性能提升。
研究人員選取了不同大小的 Llama 模型和 OPT 模型之中的線性層，在 NVIDIA A100-40GB GPU 平台上，使用 CUDA 11.8 進行了測試。
結果相比于英偉達官方的 cuBLAS（W16A16）和 TensorRT-LLM（W8A16），TC-FPx（W6A16）速度提升的最大值分别是 2.6 倍和 1.9 倍。
相比于 4bit 的 BitsandBytes（W4A16）方法，TC-FPx 的最大速度提升則是達到了 8.9 倍。
（W 和 A 分别代表權重量化位寬和激活量化位寬）
△歸一化數據，以 cuBLAS 結果爲 1
同時，TC-FPx 内核還減少了對 DRAM 内存的訪問，并提高了 DRAM 帶寬利用率和 Tensor Cores 利用率，以及 ALU 和 FMA 單元的利用率。
在 TC-FPx 基礎之上設計的端到端推理框架 FP6-LLM，也給大模型帶來了顯著的性能提高。
以 Llama-70B 爲例，用 FP6-LLM 在單卡上的運行吞吐量，比 FP16 在雙卡上還要高出 2.65 倍，在 16 以下的批大小中的延遲也低于 FP16。
而對于參數量小一些的模型 OPT-30B（FP16 也使用單卡），FP6-LLM 同樣帶來了明顯的吞吐量提升和延遲降低。
而且單卡 FP16 在這種條件下最多支持的批大小隻有 4，FP6-LLM 卻可以在批大小爲 16 的情況下正常運行。
那麽，微軟團隊是怎樣實現在 A100 上運行 FP16 量化的呢？
重新設計内核方案
爲了實現對包括 6bit 在内精度的支持，TC-FPx 團隊設計了一個統一的内核方案，可以支持不同位寬的量化權重。
相比于傳統的雙内核方法，TC-FPx 通過将去量化和矩陣乘法融合在單個内核中，減少了内存訪問次數，提高了性能。
實現低精度量化的核心奧義則是通過去量化方式，将 FP6 精度的數據 " 僞裝 " 成 FP16，然後按照 FP16 的格式交給 GPU 進行運算。
同時團隊還利用了位級預打包技術，解決 GPU 内存系統對非 2 的幂次位寬（如 6-bit）不友好的問題。
具體來說，位級預打包是在模型推理之前對權重數據進行重新組織，包括将 6-bit 量化的權重重新排列，以便它們能夠以 GPU 内存系統友好的方式進行訪問。
此外，由于 GPU 内存系統通常以 32 位或 64 位的塊進行數據訪問，位級預打包技術将還會 6-bit 權重打包，使得它們能夠以這些對齊的塊的形式存儲和訪問。
預打包完成後，研究團隊使用 SIMT 核心的并行處理能力，對寄存器中的 FP6 權重執行并行去量化，生成 FP16 格式的權重。
去量化後的 FP16 權重在寄存器中被重構，然後送入 Tensor Core，使用重構後的 FP16 權重執行矩陣乘法運算，完成線性層的計算。
在此過程中，團隊利用了 SMIT 核心的位級并行性，提高了整個去量化過程的效率。
而爲了權重重構任務能夠并行運行，團隊還使用了一種并行權重拼接技術。
具體來說，每個權重被分割成幾個部分，每個部分的位寬是 2 的幂次（如把 6 分割成 2+4 或 4+2）。
在去量化之前，權重首先從共享内存加載到寄存器中。由于每個權重被分割成多個部分，需要在運行時在寄存器級别重構完整的權重。
爲了減少運行時的開銷，TC-FPx 提出了一種并行提取和拼接權重的方法。這種方法使用兩組寄存器來存儲 32 個 FP6 權重的片段，并行地重構這些權重。
同時，爲了并行提取和拼接權重，需要确保初始數據布局滿足特定的順序要求，因此 TC-FPx 通過在運行前對權重片段進行重排。
此外，TC-FPx 還設計了一個軟件流水線，将去量化步驟與 Tensor Core 的矩陣乘法操作融合在一起，通過指令級并行性提高了整體的執行效率。
論文地址：
https://arxiv.org/abs/2401.14112
參考鏈接：
https://twitter.com/rohanpaul_ai/status/1784599257384727044
— 完 —
點這裏關注我，記得标星哦～
一鍵三連「分享」、「點贊」和「在看」
科技前沿進展日日相見 ~
