DeepSeek開源MoE新利器LPLB：基于線性規劃，專攻訓練動態負載不平衡

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

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

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

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

核心功能與實現

LPLB主要通過以下機制實現動態負載均衡：

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

副本構建：考慮靜態拓撲結構構建專家副本

最優Token分配：針對每個批次（Batch）求解最優Token分配方案

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

工作負載統計信息可由用戶提供，通過torch.distributed收集，或從Deep-EP緩沖區的內部通信器獲取。

工作原理

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

EPLB：主要處理由數據分布引起的靜態不平衡（如某些專家持續過載）。

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

具體機制：

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

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

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

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

支持的拓撲結構

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

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

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

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

局限性

成本估算：目前的規劃器僅平衡總Token數量，未考慮分組矩陣乘法時間成本的非線性，可能導致次優性能

求解延遲：求解器進行節點內優化耗時約100 μs（節點間更長），對于小批次任務，此開銷不可忽略

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

安裝與使用

預備條件：

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

DeepEP（可選，但強烈建議用于實際生產）。

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,
)

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

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