昨天，DeepSeek在GitHub上線了一個新的代碼庫：LPLB。

項目地址：https://github.com/deepseek-ai/LPLB
它沒有發(fā)推文，也沒有公眾號更新，少有的幾個技術(shù)博主分享的推文關(guān)注量也不多。截至目前，該項目的star數(shù)量還沒超過200。
但仔細(xì)探究，這個項目似乎并不簡單，值得更多關(guān)注。X網(wǎng)友gm8xx8評論認(rèn)為這表明DeepSeek正在解決正確性和吞吐量瓶頸問題，為下一版模型發(fā)布做準(zhǔn)備。

項目簡介

LPLB，全稱Linear-Programming-Based Load Balancer，即基于線性規(guī)劃的負(fù)載均衡器。
顧名思義，LPLB是一個并行負(fù)載均衡器，它利用線性規(guī)劃（Linear Programming）算法來優(yōu)化MoE（混合專家）模型中的專家并行工作負(fù)載分配。
具體而言，LPLB通過以下三個步驟實現(xiàn)動態(tài)負(fù)載均衡：
動態(tài)重排序：基于工作負(fù)載統(tǒng)計信息對專家進(jìn)行重排序（Reordering）。
構(gòu)建副本：結(jié)合靜態(tài)拓?fù)浣Y(jié)構(gòu)構(gòu)建專家副本（Replicas）。
求解最優(yōu)分配：針對每個批次（Batch）的數(shù)據(jù)，求解最優(yōu)的Token分配方案。
更詳細(xì)地說，LPLB的專家重排序過程由EPLB協(xié)助完成。而實時工作負(fù)載統(tǒng)計信息可以由用戶提供、通過torch.distributed收集，或直接從Deep - EP緩沖區(qū)的內(nèi)部通信器中獲取。至于求解器，則使用了內(nèi)置的LP（線性規(guī)劃）求解器，其實現(xiàn)了單SM（Streaming Multiprocessor）內(nèi)點法（IPM），并利用了NVIDIA的cuSolverDx和cuBLASDx庫進(jìn)行高效的線性代數(shù)運算。
這樣一來，MoE負(fù)載不均的問題能得到有效解決。在MoE模型中，某些「專家」可能比其他專家接收到更多的Token，導(dǎo)致某些GPU忙碌而其他GPU空閑。
X網(wǎng)友big goose指出這與英偉達(dá)的用于調(diào)度SM (Streaming Multiprocessor，是英偉達(dá)GPU的核心計算單元) 的方案非常相似，只是將抽象提升到了pipeline層級。LPLB強調(diào)「單SM」，意味著它的求解過程非常輕量化，不會占用過多計算資源。

不過需要指出，LPLB目前應(yīng)該還未被用于生產(chǎn)流程。DeepSeek在Readme文件中表示：「LPLB目前處于早期研究階段，性能改進(jìn)情況仍在評估中?！?h2>LPLB的工作原理LPLB是在EPLB（專家并行負(fù)載均衡器）基礎(chǔ)上的擴展，旨在解決MoE訓(xùn)練中的動態(tài)負(fù)載不均衡問題。
EPLB vs. LPLB
EPLB：主要處理靜態(tài)不均衡（例如，由于數(shù)據(jù)分布特性，某些專家總是長期過載）。
LPLB：專注于處理動態(tài)波動（由訓(xùn)練過程中小批次數(shù)據(jù)的隨機性引起的瞬時負(fù)載抖動）。
核心機制
冗余專家 (Redundant Experts)：每個冗余專家（副本）都鏈接到一個原始專家，從而在GPU之間形成連接邊。
邊容量 (Edge Capacity)：一條邊的容量定義為當(dāng)前批次中分配給該冗余專家的Token數(shù)量，這決定了用于平衡負(fù)載的最大Token流量。
LP優(yōu)化 (LP Optimization)：LPLB求解一個線性規(guī)劃問題，在遵守邊容量限制的前提下，沿著這些邊重新分配Token，以最小化專家并行（EP）組內(nèi)的負(fù)載不均衡。

實現(xiàn)流程

首先通過EPLB選擇需要復(fù)制的專家（僅重排序，此時未復(fù)制）。
然后根據(jù)選定的LPLB拓?fù)浣Y(jié)構(gòu)，復(fù)制負(fù)載最重的專家。
通信優(yōu)化：實時工作負(fù)載的同步使用NVLINK和NVSHMEM進(jìn)行優(yōu)化，替代了傳統(tǒng)的torch.distributed.allreduce，從而大幅降低通信開銷。這正是需要預(yù)裝DeepEP的原因。

局限性

盡管LPLB提供了動態(tài)優(yōu)化，但目前仍存在一些局限：
忽略非線性計算成本：當(dāng)前的規(guī)劃器僅平衡Token總數(shù)，未考慮分組矩陣乘法（Grouped GEMM）時間成本的非線性特征。這可能導(dǎo)致在某些情況下性能并非絕對最優(yōu)。
求解延遲：求解器在節(jié)點內(nèi)（intra - node）優(yōu)化大約需要100 μs（跨節(jié)點時間更長）。對于非常小的Batch Size，這個延遲可能不可忽略。
極端不均衡情況：在全局負(fù)載極端不均衡的情況下，LPLB的表現(xiàn)可能不如EPLB。這是因為LPLB在分配冗余專家時存在差異（LPLB避免將多個副本分配給同一個原始專家）。

典型拓?fù)浣Y(jié)構(gòu)

LPLB允許通過修改r2o矩陣來定義專家副本的分布方式。以下是幾種典型的拓?fù)洌?br>立方體 (Cube)：在GPU子集上復(fù)制專家，形成帶有對角邊的立方體圖。這要求每個GPU至少2個專家。適用場景：適合在8 GPU的EP子組內(nèi)進(jìn)行平衡，且不會犧牲跨節(jié)點通信性能。
超立方體 (Hypercube)：類似于Cube，但不包含對角邊。這需要16個GPU。適用場景：適合跨16個GPU的專家并行。
環(huán)面 (Torus)：在同一節(jié)點內(nèi)的鄰居GPU上復(fù)制一個專家，在鄰節(jié)點的GPU上復(fù)制另一個專家，形成環(huán)面圖。其要求每個GPU至少2個專家。優(yōu)缺點：對全局平衡有效，但由于涉及更多的節(jié)點內(nèi)通信，效率通常低于Cube。

結(jié)語

DeepSeek開源的這個LPLB庫，本質(zhì)上是在試圖解決大模型訓(xùn)練中「木桶效應(yīng)」的問題，即訓(xùn)練速度往往取決于最慢（負(fù)載最重）的那個GPU。
它的創(chuàng)新點在于引入了線性規(guī)劃這一數(shù)學(xué)工具來實時計算最優(yōu)分配，并利用底層的NVSHMEM技術(shù)來打破通信瓶頸。對于正在研究MoE架構(gòu)訓(xùn)練加速的開發(fā)者來說，這是一個非常有價值的參考實現(xiàn)。
具體的安裝和測試指南請訪問原代碼庫。
本文來自微信公眾號“機器之心”（ID：almosthuman2014），作者：機器之心，36氪經(jīng)授權(quán)發(fā)布。

本文僅代表作者觀點，版權(quán)歸原創(chuàng)者所有，如需轉(zhuǎn)載請在文中注明來源及作者名字。

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