允中 發(fā)自 凹非寺量子位 | 公眾號(hào) QbitAI
FP8通過其獨(dú)特的數(shù)值表示方式，能夠在保持一定精度的同時(shí)，在大模型訓(xùn)練中提高訓(xùn)練速度、節(jié)省內(nèi)存占用，最終降低訓(xùn)練成本。
AI大模型開發(fā)系統(tǒng)Colossal-AI的混合精度訓(xùn)練再度升級(jí)，支持主流的BF16(O2) + FP8(O1)的新一代混合精度訓(xùn)練方案。
僅需一行代碼，即可對(duì)主流LLM模型能夠獲得平均30%的加速效果，降低相應(yīng)大模型開發(fā)成本，并保證訓(xùn)練收斂性。
無需引入額外的手寫CUDA算子，避免了較長(zhǎng)的AOT編譯時(shí)間和復(fù)雜的編譯環(huán)境配置。
開源地址：https://github.com/hpcaitech/ColossalAI
FP8混合精度訓(xùn)練低精度計(jì)算一直是GPU硬件發(fā)展趨勢(shì)。
從最早的FP32，到目前通用的FP16/BF16，再到Hopper系列芯片(H100, H200, H800等）支持的FP8，低精度計(jì)算速度越來越快，所需的內(nèi)存也越來越低，非常符合大模型時(shí)代對(duì)硬件的需求。
目前FP8混合精度訓(xùn)練影響訓(xùn)練結(jié)果的最大因素就是scaling方案，常見的方案有兩種：
延遲scaling
實(shí)時(shí)scaling
延遲scaling采用之前一段時(shí)間窗口內(nèi)的scaling值來估計(jì)當(dāng)前scaling，同時(shí)將scaling的更新和矩陣乘法(gemm)融合起來。這種計(jì)算方法效率較高，但由于是估算的scaling，所以對(duì)收斂性影響較大。
實(shí)時(shí)scaling直接采用當(dāng)前的張量值來計(jì)算scaling，所以計(jì)算效率較低，但是對(duì)收斂性影響較校根據(jù)英偉達(dá)的報(bào)告，這兩種scaling方案的計(jì)算效率差距在10%以內(nèi)。
Colossal-AI采用了對(duì)訓(xùn)練收斂性影響較小的實(shí)時(shí)scaling方案，同時(shí)實(shí)現(xiàn)有著不輸其他延遲scaling實(shí)現(xiàn)的性能。
在單卡H100上對(duì)矩陣乘法進(jìn)行的測(cè)試，可以看到矩陣的維度越大，F(xiàn)P8的加速效果越明顯，而且Colossal-AI的實(shí)現(xiàn)與Transformer Engine的性能幾乎一致，如圖1所示。但Transformer Engine需要復(fù)雜的AOT編譯環(huán)境配置和較長(zhǎng)的編譯時(shí)間。
△圖1. 單卡GEMM性能測(cè)試
為了實(shí)驗(yàn)結(jié)果更貼近現(xiàn)實(shí)，Colossal-AI直接在主流LLM上進(jìn)行了實(shí)際訓(xùn)練的測(cè)試。
首先在H100單卡上進(jìn)行了測(cè)試，以下測(cè)試中Transformer Engine (TE)采用的其默認(rèn)的延遲scaling方案。
同時(shí)進(jìn)行了收斂性測(cè)試，可以看到FP8混合精度訓(xùn)練的loss曲線與bf16的基本一致，如圖4所示：
△圖4. H100單卡 LLaMA2-7B 混合精度訓(xùn)練loss曲線
Colossal-AI還測(cè)試了H800多卡并行訓(xùn)練場(chǎng)景下的性能。在單機(jī)8卡H800上訓(xùn)練LLaMA2-7B，Colossal-AI FP8對(duì)比Colossal-AI BF16有35%的吞吐提升，對(duì)比Torch FSDP BF16有94%的吞吐提升。
在單機(jī)8卡H800上訓(xùn)練LLaMA2-13B，Colossal-AI FP8對(duì)比Colossal-AI BF16有39%的吞吐提升。
在2機(jī)16卡H800上訓(xùn)練Cohere Command-R 35B，Colossal-AI FP8對(duì)比Colossal-AI BF16有10%的吞吐提升，如圖7所示：
根據(jù)英偉達(dá)的報(bào)告和測(cè)試經(jīng)驗(yàn)，對(duì)FP8混合精度訓(xùn)練性能調(diào)優(yōu)有一些初步的認(rèn)識(shí)：
盡量少使用張量并行，用流水線并行代替張量并行
模型hidden size越大，加速效果越明顯
矩陣乘法占比高的模型加速效果大
由于上述實(shí)驗(yàn)中Command-R 35B采用了張量并行，所以加速效果不太明顯。
Colossal-AI對(duì)FP8的支持較為廣泛，各種并行方式都能和FP8混合精度訓(xùn)練兼容。使用時(shí)，僅需在初始化plugin時(shí)開啟FP8即可：
from colossalai.booster.plugin import GeminiPlugin, HybridParallelPlugin, LowLevelZeroPlugin ... plugin = LowLevelZeroPlugin(..., use_fp8=True) plugin = GeminiPlugin(..., use_fp8=True) plugin = HybridParallelPlugin(..., use_fp8=True)
除此之外，無需多余的代碼和AOT編譯。
開源地址：https://github.com/hpcaitech/ColossalAI
