來源:量子位
DeepSeek-R1火遍海內(nèi)外,但推理服務(wù)器頻頻宕機(jī),專享版按GPU小時計費的天價成本更讓中小團(tuán)隊望而卻步。
而市面上所謂“本地部署”方案,多為參數(shù)量縮水90%的蒸餾版,背后原因是671B參數(shù)的MoE架構(gòu)對顯存要求極高——即便用8卡A100也難以負(fù)荷。因此,想在本地小規(guī)模硬件上跑真正的DeepSeek-R1,被認(rèn)為基本不可能。
但就在近期,清華大學(xué)KVCache.AI團(tuán)隊聯(lián)合趨境科技發(fā)布的KTransformers開源項目公布更新:
支持24G顯存在本地運行DeepSeek-R1、V3的671B滿血版。預(yù)處理速度最高達(dá)到286 tokens/s,推理生成速度最高能達(dá)到14 tokens/s。
其實早在DeepSeek-V2 時代,這個項目就因“專家卸載”技術(shù)而備受關(guān)注——它支持了236B的大模型在僅有24GB顯存的消費級顯卡上流暢運行,把顯存需求砍到10分之一。
HuggingFace 的開源負(fù)責(zé)人的點贊
隨著DeepSeek-R1的發(fā)布,社區(qū)的需求迅速激增,在GitHub蓋起上百樓的issue,呼吁對其進(jìn)行支持。
版本更新發(fā)布后,不少開發(fā)者也紛紛用自己的3090顯卡和200GB內(nèi)存進(jìn)行實測,借助與Unsloth優(yōu)化的組合,Q2_K_XL模型的推理速度已達(dá)到9.1 tokens/s,真正實現(xiàn)了千億級模型的“家庭化”。
此外,KTransformers團(tuán)隊還公布了v0.3預(yù)覽版的性能指標(biāo),將通過整合Intel AMX指令集,CPU預(yù)填充速度最高至286 tokens/s,相比llama.cpp快了近28倍。對于那些需要處理上萬級Token上下文的長序列任務(wù)(比如大規(guī)模代碼庫分析)來說,相當(dāng)于能夠從“分鐘級等待”瞬間邁入“秒級響應(yīng)”,徹底釋放CPU的算力潛能。
另外,KTransformers還提供了兼容Hugginface Transformers的API與ChatGPT式Web界面,極大降低了上手難度。同時,其基于YAML的“模板注入框架”能夠靈活切換量化策略、內(nèi)核替換等多種優(yōu)化方式。
目前,KTransformers在localLLaMa社區(qū)持續(xù)位居熱榜第一,有上百條開發(fā)者的討論。
項目背后的技術(shù)細(xì)節(jié),團(tuán)隊也給出了詳細(xì)介紹。
利用MoE架構(gòu)的稀疏性
DeepSeek-R1/V3均采用了MoE(混合專家)架構(gòu),這種架構(gòu)的核心是將模型中的任務(wù)分配給不同的專家模塊,每個專家模塊專注于處理特定類型的任務(wù)。MoE結(jié)構(gòu)的模型具有很強的稀疏性,在執(zhí)行推理任務(wù)的時候,每次只會激活其中一部分的模型參數(shù)。
因此,MoE架構(gòu)需要大量的存儲空間,但是并不需要很多的計算資源。
基于此,團(tuán)隊采用了GPU/CPU的異構(gòu)計算劃分策略:僅將非Shared部分的稀疏MoE矩陣放在CPU/DRAM上并使用llamafile提供的高速算子處理,剩余稠密部分放在GPU上使用Marlin算子處理。
在這樣的情況下,同樣使用4bit量化,GPU上的參數(shù)只需要24GB的顯存環(huán)境,這樣的消耗只需要一張4090就能滿足。
此外通過這樣的組合,還能夠大幅度提升整個推理的性能,達(dá)到286 token/s的預(yù)填充和14 token/s的生成速度,比llama.cpp快28倍。
具體到技術(shù)實現(xiàn)中,團(tuán)隊采用了基于計算強度的offload策略、高性能的CPU和GPU算子、CUDA Graph加速的多種方式來加速推理速度。
基于計算強度的offload策略
在Attention的核心,DeepSeek引入了一種新的MLA算子,它能夠充分利用顯卡算力,能夠很大程度提升效率。然而,MLA運算符在官方開源的v2版本中,是將MLA展開成MHA進(jìn)行的計算,這個過程不僅擴(kuò)大了KV cache大小,還降低了推理性能。
為了真正發(fā)揮MLA的性能,在KTransformers推理框架中,團(tuán)隊將矩陣直接吸收到q_proj和out_proj權(quán)重中。因此,壓縮表示不需要解壓縮來計算Attention。
這種調(diào)整顯著減少了KV緩存大小,并增加了該運算符的算術(shù)強度,這非常顯著地優(yōu)化了GPU計算能力的利用率。
在計算中,MLA和Expert的計算強度相差數(shù)千倍。因此,團(tuán)隊通過計算強度來決定劃分策略,優(yōu)先將計算強度高的放入GPU(MLA > Shared Expert > Routed Expert),直到GPU放不下為止。
引入CPU和GPU的高性能算子
在CPU算子中,團(tuán)隊使用llamafile作為CPU內(nèi)核,使用expert并行和其他優(yōu)化,組成高性能算子框架CPUInfer。此外增加多線程、任務(wù)調(diào)度、負(fù)載均衡、NUMA感知等優(yōu)化。
在GPU算子的使用上,團(tuán)隊引入Marlin算子作為GPU計算的內(nèi)核,它能夠非常高效地進(jìn)行量化后的矩陣計算,和torch這些計算量化后的矩陣乘法的庫相比,使用Marlin算子完成在GPU上面的計算大概可以達(dá)到3.87倍的理想加速效果。
CUDA Graph的改進(jìn)和優(yōu)化
為了平衡推理性能和框架本身的易用性/可擴(kuò)展性,基于Python構(gòu)建KTransformers框架,同時使用CUDA Graph降低Python調(diào)用開銷是一個必然的選擇。
KTransformers中使用CUDA Graph過程中盡可能地減少了CPU/GPU通訊造成的斷點,在CUDA Graph中摻雜和CPU異構(gòu)算子通訊,最終實現(xiàn)一次decode僅有一個完整的CUDA Graph調(diào)用的結(jié)果。
靈活高效的推理實驗平臺
值得關(guān)注的是,KTransformers不止是一個固定的推理框架,也不只能推理DeepSeek的模型,它可以兼容各式各樣的MoE模型和算子,能夠集成各種各樣的算子,做各種組合的測試。
此外還同時提供了Windows、Linux的平臺的支持,方便運行。
當(dāng)大模型不斷往上卷,KTransformers用異構(gòu)計算打開一條新的推理路徑?;诖耍蒲泄ぷ髡邿o需巨額預(yù)算也能夠探索模型本質(zhì)。
GitHub 地址:https://github.com/kvcache-ai/ktransformers
具體技術(shù)細(xì)節(jié)指路:https://zhuanlan.zhihu.com/p/714877271
特別聲明:以上內(nèi)容(如有圖片或視頻亦包括在內(nèi))為自媒體平臺“網(wǎng)易號”用戶上傳并發(fā)布,本平臺僅提供信息存儲服務(wù)。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.