![]()
刚刚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
可选
pip install --no-build-isolation .
接口示例:
规划器返回物理专家索引
redirected_indices = planner.run(indices, avail_counter, N_SMS)
项目地址:https://github.com/deepseek-ai/LPLB





京公网安备 11011402013531号