DeepSeek開源MoE新利器LPLB：基于線性規(guī)劃，專攻訓(xùn)練動態(tài)負載不平衡

剛剛DeepSeek在GitHub開源了LPLB（Linear-Programming-Based Load Balancer）。這是一個基于線性規(guī)劃的并行負載均衡器，旨在優(yōu)化MoE（混合專家）模型的專家并行工作負載分配

看起來 DeepSeek 和老黃的思路是一致的

英偉達在一個由 NVLink 連接的 1 萬張 GPU 集群里，用完全一樣的機制來給不同 kernel 分配 SM（GPU 的計算單元：Streaming Multiprocessors）。DeepSeek 做的事也一樣，只不過它把這個調(diào)度機制往上抽象了一層，做到了整個 pipeline 級別

目前該項目處于早期研究階段，性能提升仍在評估中。

核心功能與實現(xiàn)

LPLB主要通過以下機制實現(xiàn)動態(tài)負載均衡：

動態(tài)重排序：基于工作負載統(tǒng)計信息對專家進行動態(tài)重排序（該過程由嵌入的EPLB輔助）

副本構(gòu)建：考慮靜態(tài)拓撲結(jié)構(gòu)構(gòu)建專家副本

最優(yōu)Token分配：針對每個批次（Batch）求解最優(yōu)Token分配方案

在技術(shù)實現(xiàn)上，其內(nèi)置的LP（線性規(guī)劃）求解器實現(xiàn)了單SM內(nèi)點法（Interior Point Method, IPM），并利用NVIDIA的cuSolverDx和cuBLASDx庫進行線性代數(shù)運算。

工作負載統(tǒng)計信息可由用戶提供，通過torch.distributed收集，或從Deep-EP緩沖區(qū)的內(nèi)部通信器獲取。

工作原理

LPLB是對EPLB（Expert Parallelism Load Balancer）的擴展，旨在解決MoE訓(xùn)練中的動態(tài)負載不平衡問題：

EPLB：主要處理由數(shù)據(jù)分布引起的靜態(tài)不平衡（如某些專家持續(xù)過載）。

LPLB：針對訓(xùn)練過程中小批次隨機性引起的每批次波動

具體機制：

1.冗余專家：每個冗余專家鏈接到一個原始專家，在GPU之間形成邊（Edge）

2.邊容量：邊的容量定義為當(dāng)前批次分配給冗余專家的Token數(shù)量，即用于平衡的最大Token流

3.LP優(yōu)化：LPLB求解線性規(guī)劃問題，在尊重邊容量的前提下沿這些邊重新分配Token，以最小化專家并行（EP）組內(nèi)的負載不平衡。

在該過程中，待復(fù)制的專家通過EPLB選擇（僅重排序，不復(fù)制），最重的專家根據(jù)選定的LPLB拓撲進行復(fù)制。為了減少通信開銷，實時工作負載同步利用NVLINK和NVSHMEM（需預(yù)裝DeepEP），而非torch.distributed.allreduce。

支持的拓撲結(jié)構(gòu)

LPLB支持通過修改r2o矩陣探索自定義拓撲，典型拓撲包括：

Cube：在GPU子集上復(fù)制專家，形成帶有對角邊的立方體圖。每GPU至少需要2個專家。適用于8-GPU EP子組內(nèi)的平衡，且不犧牲節(jié)點間通信

Hypercube：類似于Cube，但排除對角邊，需要16個GPU。適用于跨16個GPU的專家并行

Torus：在同一節(jié)點的鄰居GPU和鄰居節(jié)點的GPU上各復(fù)制一個專家，形成環(huán)面圖。每GPU至少需要2個專家。適用于全局平衡，但由于節(jié)點內(nèi)通信效率原因，效果可能不如Cube

局限性

成本估算：目前的規(guī)劃器僅平衡總Token數(shù)量，未考慮分組矩陣乘法時間成本的非線性，可能導(dǎo)致次優(yōu)性能

求解延遲：求解器進行節(jié)點內(nèi)優(yōu)化耗時約100 μs（節(jié)點間更長），對于小批次任務(wù)，此開銷不可忽略

極端不平衡：在全局負載極端不平衡的情況下，由于LPLB避免將多個副本分配給同一原始專家，其表現(xiàn)可能不如EPLB

安裝與使用

預(yù)備條件：

CUDA Toolkit >= 12.6.3（包含cuSolverDx依賴）。

DeepEP（可選，但強烈建議用于實際生產(chǎn)）。

EPLB（已嵌入）

安裝命令：

                                                                    ./download-mathdx.sh
# export NVSHMEM_DIR=...  # 可選
pip install --no-build-isolation .

接口示例：

                                                                    # 定義冗余專家拓撲
r2o = torch.tensor(
    [
        [3, 0, 1, 2, 7, 4, 5, 6],
        [6, 7, 4, 5, 0, 1, 2, 3],
    ]
).T.int().cuda()

planner = Planner(
    r2o,
    n_logical_experts + n_redundants_per_rank * ep_size,
    n_logical_experts,
    group=ep_group,
)

# 規(guī)劃器返回物理專家索引
redirected_indices = planner.run(indices, avail_counter, N_SMS)

項目地址：https://github.com/deepseek-ai/LPLB