用異構計算打開一條新的推理路徑
原標題:4090單卡跑滿血版DeepSeek-R1,清華團隊開源項目再破大模型推理門檻
文章來源:量子位
內容字數:4243字
清華大學開源項目KTransformers:讓千億參數大模型“飛入尋常百姓家”
DeepSeek-R1等千億參數大模型因其強大的性能而備受關注,但高昂的推理成本和對硬件資源的極高需求,使其難以在中小團隊甚至個人手中運行。清華大學KVCache.AI團隊聯合趨境科技發布的KTransformers開源項目,為這一難題提供了一種有效的解決方案。該項目通過一系列技術創新,實現了在24GB顯存的消費級顯卡上運行DeepSeek-R1/V3的671B參數滿血版,將千億級模型的推理帶入“家庭化”時代。
突破顯存限制:專家卸載技術與異構計算
KTransformers的核心技術在于“專家卸載”技術,它巧妙地利用MoE(混合專家)架構的稀疏性。DeepSeek-R1/V3采用MoE架構,模型的任務分配給不同的專家模塊,每次推理僅激活部分參數。KTransformers將非共享的稀疏MoE矩陣放在CPU/DRAM上,使用llamafile的高速算子處理;將稠密部分放在GPU上使用Marlin算子處理,實現了GPU/CPU的異構計算劃分策略。這種策略使得在僅需24GB顯存的環境下就能運行671B參數的模型,例如單張4090顯卡即可滿足需求。
顯著提升推理速度:高性能算子與CUDA Graph優化
KTransformers不僅解決了顯存問題,還大幅提升了推理速度。它采用基于計算強度的offload策略,優先將計算強度高的部分(MLA > Shared Expert > Routed Expert)放入GPU。同時,引入llamafile作為CPU內核和Marlin算子作為GPU內核,并進行多線程、任務調度等優化,實現了高達286 tokens/s的預填充速度和14 tokens/s的生成速度。此外,KTransformers通過優化CUDA Graph,減少了CPU/GPU通訊開銷,進一步提升了性能。
易用性與擴展性:友好的API和靈活的框架
為了降低使用門檻,KTransformers提供了兼容HuggingFace Transformers的API和ChatGPT式Web界面。同時,其基于YAML的“模板注入框架”支持靈活切換量化策略、內核替換等多種優化方式,方便用戶進行模型和算子的組合測試。該項目支持Windows和Linux平臺,并已在localLLaMa社區持續位居熱榜第一,擁有活躍的開發者社區。
技術細節:MLA算子優化與算子選擇策略
KTransformers對MLA算子進行了優化,將矩陣直接吸收到q_proj和out_proj權重中,減少了KV緩存大小并提升了GPU利用率。通過計算強度來決定劃分策略,將高計算強度任務優先分配給GPU,充分發揮GPU和CPU的算力。
總而言之,KTransformers通過巧妙的架構設計和優化策略,成功地將千億參數大模型的推理門檻大幅降低,為廣大開發者和科研人員提供了更便捷、更經濟的探索和應用途徑,推動了大模型技術的普及和發展。
聯系作者
文章來源:量子位
作者微信:
作者簡介:追蹤人工智能新趨勢,關注科技行業新突破