DeepSeek悄悄开源LPLB：用线性规划解决MoE负载不均

昨天，DeepSeek 在 GitHub 上线了一个新的代码库：LPLB。

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

没有发推文，也没有公众号更新，少有的几个技术博主分享的推文也关注不多。截至目前，该项目的 star 数量也还没超过 200。

但仔细一看，这个项目却似乎并不简单，值得更多关注。X 网友 gm8xx8 评论认为这表明 DeepSeek 正在解决正确性和吞吐量瓶颈问题，为下一版模型发布做准备。

项目简介

LPLB，全称 Linear-Programming-Based Load Balancer，即基于线性规划的负载均衡器。

顾名思义，LPLB 是一个并行负载均衡器，它利用线性规划（Linear Programming）算法来优化 MoE（混合专家）模型中的专家并行工作负载分配。

具体来说，LPLB 通过以下三个步骤实现动态负载均衡：

1. 动态重排序： 基于工作负载统计信息对专家进行重排序（Reordering）。
2. 构建副本： 结合静态拓扑结构构建专家副本（Replicas）。
3. 求解最优分配： 针对每个批次（Batch）的数据，求解最优的 Token 分配方案。

更具体而言，LPLB 的专家重排序过程由 EPLB 协助完成。而实时工作负载统计信息可以由用户提供、通过 torch.distributed 收集，或直接从 Deep-EP 缓冲区的内部通信器中获取。至于求解器，则使用了内置的 LP（线性规划）求解器，其实现了单 SM（Streaming Multiprocessor）内点法（IPM），并利用了 NVIDIA 的 cuSolverDx 和 cuBLASDx 库进行高效的线性代数运算。

如此一来，MoE 负载不均的问题可以得到有效解决，即在 MoE 模型中，某些「专家」可能比其他专家接收到更多的 Token，导致某些 GPU 忙碌而其他 GPU 空闲。

X 网友 big goose 指出这与英伟达的用于调度 SM (Streaming Multiprocessor，是英伟达 GPU 的核心计算单元) 的方案非常相似，只是将抽象提升到了 pipeline 层级。LPLB 强调「单 SM」，意味着它的求解过程非常轻量化，不会占用过多计算资源。

不过需要指出，LPLB 目前应该还未被用于生产流程。DeepSeek 在 Readme 文件中表示：「LPLB 目前处于早期研究阶段，性能改进情况仍在评估中。」

LPLB 的工作原理

LPLB 是在 EPLB（专家并行负载均衡器）基础上的扩展，旨在解决 MoE 训练中的动态负载不均衡问题。

EPLB vs. LPLB

• EPLB： 主要处理静态不均衡（例如，由于数据分布特性，某些专家总是长期过载）。
• LPLB： 专注于处理动态波动（由训练过程中小批次数据的随机性引起的瞬时负载抖动）。

核心机制

• 冗余专家 (Redundant Experts)： 每个冗余专家（副本）都链接到一个原始专家，从而在 GPU 之间形成连接边。
• 边容量 (Edge Capacity)： 一条边的容量定义为当前批次中分配给该冗余专家的 Token 数量，这决定了用于平衡负载的最大 Token 流量。
• LP 优化 (LP Optimization)： LPLB 求解一个线性规划问题，在遵守边容量限制的前提下，沿着这些边重新分配 Token，以最小化专家并行（EP）组内的负载不均衡。

实现流程

• 首先通过 EPLB 选择需要复制的专家（仅重排序，此时未复制）。
• 然后根据选定的 LPLB 拓扑结构，复制负载最重的专家。
• 通信优化： 实时工作负载的同步使用 NVLINK 和 NVSHMEM 进行优化，替代了传统的 torch.distributed.allreduce，从而大幅降低通信开销。这正是需要预装 DeepEP 的原因。

局限性 

尽管 LPLB 提供了动态优化，但目前仍存在一些局限：

1. 忽略非线性计算成本： 当前的规划器仅平衡 Token 总数，未考虑分组矩阵乘法（Grouped GEMM）时间成本的非线性特征。这可能导致在某些情况下性能并非绝对最优。
2. 求解延迟： 求解器在节点内（intra-node）优化大约需要 100 µs（跨节点时间更长）。对于非常小的 Batch Size，这个延迟可能不可忽略。
3. 极端不均衡情况： 在全局负载极端不均衡的情况下，LPLB 的表现可能不如 EPLB。这是因为 LPLB 在分配冗余专家时存在差异（LPLB 避免将多个副本分配给同一个原始专家）。

典型拓扑结构

LPLB 允许通过修改 r2o 矩阵来定义专家副本的分布方式。以下是几种典型的拓扑：

• 立方体 (Cube)：在 GPU 子集上复制专家，形成带有对角边的立方体图。这要求每个 GPU 至少 2 个专家。适用场景： 适合在 8 GPU 的 EP 子组内进行平衡，且不会牺牲跨节点通信性能。
• 超立方体 (Hypercube)：类似于 Cube，但不包含对角边。这需要 16 个 GPU。适用场景： 适合跨 16 个 GPU 的专家并行。
• 环面 (Torus)：在同一节点内的邻居 GPU 上复制一个专家，在邻节点的 GPU 上复制另一个专家，形成环面图。其要求 每个 GPU 至少 2 个专家。优缺点： 对全局平衡有效，但由于涉及更多的节点内通信，效率通常低于 Cube。

结语

DeepSeek 开源的这个 LPLB 库，本质上是在试图解决大模型训练中「木桶效应」的问题，即训练速度往往取决于最慢（负载最重）的那个 GPU。

它的创新点在于引入了线性规划这一数学工具来实时计算最优分配，并利用底层的 NVSHMEM 技术来打破通信瓶颈。对于正在研究 MoE 架构训练加速的开发者来说，这是一个非常有价值的参考实现。

具体的安装和测试指南请访问原代码库。

文章来自于“机器之心”，作者 “Panda”。