当前位置: 首页 » 资讯 » 新科技 » 正文

DeepSeek开源MoE新利器LPLB:基于线性规划,专攻训练动态负载不平衡

IP属地 中国·北京 AI寒武纪 时间:2025-11-20 10:14:40


刚刚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

免责声明:本网信息来自于互联网,目的在于传递更多信息,并不代表本网赞同其观点。其内容真实性、完整性不作任何保证或承诺。如若本网有任何内容侵犯您的权益,请及时联系我们,本站将会在24小时内处理完毕。

全站最新