DeepSeek悄悄開源LPLB：用線性規(guī)劃解決MoE負(fù)載不均
昨天，DeepSeek 在 GitHub 上線了一個(gè)新的代碼庫(kù)：LPLB。
項(xiàng)目地址：https://github.com/deepseek-ai/LPLB
沒(méi)有發(fā)推文，也沒(méi)有公眾號(hào)更新，少有的幾個(gè)技術(shù)博主分享的推文也關(guān)注不多。截至目前，該項(xiàng)目的 star 數(shù)量也還沒(méi)超過(guò) 200。
但仔細(xì)一看，這個(gè)項(xiàng)目卻似乎并不簡(jiǎn)單，值得更多關(guān)注。X 網(wǎng)友 gm8xx8 評(píng)論認(rèn)為這表明 DeepSeek 正在解決正確性和吞吐量瓶頸問(wèn)題，為下一版模型發(fā)布做準(zhǔn)備。
項(xiàng)目簡(jiǎn)介
LPLB，全稱 Linear-Programming-Based Load Balancer，即基于線性規(guī)劃的負(fù)載均衡器。
顧名思義，LPLB 是一個(gè)并行負(fù)載均衡器，它利用線性規(guī)劃（Linear Programming）算法來(lái)優(yōu)化 MoE（混合專家）模型中的專家并行工作負(fù)載分配。
具體來(lái)說(shuō)，LPLB 通過(guò)以下三個(gè)步驟實(shí)現(xiàn)動(dòng)態(tài)負(fù)載均衡：
- 動(dòng)態(tài)重排序： 基于工作負(fù)載統(tǒng)計(jì)信息對(duì)專家進(jìn)行重排序（Reordering）。
- 構(gòu)建副本： 結(jié)合靜態(tài)拓?fù)浣Y(jié)構(gòu)構(gòu)建專家副本（Replicas）。
- 求解最優(yōu)分配： 針對(duì)每個(gè)批次（Batch）的數(shù)據(jù)，求解最優(yōu)的 Token 分配方案。
更具體而言，LPLB 的專家重排序過(guò)程由 EPLB 協(xié)助完成。而實(shí)時(shí)工作負(fù)載統(tǒng)計(jì)信息可以由用戶提供、通過(guò) torch.distributed 收集，或直接從 Deep-EP 緩沖區(qū)的內(nèi)部通信器中獲取。至于求解器，則使用了內(nèi)置的 LP（線性規(guī)劃）求解器，其實(shí)現(xiàn)了單 SM（Streaming Multiprocessor）內(nèi)點(diǎn)法（IPM），并利用了 NVIDIA 的 cuSolverDx 和 cuBLASDx 庫(kù)進(jìn)行高效的線性代數(shù)運(yùn)算。
如此一來(lái)，MoE 負(fù)載不均的問(wèn)題可以得到有效解決，即在 MoE 模型中，某些「專家」可能比其他專家接收到更多的 Token，導(dǎo)致某些 GPU 忙碌而其他 GPU 空閑。
X 網(wǎng)友 big goose 指出這與英偉達(dá)的用于調(diào)度 SM (Streaming Multiprocessor，是英偉達(dá) GPU 的核心計(jì)算單元) 的方案非常相似，只是將抽象提升到了 pipeline 層級(jí)。LPLB 強(qiáng)調(diào)「單 SM」，意味著它的求解過(guò)程非常輕量化，不會(huì)占用過(guò)多計(jì)算資源。
不過(guò)需要指出，LPLB 目前應(yīng)該還未被用于生產(chǎn)流程。DeepSeek 在 Readme 文件中表示：「LPLB 目前處于早期研究階段，性能改進(jìn)情況仍在評(píng)估中?！?/span>
LPLB 的工作原理
LPLB 是在 EPLB（專家并行負(fù)載均衡器）基礎(chǔ)上的擴(kuò)展，旨在解決 MoE 訓(xùn)練中的動(dòng)態(tài)負(fù)載不均衡問(wèn)題。
EPLB vs. LPLB
- EPLB： 主要處理靜態(tài)不均衡（例如，由于數(shù)據(jù)分布特性，某些專家總是長(zhǎng)期過(guò)載）。
- LPLB： 專注于處理動(dòng)態(tài)波動(dòng)（由訓(xùn)練過(guò)程中小批次數(shù)據(jù)的隨機(jī)性引起的瞬時(shí)負(fù)載抖動(dòng)）。
核心機(jī)制
- 冗余專家 (Redundant Experts)： 每個(gè)冗余專家（副本）都鏈接到一個(gè)原始專家，從而在 GPU 之間形成連接邊。
- 邊容量 (Edge Capacity)： 一條邊的容量定義為當(dāng)前批次中分配給該冗余專家的 Token 數(shù)量，這決定了用于平衡負(fù)載的最大 Token 流量。
- LP 優(yōu)化 (LP Optimization)： LPLB 求解一個(gè)線性規(guī)劃問(wèn)題，在遵守邊容量限制的前提下，沿著這些邊重新分配 Token，以最小化專家并行（EP）組內(nèi)的負(fù)載不均衡。
實(shí)現(xiàn)流程
- 首先通過(guò) EPLB 選擇需要復(fù)制的專家（僅重排序，此時(shí)未復(fù)制）。
- 然后根據(jù)選定的 LPLB 拓?fù)浣Y(jié)構(gòu)，復(fù)制負(fù)載最重的專家。
- 通信優(yōu)化： 實(shí)時(shí)工作負(fù)載的同步使用 NVLINK 和 NVSHMEM 進(jìn)行優(yōu)化，替代了傳統(tǒng)的 torch.distributed.allreduce，從而大幅降低通信開銷。這正是需要預(yù)裝 DeepEP 的原因。
局限性
盡管 LPLB 提供了動(dòng)態(tài)優(yōu)化，但目前仍存在一些局限：
- 忽略非線性計(jì)算成本： 當(dāng)前的規(guī)劃器僅平衡 Token 總數(shù)，未考慮分組矩陣乘法（Grouped GEMM）時(shí)間成本的非線性特征。這可能導(dǎo)致在某些情況下性能并非絕對(duì)最優(yōu)。
- 求解延遲： 求解器在節(jié)點(diǎn)內(nèi)（intra-node）優(yōu)化大約需要 100 μs（跨節(jié)點(diǎn)時(shí)間更長(zhǎng)）。對(duì)于非常小的 Batch Size，這個(gè)延遲可能不可忽略。
- 極端不均衡情況： 在全局負(fù)載極端不均衡的情況下，LPLB 的表現(xiàn)可能不如 EPLB。這是因?yàn)?LPLB 在分配冗余專家時(shí)存在差異（LPLB 避免將多個(gè)副本分配給同一個(gè)原始專家）。
典型拓?fù)浣Y(jié)構(gòu)
LPLB 允許通過(guò)修改 r2o 矩陣來(lái)定義專家副本的分布方式。以下是幾種典型的拓?fù)洌?/span>
- 立方體 (Cube)：在 GPU 子集上復(fù)制專家，形成帶有對(duì)角邊的立方體圖。這要求每個(gè) GPU 至少 2 個(gè)專家。適用場(chǎng)景： 適合在 8 GPU 的 EP 子組內(nèi)進(jìn)行平衡，且不會(huì)犧牲跨節(jié)點(diǎn)通信性能。
- 超立方體 (Hypercube)：類似于 Cube，但不包含對(duì)角邊。這需要 16 個(gè) GPU。適用場(chǎng)景： 適合跨 16 個(gè) GPU 的專家并行。
- 環(huán)面 (Torus)：在同一節(jié)點(diǎn)內(nèi)的鄰居 GPU 上復(fù)制一個(gè)專家，在鄰節(jié)點(diǎn)的 GPU 上復(fù)制另一個(gè)專家，形成環(huán)面圖。其要求 每個(gè) GPU 至少 2 個(gè)專家。優(yōu)缺點(diǎn)： 對(duì)全局平衡有效，但由于涉及更多的節(jié)點(diǎn)內(nèi)通信，效率通常低于 Cube。
結(jié)語(yǔ)
DeepSeek 開源的這個(gè) LPLB 庫(kù)，本質(zhì)上是在試圖解決大模型訓(xùn)練中「木桶效應(yīng)」的問(wèn)題，即訓(xùn)練速度往往取決于最慢（負(fù)載最重）的那個(gè) GPU。
它的創(chuàng)新點(diǎn)在于引入了線性規(guī)劃這一數(shù)學(xué)工具來(lái)實(shí)時(shí)計(jì)算最優(yōu)分配，并利用底層的 NVSHMEM 技術(shù)來(lái)打破通信瓶頸。對(duì)于正在研究 MoE 架構(gòu)訓(xùn)練加速的開發(fā)者來(lái)說(shuō)，這是一個(gè)非常有價(jià)值的參考實(shí)現(xiàn)。
具體的安裝和測(cè)試指南請(qǐng)?jiān)L問(wèn)原代碼庫(kù)。
