吞吐量比FP16提升2.65倍
克雷西 發自 凹非寺
量子位 | 公眾號 QbitAI
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
本文鏈接：單卡跑Llama 70B快過雙卡，微軟硬生生把FP6搞到了A100里 | 開源http://m.sq15.cn/show-2-5562-0.html
聲明：本網站為非營利性網站，本網頁內容由互聯網博主自發貢獻，不代表本站觀點，本站不承擔任何法律責任。天上不會到餡餅，請大家謹防詐騙！若有侵權等問題請及時與本網聯系，我們將在第一時間刪除處理。
上一篇： 字節跳動反腐持續加壓 大廠高舉反舞弊大旗
下一篇： “人工智能＋”浪潮下，財務領域如何突圍？