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

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

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

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

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

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

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

副本構(gòu)建：考慮靜態(tài)拓?fù)浣Y(jié)構(gòu)構(gòu)建專家副本

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

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

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

工作原理

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

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

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

具體機(jī)制：

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

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

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

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

支持的拓?fù)浣Y(jié)構(gòu)

LPLB支持通過修改r2o矩陣探索自定義拓?fù)?，典型拓?fù)浒ǎ?/p>

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

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

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

局限性

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

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

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

安裝與使用

預(yù)備條件：

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

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

EPLB（已嵌入）

安裝命令：

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

接口示例：

                                                                    # 定義冗余專家拓?fù)?br/>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)

項(xiàng)目地址：https://github.com/deepseek-ai/LPLB