LPLB: An early research stage MoE load balancer based on linear programming
3 months ago
- #linear programming
- #Mixture-of-Experts
- #load balancing
- LPLB是一种采用线性规划优化的并行负载均衡器，专为MoE模型的工作负载分配而设计
- 它通过动态重排专家节点、构建副本并求解最优令牌分配方案，实现动态负载均衡
- EPLB组件负责专家重排序功能，工作负载统计信息可由用户提供或通过torch.distributed收集
- LPLB实现了单SM内点法(IPM)优化器，并采用NVIDIA的cuSolverDx和cuBLASDx数学库
- 运行环境要求包括：CUDA工具包>=12.6.3、推荐安装DeepEP框架及嵌入式EPLB组件
- LPLB在EPLB基础上扩展了动态负载均衡能力，重点解决MoE训练中逐批次波动问题
- 冗余专家节点与原始专家建立关联，其边容量由令牌分配方案动态定义
- 通过线性规划优化令牌重分配，最小化专家并行组内的负载不均衡
- 采用NVLINK和NVSHMEM技术优化工作负载同步，显著降低通信开销
- 当前版本局限包括：仅平衡令牌数量、求解器延迟问题，极端不均衡时可能表现欠佳
- 支持立方体(Cube)、超立方体(Hypercube)和环面(Torus)三种拓扑结构，适配不同GPU配置
- 用户可通过修改r2o矩阵探索自定义拓扑结构