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

在沒有英偉達(dá)官方支持的情況下，老黃希望大家使用INT8/INT4，微軟DeepSpeed團(tuán)隊(duì)，生活在A100上跑FP6。

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

基于這個(gè)基礎(chǔ)，還有大模型支持端到端，現(xiàn)已開源并集成到DeepSpeed等深度學(xué)習(xí)推理框架中。

這個(gè)結(jié)果對(duì)于大型模型的加速效果是立竿見影的——在這個(gè)框架下用單卡跑Llama，吞吐量比雙卡高2.65倍。

一位機(jī)器學(xué)習(xí)研究人員看了之后說，微軟的這項(xiàng)研究可以用crazy來形容。

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

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

微軟：Fine，我自己做的。

那這個(gè)框架到底能達(dá)到什么樣的效果，背后又用了什么樣的技術(shù)呢？

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

FP6精度用于A100，帶來的是提高核心級(jí)別的性能。

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

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

與4bit的BitsandBytes(W4A16)方法相比，TC-FPx的最大速度提升是8.9倍。

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

△

同時(shí)，TC-FPx核心還減少了對(duì)DRAM內(nèi)存的瀏覽，提高了DRAM帶寬利用率和Tensor ALU和FMA模塊的利用率，Cores利用率。

基于TC-FPx設(shè)計(jì)的設(shè)計(jì)FP6端到端推理框架-LLM，同時(shí)也給大型模型帶來了顯著的性能提升。

以Llama-70B為例，單卡上使用FP6-LLM的吞吐量比FP16在雙卡上的吞吐量高2.65倍，16以下批次大小的延遲也低于FP16。

FP6-LLM也帶來了明顯的吞吐量增加和延遲減少，對(duì)于參數(shù)較小的OPT-30B(FP16也使用單卡)。

而且在這種情況下，單卡FP16最支持的批次大小只有4，F(xiàn)P6-LLM可以在批次大小為16的情況下正常工作。

那微軟團(tuán)隊(duì)是如何在A100上實(shí)現(xiàn)FP16量化的呢？

02 重新設(shè)計(jì)核心方案

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

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

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

與此同時(shí)，團(tuán)隊(duì)也在使用位級(jí)預(yù)打包技術(shù)，處理GPU內(nèi)存系統(tǒng)對(duì)非2的冪次位寬(例如6-bit）不友好的問題。

具體來說，排名預(yù)包裝是在模型推理之前重新組織權(quán)重信息，包括重新排序6-bit量化權(quán)重，這樣它們就可以以友好的方式瀏覽GPU內(nèi)存系統(tǒng)。

另外，由于GPU內(nèi)存系統(tǒng)通常采用32或64位塊進(jìn)行數(shù)據(jù)訪問，所以等級(jí)預(yù)包裝技術(shù)也將采用6-bit權(quán)重進(jìn)行包裝，使其能夠以這些對(duì)齊塊的形式進(jìn)行存儲(chǔ)和瀏覽。

預(yù)包裝完成后，研究小組利用SIMT核心的并行計(jì)算能力，并行量化存儲(chǔ)器中的FP6權(quán)重，生成FP16格式的權(quán)重。

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

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

同時(shí)，團(tuán)隊(duì)還使用了一種權(quán)重重構(gòu)任務(wù)可以并行運(yùn)行的方法。并行權(quán)重拼接技術(shù)。

具體而言，每個(gè)權(quán)重被分成幾個(gè)部分，每個(gè)部分的位寬是2的冪(例如，將6分成2。 4或4 2）。

在量化之前，重量首先從共享內(nèi)存加載到存儲(chǔ)器。由于每個(gè)重量分為多個(gè)部分，在運(yùn)行過程中需要在存儲(chǔ)器等級(jí)中重構(gòu)完整的重量。

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

與此同時(shí)，TC-FPx為了平行提取和拼接權(quán)重，需要保證初始數(shù)據(jù)布局符合特定的順序要求，所以TC-FPx在運(yùn)行前重新排列權(quán)重片段。

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

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

參考鏈接：https://twitter.com/rohanpaul_ai/status/1784599257384727044

本文來自微信微信官方賬號(hào)“量子位”（ID:QbitAI），作者：關(guān)注前沿技術(shù)，授權(quán)發(fā)布36氪。

本文僅代表作者觀點(diǎn)，版權(quán)歸原創(chuàng)者所有，如需轉(zhuǎn)載請(qǐng)?jiān)谖闹凶⒚鱽碓醇白髡呙帧?/p>

免責(zé)聲明：本文系轉(zhuǎn)載編輯文章，僅作分享之用。如分享內(nèi)容、圖片侵犯到您的版權(quán)或非授權(quán)發(fā)布，請(qǐng)及時(shí)與我們聯(lián)系進(jìn)行審核處理或刪除，您可以發(fā)送材料至郵箱：service@tojoy.com