DeepSeek開源放大招:FlashMLA讓H800算力狂飆!曝光低成本秘笈
就在剛剛,DeepSeek放出了開源周首日的重磅炸彈——FlashMLA。
這是DeepSeek專為英偉達(dá)Hopper GPU打造的高效MLA解碼內(nèi)核,特別針對(duì)變長序列進(jìn)行了優(yōu)化,目前已正式投產(chǎn)使用。
經(jīng)實(shí)測(cè),F(xiàn)lashMLA在H800 SXM5平臺(tái)上(CUDA 12.6),在內(nèi)存受限配置下可達(dá)最高3000GB/s,在計(jì)算受限配置下可達(dá)峰值580 TFLOPS。。
圖片
圖片
開源地址:https://github.com/deepseek-ai/FlashMLA
當(dāng)前已經(jīng)發(fā)布的內(nèi)容為:
- 對(duì)BF16精度的支持
- 塊大小為64的分頁KV緩存
團(tuán)隊(duì)在致謝部分表示,F(xiàn)lashMLA的設(shè)計(jì)參考了FlashAttention-2、FlashAttention-3以及CUTLASS的技術(shù)實(shí)現(xiàn)。
有網(wǎng)友對(duì)此表示,「DeepSeek王炸開局,F(xiàn)lashMLA是真正能加速AGI進(jìn)程的」。
圖片
快速入門
首先,需要打開終端,輸入下面代碼安裝setup.py文件:
這是一個(gè)基于Python的安裝命令,用于編譯和安裝FlashMLA模塊,確保其高效運(yùn)行于特定硬件。
python setup.py install
基準(zhǔn)測(cè)試:
這段代碼是一個(gè)測(cè)試腳本,用于驗(yàn)證FlashMLA的功能和性能,并與PyTorch的基準(zhǔn)實(shí)現(xiàn)進(jìn)行對(duì)比。
python tests/test_flash_mla.py
使用方法:
下面是一段使用的示例代碼。
from flash_mla import get_mla_metadata, flash_mla_with_kvcache
tile_scheduler_metadata, num_splits = get_mla_metadata(cache_seqlens, s_q * h_q // h_kv, h_kv)
for i in range(num_layers):
...
o_i, lse_i = flash_mla_with_kvcache(
q_i, kvcache_i, block_table, cache_seqlens, dv,
tile_scheduler_metadata, num_splits, causal=True,
)
...
DeepSeek訓(xùn)練成本如此之低的兩大關(guān)鍵
DeepSeek的成本涉及兩項(xiàng)關(guān)鍵的技術(shù):一個(gè)是MoE,一個(gè)就是MLA(多頭潛注意力)。
其中,MLA的開發(fā)耗時(shí)數(shù)月,可將每個(gè)查詢KV緩存量減少93.3%,顯著減少了推理過程中的內(nèi)存占用(在訓(xùn)練過程也是如此)。
圖片
MLA架構(gòu)需要一些巧妙的設(shè)計(jì),因此實(shí)現(xiàn)的復(fù)雜性大大增加。而DeepSeek成功地將這些技術(shù)整合在一起,表明他們?cè)诟咝дZ言模型訓(xùn)練方面走在了前沿
多頭潛注意力(MLA)
KV緩存是Transforme模型中的一種內(nèi)存機(jī)制,用于存儲(chǔ)表示對(duì)話上下文的數(shù)據(jù),從而減少不必要的計(jì)算開銷。
隨著對(duì)話上下文的增長,KV緩存會(huì)不斷擴(kuò)大,從而造成顯著的內(nèi)存限制。
通過大幅減少每次查詢所需的KV緩存量,可以相應(yīng)減少每次查詢所需的硬件資源,從而降低運(yùn)營成本。
與標(biāo)準(zhǔn)注意力機(jī)制相比,MLA將每次查詢所需的KV緩存減少了約93.3%。
圖片
MLA這種全新多頭潛注意力,可以將注意力機(jī)制的內(nèi)存占用減少大約80%到90%,尤其有助于處理長上下文
此外,由于H20芯片比H100具有更高的內(nèi)存帶寬和容量,DeepSeek在推理工作負(fù)載方面獲得了更多效率提升。
除了MLA,DeepSeek其他突破性進(jìn)展還有哪些?
訓(xùn)練(前期和后期)
不是「下一個(gè)token預(yù)測(cè)」,而是「多token預(yù)測(cè)」
DeepSeek V3以前所未見的規(guī)模實(shí)現(xiàn)了多Token預(yù)測(cè)(MTP)技術(shù),這些新增的注意力模塊可以預(yù)測(cè)接下來的多個(gè)Token,而不是傳統(tǒng)的單個(gè)Token。
這顯著提高了訓(xùn)練階段的模型性能,且這些模塊可以在推理階段移除。
這是一個(gè)典型的算法創(chuàng)新案例,實(shí)現(xiàn)了在更低計(jì)算資源消耗下的性能提升。
其他方面,雖然DeepSeek在訓(xùn)練中采用了FP8精度,但像全球一些頂尖的實(shí)驗(yàn)室已經(jīng)采用這項(xiàng)技術(shù)相當(dāng)長時(shí)間了。
DeepSeek V3采用了我們常見的「混合專家模型」(MoE)架構(gòu),個(gè)由多個(gè)專門處理不同任務(wù)的小型專家模型組成的大模型,展現(xiàn)出強(qiáng)大的涌現(xiàn)能力。
MoE模型面臨的主要挑戰(zhàn)是,如何確定將哪個(gè)Token分配給哪個(gè)子模型(即「專家」)。
DeepSeek創(chuàng)新性地采用了一個(gè)「門控網(wǎng)絡(luò)」(gating network),能夠高效且平衡地將Token路由到相應(yīng)的專家,同時(shí)保持模型性能不受影響。
這意味著路由過程非常高效,在訓(xùn)練過程中每個(gè)Token只需要調(diào)整小量參數(shù)(相較于模型整體規(guī)模)。
這既提高了訓(xùn)練效率,又降低了推理成本。
盡管有人擔(dān)心MoE帶來的效率提升,可能降低投資意愿,但Dario指出,更強(qiáng)大的AI模型帶來的經(jīng)濟(jì)效益非??捎^,任何節(jié)省的成本都會(huì)立即被投入到開發(fā)更大規(guī)模的模型中。
因此,MoE效率提升不會(huì)減少總體投資,反而會(huì)加速模型Scaling的進(jìn)程。
當(dāng)前,包括OpenAI、谷歌、Anthropic等一些公司正專注于擴(kuò)大模型的計(jì)算規(guī)模,并提高算法效率。
V3打好了基礎(chǔ),RL立大功
對(duì)于R1而言,它極大地受益于其強(qiáng)大的基礎(chǔ)模型——V3,這在很大程度上要?dú)w功于強(qiáng)化學(xué)習(xí)(RL)。
RL主要關(guān)注兩個(gè)方面:格式化(確保輸出連貫性)以及有用性與安全性(確保模型實(shí)用且無害)。
模型的推理能力,是在對(duì)合成數(shù)據(jù)集進(jìn)行微調(diào)過程中自然涌現(xiàn)的,這與o1的情況類似。
值得注意的是,R1論文中并沒有提及具體的計(jì)算量,因?yàn)榕妒褂玫挠?jì)算資源,會(huì)暴露DeepSeek實(shí)際擁有的GPU數(shù)量遠(yuǎn)超過其對(duì)外宣稱的規(guī)模。
這種規(guī)模的強(qiáng)化學(xué)習(xí)需要龐大的計(jì)算資源,特別是在生成合成數(shù)據(jù)時(shí)。
談到蒸餾,R1論文最引人注目的發(fā)現(xiàn)可能是,通過具有推理能力的模型輸出來微調(diào)較小的非推理模型,使其獲得推理能力。
數(shù)據(jù)集包含了約80萬個(gè)樣本,現(xiàn)在研究人員可以利用R1的思維鏈(CoT)輸出創(chuàng)建自己的數(shù)據(jù)集,并借此開發(fā)具有推理能力的模型。
未來,我們可能會(huì)看到更多小模型展現(xiàn)出推理能力,從而提升小模型的整體性能。
參考資料: