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