FP8和更低的浮點(diǎn)數(shù)量化精度，不再是H100的“專利”了！

老黃想讓大家用INT8/INT4，微軟DeepSpeed團(tuán)隊(duì)在沒有英偉達(dá)官方支持的條件下，硬生生在A100上跑起FP6。

測(cè)試結(jié)果表明，新方法TC-FPx在A100上的FP6量化，速度接近甚至偶爾超過INT4，而且擁有比后者更高的精度。

在此基礎(chǔ)之上，還有端到端的大模型支持，目前已經(jīng)開源并集成到了DeepSpeed等深度學(xué)習(xí)推理框架中。

這一成果對(duì)大模型的加速效果也是立竿見影——在這種框架下用單卡跑Llama，吞吐量比雙卡還要高2.65倍。

一名機(jī)器學(xué)習(xí)研究人員看了后表示，微軟的這項(xiàng)研究簡(jiǎn)直可以用crazy來形容。

表情包也第一時(shí)間上線，be like：

英偉達(dá)：只有H100支持FP8。

微軟：Fine，我自己搞定。

那么，這個(gè)框架到底能實(shí)現(xiàn)什么樣的效果，背后又采用了什么樣的技術(shù)呢？

用FP6跑Llama，單卡比雙卡還快

在A100上使用FP6精度，帶來的是內(nèi)核級(jí)的性能提升。

研究人員選取了不同大小的Llama模型和OPT模型之中的線性層，在NVIDIA A100-40GB GPU平臺(tái)上，使用CUDA 11.8進(jìn)行了測(cè)試。

結(jié)果相比于英偉達(dá)官方的cuBLAS（W16A16）和TensorRT-LLM（W8A16），TC-FPx（W6A16）速度提升的最大值分別是2.6倍和1.9倍。

相比于4bit的BitsandBytes（W4A16）方法，TC-FPx的最大速度提升則是達(dá)到了8.9倍。

（W和A分別代表權(quán)重量化位寬和激活量化位寬）

△歸一化數(shù)據(jù)，以cuBLAS結(jié)果為1

同時(shí)，TC-FPx內(nèi)核還減少了對(duì)DRAM內(nèi)存的訪問，并提高了DRAM帶寬利用率和Tensor Cores利用率，以及ALU和FMA單元的利用率。

在TC-FPx基礎(chǔ)之上設(shè)計(jì)的端到端推理框架FP6-LLM，也給大模型帶來了顯著的性能提高。

以Llama-70B為例，用FP6-LLM在單卡上的運(yùn)行吞吐量，比FP16在雙卡上還要高出2.65倍，在16以下的批大小中的延遲也低于FP16。

而對(duì)于參數(shù)量小一些的模型OPT-30B（FP16也使用單卡），F(xiàn)P6-LLM同樣帶來了明顯的吞吐量提升和延遲降低。

而且單卡FP16在這種條件下最多支持的批大小只有4，F(xiàn)P6-LLM卻可以在批大小為16的情況下正常運(yùn)行。

那么，微軟團(tuán)隊(duì)是怎樣實(shí)現(xiàn)在A100上運(yùn)行FP16量化的呢？

重新設(shè)計(jì)內(nèi)核方案

為了實(shí)現(xiàn)對(duì)包括6bit在內(nèi)精度的支持，TC-FPx團(tuán)隊(duì)設(shè)計(jì)了一個(gè)統(tǒng)一的內(nèi)核方案，可以支持不同位寬的量化權(quán)重。

相比于傳統(tǒng)的雙內(nèi)核方法，TC-FPx通過將去量化和矩陣乘法融合在單個(gè)內(nèi)核中，減少了內(nèi)存訪問次數(shù)，提高了性能。

實(shí)現(xiàn)低精度量化的核心奧義則是通過去量化方式，將FP6精度的數(shù)據(jù)“偽裝”成FP16，然后按照FP16的格式交給GPU進(jìn)行運(yùn)算。

同時(shí)團(tuán)隊(duì)還利用了位級(jí)預(yù)打包技術(shù)，解決GPU內(nèi)存系統(tǒng)對(duì)非2的冪次位寬（如6-bit）不友好的問題。

具體來說，位級(jí)預(yù)打包是在模型推理之前對(duì)權(quán)重?cái)?shù)據(jù)進(jìn)行重新組織，包括將6-bit量化的權(quán)重重新排列，以便它們能夠以GPU內(nèi)存系統(tǒng)友好的方式進(jìn)行訪問。

此外，由于GPU內(nèi)存系統(tǒng)通常以32位或64位的塊進(jìn)行數(shù)據(jù)訪問，位級(jí)預(yù)打包技術(shù)將還會(huì)6-bit權(quán)重打包，使得它們能夠以這些對(duì)齊的塊的形式存儲(chǔ)和訪問。

預(yù)打包完成后，研究團(tuán)隊(duì)使用SIMT核心的并行處理能力，對(duì)寄存器中的FP6權(quán)重執(zhí)行并行去量化，生成FP16格式的權(quán)重。

去量化后的FP16權(quán)重在寄存器中被重構(gòu)，然后送入Tensor Core，使用重構(gòu)后的FP16權(quán)重執(zhí)行矩陣乘法運(yùn)算，完成線性層的計(jì)算。

在此過程中，團(tuán)隊(duì)利用了SMIT核心的位級(jí)并行性，提高了整個(gè)去量化過程的效率。

而為了權(quán)重重構(gòu)任務(wù)能夠并行運(yùn)行，團(tuán)隊(duì)還使用了一種并行權(quán)重拼接技術(shù)。

具體來說，每個(gè)權(quán)重被分割成幾個(gè)部分，每個(gè)部分的位寬是2的冪次（如把6分割成2 4或4 2）。

在去量化之前，權(quán)重首先從共享內(nèi)存加載到寄存器中。由于每個(gè)權(quán)重被分割成多個(gè)部分，需要在運(yùn)行時(shí)在寄存器級(jí)別重構(gòu)完整的權(quán)重。

為了減少運(yùn)行時(shí)的開銷，TC-FPx提出了一種并行提取和拼接權(quán)重的方法。這種方法使用兩組寄存器來存儲(chǔ)32個(gè)FP6權(quán)重的片段，并行地重構(gòu)這些權(quán)重。

同時(shí)，為了并行提取和拼接權(quán)重，需要確保初始數(shù)據(jù)布局滿足特定的順序要求，因此TC-FPx通過在運(yùn)行前對(duì)權(quán)重片段進(jìn)行重排。

此外，TC-FPx還設(shè)計(jì)了一個(gè)軟件流水線，將去量化步驟與Tensor Core的矩陣乘法操作融合在一起，通過指令級(jí)并行性提高了整體的執(zhí)行效率。

論文地址：https://arxiv.org/abs/2401.14112