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