Hotdry.
归档
ai-systems

LPLB：基于线性规划的MoE负载均衡器

将MoE token到专家路由建模为线性规划，利用单SM IPM求解器实现动态负载均衡、最小化GPU间排队延迟和高吞吐分布式推理。

在 MoE（Mixture-of-Experts）模型的分布式训练中，专家负载不均衡是常见瓶颈，尤其是小批量训练时，每批 token 分配的随机波动会导致部分 GPU 空闲、整体吞吐低下。LPLB（Linear-Programming-Based Load Balancer）通过将 token-to-expert 路由问题形式化为线性规划（LP）问题，并在 GPU 上高效求解，提供了一种动态负载均衡方案。该方法针对每批数据重新分配 token 至冗余专家副本，实现最小化最大负载的目标，同时尊重边容量约束。

LP 问题形式化：从图模型到优化目标

LPLB 的核心是将 MoE 专家并行（EP）组内的负载均衡建模为一个带容量约束的最大流最小成本问题，但简化成 min-max 负载的 LP 形式。给定 EP 组大小为 ep_size，每个原始专家 i 有 d_i 个 token 需求，通过 EPLB 预先重排序选择最重专家，并构建 r2o 矩阵定义冗余专家拓扑（redundant-to-original）。每个冗余专家 j 链接至原始专家，形成 GPU 间 “边”，边容量 c_j 为其在本批可处理的 token 上限（基于可用 SM 数 N_SMS）。

优化目标：min T（最大负载），约束包括：

  • 流量守恒：每个原始专家流出量等于 d_i。
  • 容量约束：每边流量 ≤ c_j。
  • 非负变量：流量 x_e ≥ 0。

变量数通常 < 128（典型 Cube 拓扑 8GPU×2 冗余 / 卡），求解器采用单 SM 内点法（IPM），调用 cuSolverDx 进行 Cholesky 分解、cuBLASDx 加速 GEMM，全程 GPU kernel 内完成，无 CPU 介入。intra-node 求解～100µs，适合大批量；inter-node 稍长，但 NVLINK/NVSHMEM 优化通信（替代 torch.distributed.allreduce）。

证据显示，在 DeepSeek 的早期评估中，此 LP 路由可显著降低动态抖动引起的 “木桶效应”。例如，r2o 矩阵示例：

r2o = torch.tensor([
    [3, 0, 1, 2, 7, 4, 5, 6],
    [6, 7, 4, 5, 0, 1, 2, 3],
]).T.int().cuda()

定义 Cube 拓扑，对角边提升 intra-node 效率。

集成与参数调优：可落地部署清单

  1. 环境准备：

    • CUDA ≥12.6.3（含 cuSolverDx/cuBLASDx）。
    • DeepEP（强烈推荐，用于 buffer 内统计采集与 NVSHMEM 计数器）。
    • EPLB 嵌入式，无需额外安装。
    ./download-mathdx.sh
pip install --no-build-isolation -e .

  2. Planner 初始化：

    planner = Planner(r2o, n_physical, n_logical, group=ep_group)
# 可选：planner.init_from_deep_ep(buffer)  # 零额外延迟采集
    • n_physical = n_logical + n_redundants_per_rank * ep_size（e.g., 8 逻辑 + 16 冗余 = 24）。
    • group=ep_group，确保 EP 内 allgather。

  3. 每批调用：

    redirected_indices = planner.run(logical_indices, avail_counter, N_SMS=100)
    • N_SMS：占用 SM 数，调优阈值：batch_size * seq_len /hidden_dim 估算专家负载，N_SMS=64~128（H100 单卡 108SM，预留裕量）。
    • avail_counter：torch.int64 tensor，累积成功调度计数。

  4. 拓扑选择参数：

    拓扑 GPU 需求 优势 适用场景
    Cube ≥8/GPU≥2 专家 对角边，intra 高效 单节点 8 卡 EP
    Hypercube 16 均匀跨节点 多节点 16 卡
    Torus ≥8/GPU≥2 专家 全局平衡 跨机集群

    自定义 r2o：确保无多副本指同一原始专家，避免 LPLB 退化。

  5. 监控与阈值：

    • 指标：max_load /avg_load <1.2（目标）；solver_time < batch_compute_time * 0.01。
    • 回滚策略：若 solver_time >50µs 或 imbalance>2x，fallback EPLB（静态重排序）。
    • 风险阈值：小 batch (<1k token/GPU) 禁用 LPLB；极端全局倾斜（全 EP 超载）时 EPLB 优于 LPLB。

局限与优化路径

当前仅均衡 token 计数，忽略 grouped GEMM 非线性（短序列 padding 开销），可扩展为带权 LP（权重 = 估算 flops）。求解延迟对 micro-batch 占比高，未来 warm-start（上批解初始化）或学习式剪枝可降至 10µs。集成 DeepEP buffer 采集路径零开销，NVSHMEM 放计数器于 NVLINK，提升多节点同步。

实际落地中，从 8 卡 Cube 起步，监控 1epoch 后对比 baseline（无 LPLB）MFU 提升 5-15%。结合 DeepSpeed-MoE 或 vLLM 推理栈，扩展至 serving 场景（动态路由高吞吐）。

资料来源： [1] https://github.com/deepseek-ai/LPLB (官方 repo 与 README)。 [2] EPLB/DeepEP 相关文档，用于拓扑与统计集成。

查看归档