diff --git a/docs/CN/source/framework/eplb.md b/docs/CN/source/framework/eplb.md new file mode 100644 index 0000000000..23fea0c522 --- /dev/null +++ b/docs/CN/source/framework/eplb.md @@ -0,0 +1,546 @@ +# EPLB 专家负载均衡实现 + +本文介绍 LightLLM 中 Expert Parallelism Load Balancer(EPLB)的完整实现,包括物理专家槽位、在线路由、负载采集、布局规划、权重迁移、运行时状态机、布局持久化,以及如何扩展新的规划算法。 + +EPLB 的目标是在不改变模型逻辑专家语义的前提下,利用额外的物理专家副本缓解热点专家造成的 EP rank 负载不均。它把“模型选择了哪个逻辑专家”和“本次由哪个物理副本执行”分成两个阶段,并允许服务运行期间重新安排物理副本。 + +## 1. 核心概念 + +设: + +- `E`:模型每层的逻辑专家数; +- `W`:EP world size; +- `R`:每个 rank 配置的冗余专家数; +- `E / W`:每个 rank 原本持有的专家数; +- `E / W + R`:每个 rank 实际分配的物理专家槽位数; +- `E + W * R`:一个 MoE 层在整个 EP world 中的物理槽位总数。 + +逻辑专家 ID 来自模型路由器,范围固定为 `[0, E)`。物理专家 ID 标识实际执行权重所在的槽位: + +```text +physical_expert_id = rank * num_physical_experts_per_rank + local_slot +``` + +一个逻辑专家可以拥有多个物理副本,但同一个 rank 上不会重复放置同一个逻辑专家。任意合法布局还必须满足: + +1. 每个 rank 的物理槽位数相同; +2. 所有逻辑专家至少有一个物理副本; +3. 所有逻辑专家 ID 都在 `[0, E)` 范围内; +4. 同一 rank 内的逻辑专家 ID 不重复。 + +## 2. 启用方式 + +EPLB 通过冗余专家数量开启: + +```bash +python -m lightllm.server.api_server \ + --model_dir /path/to/model \ + --enable_ep_moe \ + --eplb_num_redundant_experts_per_rank 2 \ + --eplb_plan_mode greedy \ + --eplb_rebalance_count 1 \ + --eplb_config_path /path/to/eplb-placement.json +``` + +主要参数如下: + +| 参数 | 默认值 | 作用 | +| --- | --- | --- | +| `--enable_ep_moe` | 关闭 | 启用专家并行;EPLB 的前置条件 | +| `--eplb_num_redundant_experts_per_rank` | `0` | 每个 rank 的额外物理专家槽位数;大于 0 时启用 EPLB | +| `--eplb_plan_mode` | `greedy` | 选择动态布局规划算法;当前支持 `greedy` | +| `--eplb_rebalance_count` | `1` | 最多完成的动态重排次数;`-1` 表示不限次数,`0` 表示不动态重排 | +| `--eplb_config_path` | `None` | 可选的布局加载与回写路径 | + +完整命令行说明见 {doc}`../tutorial/api_server_args`。 + +在 PD 分离部署中,prefill 和 decode 进程各自拥有独立的 EPLB manager,可以分别设置 `--eplb_plan_mode`。同一个 EP 通信组内的所有 rank 必须使用相同配置。非 PD 部署只有一个 manager,它根据该进程采集到的全部路由负载生成统一布局。 + +## 3. 总体架构 + +```text +模型路由器 + │ + │ logical top-k IDs + v +EPLB 路由 kernel + ├── 把本次 prefill 的 logical expert 负载写入环形采样行 + ├── 查询 logical_to_physical_map + └── 为每个 token 选择 physical expert ID + │ + v + MoE 执行 kernel + +周期性控制面: + +prefill_route_counter(最近 24 次 prefill 采样) + -> 全局负载汇总 + -> placement planner + -> target placement + -> transfer planner + -> 后台权重传输 + -> 安全边界提交权重和路由 metadata +``` + +主要实现位置: + +| 模块 | 职责 | +| --- | --- | +| `fused_moe/impl/deepgemm_impl.py` | 初始化 EPLB 运行态,在 MoE 执行前修复 logical top-k IDs | +| `triton_kernel/fused_moe/eplb_topk_ids.py` | 统计逻辑专家负载,并把 logical ID 映射为 physical ID | +| `eplb/placement/initial.py` | 构建确定性的初始专家布局 | +| `eplb/placement/routing.py` | 根据完整布局构建紧凑路由表 | +| `eplb/placement/planner.py` | 布局规划器抽象接口 | +| `eplb/placement/factory.py` | 根据 `eplb_plan_mode` 创建具体规划器 | +| `eplb/placement/greedy.py` | 默认的贪心布局算法 | +| `eplb/async_task.py` | 统一后台线程任务的启动、完成与异常处理 | +| `eplb/async_load_gather_task.py` | 在独立 Gloo 通信组中后台汇集逐 rank、逐 sample 的原始负载 | +| `eplb/async_placement_plan_task.py` | 在后台根据负载生成目标专家布局 | +| `eplb/async_transfer_planner.py` | 在后台生成跨层传输批次 | +| `eplb/async_expert_transfer.py` | 规划槽位依赖并在后台执行专家权重传输 | +| `eplb/runtime_manager.py` | 驱动状态机,协调采集、规划、传输和提交 | + +## 4. 初始化布局与权重加载 + +### 4.1 默认布局 + +启动时首先把逻辑专家连续划分到各 rank,然后从下一个 rank 的主专家区间开始循环选择冗余副本。例如 `E=8`、`W=4`、`R=2` 时: + +```text +rank 0: [0, 1, 2, 3] +rank 1: [2, 3, 4, 5] +rank 2: [4, 5, 6, 7] +rank 3: [6, 7, 0, 1] +``` + +每行前两个槽位来自原始连续划分,后两个槽位是启动时已经加载完成的冗余副本。运行期允许重新分配所有物理槽位,不再区分不可移动的“主槽位”和只能替换的“冗余槽位”。 + +### 4.2 从历史布局启动 + +指定 `--eplb_config_path` 后,每个 MoE 层会尝试加载历史布局。配置必须同时匹配: + +- 配置版本; +- 逻辑专家数; +- world size; +- 每个 rank 的冗余专家数; +- 模型层号; +- 每层布局形状、专家 ID 范围、rank 内唯一性和全专家覆盖关系。 + +任意校验失败都会记录 warning,并仅对受影响的层回退到默认布局。校验通过时,专家权重会直接按照历史布局加载,不需要服务启动后再执行一次恢复迁移。 + +### 4.3 本地运行态 + +每个 MoE 实现对象持有: + +- `local_logics_expert_ids_list`:本 rank 每个物理槽对应的逻辑专家; +- `logical_to_physical_map`:logical ID 到可用 physical IDs 的设备路由表; +- `prefill_route_counter`:shape 为 `[24, E]` 的 `int64` GPU 环形采样缓冲区; +- `prefill_route_sample_index`:shape 为 `[2]` 的 `int64` GPU 状态,分别保存 sample index 和核内同步计数; +- `num_redundant_experts_per_rank`:本 rank 的额外槽位数。 + +目前启用 EP MoE 时使用 `FuseMoeDeepGEMM` 实现。EPLB manager 只收集启用了 EP 的 `layer.experts`,并保留模型中的层顺序。 + +## 5. 在线路由与负载采集 + +### 5.1 logical ID 与 physical ID 分离 + +MoE 路由器首先只在模型的逻辑专家空间中计算 top-k: + +```text +_select_experts + -> topk_weights + logical_topk_ids + -> capture callback + -> _prepare_expert_execution + -> EPLB logical-to-physical 映射 + -> _fused_experts +``` + +逻辑 ID tensor 不会被原地修改。监控和 capture callback 始终看到模型语义上的 logical expert;只有实际执行 MoE kernel 前才生成新的 physical ID tensor。 + +### 5.2 路由表布局 + +每个 logical expert 对应一行固定宽度 metadata: + +```text +[global_count, node_count, current_gpu_count, + physical_ids..., -1 padding...] +``` + +- `global_count`:整个 EP world 中的有效副本数; +- `node_count`:当前节点内的有效副本数,包含本卡; +- `current_gpu_count`:当前 GPU 上的有效副本数; +- `physical_ids`:按“本卡、本节点其他卡、其他节点”的顺序稳定排列; +- `padding`:未使用槽位填 `-1`,kernel 不会读取。 + +路由槽位上限等于整个 world 的物理槽位总数,因此布局变化不会改变 tensor 的 shape。 + +### 5.3 副本分发模式 + +路由算子要求调用方显式指定分发模式: + +| 模式 | 候选副本 | +| --- | --- | +| `current_gpu_first` | 本卡存在副本时只在本卡副本间选择,否则回退到全局副本 | +| `current_node_first` | 本节点存在副本时在节点内选择,否则回退到全局副本 | +| `global_first` | 直接在全局全部有效副本间选择 | + +当前 DeepGEMM EPLB 路径使用 `global_first`。未来如果要支持“本卡 -> 本节点 -> 全局”的三级回退,需要布局规划算法同时具备节点拓扑感知能力。 + +### 5.4 副本哈希 + +同一候选集合内使用 `(token_index, logical_expert_id)` 生成 32 位哈希,再对有效副本数取模。实现先用 logical expert ID 给 token index 加盐,然后执行 32 位 avalanche finalizer。 + +该变换由奇数乘法和可逆的异或移位组成,可以显著降低规律性 token 间隔与副本数之间的低位相关性。例如同一专家每隔 4 个 token 出现且有 4 个副本时,简单线性哈希可能退化到单一副本,avalanche mix 能将流量重新打散。 + +### 5.5 prefill 环形采样 + +负载采样只在调用方明确传入 `is_prefill=True` 时启用。decode 仍然执行 logical-to-physical 映射,但不会更新采样缓冲区,也不会推进 sample index。这样可以只使用吞吐量较大、统计稳定性更好的 prefill 路由结果,同时避免给高频 decode 路径增加原子操作。 + +每层使用一个 `[24, E]` 的 `prefill_route_counter`。其中每一行表示一次 prefill 路由 kernel 调用的 logical expert 直方图,24 表示最多保留最近 24 次采样,而不是 24 个 token 或 24 个 manager step。一次 manager step 内如果发生多次 prefill dispatch,它们会分别占用不同的采样行;第 25 次采样开始按环形方式覆盖最旧的数据: + +```text +sample_row = sample_index % 24 + +prefill_route_counter + row 0 -> 一次完整 prefill dispatch 的 [expert_0, ..., expert_E-1] 计数 + row 1 -> 下一次完整 prefill dispatch 的计数 + ... + row 23 -> 最近 24 次采样中的一行 +``` + +固定 24 行可以限制设备内存和 CPU 快照成本,并让规划器观察最近一段时间的流量,而不是让很早以前的流量永久影响当前布局。当前 manager 在评估时沿 sample 维求和,将 `[24, E]` 聚合回 `[E]`,因此现有 planner 接口无需感知环形缓冲区。 + +计数始终使用 logical expert ID,而不是最终选中的 physical expert ID。同一逻辑专家即使拥有多个物理副本,规划器看到的仍然是一份完整需求量,不会因为副本分发而被拆散。 + +### 5.6 无额外清零 kernel 的采样事务 + +环形行在复用前必须清零,否则新旧两次采样会叠加。为避免每次 prefill 额外发射一个清零 kernel,清零、路由计数和 sample index 提交都融合在 `_eplb_repair_topk_ids_kernel` 内;其中 `_record_prefill_route_sample` 是 Triton 子 JIT 函数,不会形成独立的 kernel launch。 + +`prefill_route_sample_index` 的两个元素含义如下: + +```text +[0] sample index:单调递增;对 24 取余得到当前环形行 +[1] sync state :0 表示目标行尚未清零 + 1 表示清零完成,采样可以开始 + 1 + completed_programs 表示已经完成的 program 数 +``` + +一次采样事务按以下顺序执行: + +1. 所有 program 读取同一个 sample index,并计算本次目标行。sample index 只由最后完成者推进,因此在本次 kernel 生命周期内保持不变。 +2. `program_id == 0` 清空目标行,然后通过带 `release` 语义的原子加一把 sync state 从 0 发布为 1。 +3. 其他 program 使用带 `acquire` 语义的原子读等待 sync state 达到 1,确保不会在清零完成前向目标行累加。 +4. 屏障通过后,每个 program 根据自己处理的有效 top-k 元素,对对应 logical expert 执行 `atomic_add(1)`。 +5. 每个 program 完成 physical ID 写回和负载计数后,再对 sync state 原子加一,提交一个完成信号。 +6. Triton 的 `atomic_add` 返回加法前的旧值,因此用 `old_value + 1 == num_programs + 1` 判断唯一的最后完成者。额外的 1 是步骤 2 发布的 ready 标记。 +7. 最后完成者先把 sample index 加一,再把 sync state 交换为 0,使下一次 kernel 可以复用后续环形行。 + +完整状态变化如下: + +```text +sync=0 + -> program 0 清零目标行 + -> sync=1(ready) + -> 所有 program 统计 logical expert 并分别提交完成信号 + -> sync=1+num_programs + -> 唯一最后完成者推进 sample index,并复位 sync=0 +``` + +ready 发布使用 `release`、等待方使用 `acquire`,最后完成信号使用 `acq_rel`,从而约束目标行清零和后续原子计数的可见顺序。所有调用还必须在同一 CUDA stream 上串行复用同一组 counter 和同步状态;当前 MoE forward 与采样都位于 overlap stream,满足这一约束。 + +### 5.7 manager 聚合与重置 + +manager 在安全推理边界一次性堆叠各层 `[24, E]` 环形缓冲区,并完整复制为 `[layer, sample, logical_expert]` CPU 快照。rank 0 先沿 sample 维聚合本地负载并判断样本量,再向所有 rank 广播是否继续规划;样本不足时直接返回采集状态,不发起大块通信。样本充足时,后台任务使用独立的 Gloo 通信组执行 all-gather,得到 `[rank, layer, sample, logical_expert]`,并把这个四维 Tensor 直接交给 planner。独立通信组使长时间运行的后台 all-gather 不会打乱主线程控制 collective 的调用顺序。 + +如果样本量不足或规划结果未改变布局,不主动清空缓冲区;后续 prefill 会继续写入,并在容量用满后滚动覆盖最旧行。 + +初始化、成功切换到新布局,以及达到重排次数上限后开始下一轮指标窗口时,manager 会同时清零 `prefill_route_counter` 和 `prefill_route_sample_index`。清零提交到 overlap stream,自然排在此前 forward 之后、后续 forward 之前,不需要额外的全设备同步。 + +## 6. EPLB 状态机 + +`EPLBManager.step()` 在安全的推理边界被调用,每次最多处理一个状态: + +```text +[COLLECTING] + 将 prefill logical route 写入 24 行环形采样,等待评估周期 + | + v +[EVALUATING] + CPU 原始样本快照、指标上报;rank 0 判断样本量并广播结果 + 样本充足且次数允许时启动后台 load all-gather + | + v +[WAIT_LOAD_GATHER_FINISHED] + 等待各 rank 汇集完成,保留 planner 需要的原始四维输入 + | + v +[PLAN_PLACEMENT] + rank 0 启动后台布局规划 + | + v +[WAIT_PLAN_PLACEMENT_FINISHED] + 轮询规划结果,并广播目标布局 + | + v +[PLAN_TRANSFER] + 各 rank 根据相同布局启动后台传输规划 + | + v +[WAIT_PLAN_TRANSFER_FINISHED] + 等待所有 rank 生成一致的传输批次 + | + v +[TRANSFERRING] + 分批启动/轮询权重传输,在安全边界统一提交 + | + `-------------------------------> COLLECTING +``` + +提前返回 `COLLECTING` 的分支: + +```text +EVALUATING + |-- rank 0 平均 token 数不足 ----> 保留环形窗口,继续滚动采样 + `-- 达到重排次数上限 ----------> 清空采样,只做周期性指标上报 + +WAIT_PLAN_PLACEMENT_FINISHED + `-- 目标布局与当前布局相同 -----> 保留环形窗口,等待下次评估 +``` + +默认每 20 个采样 step 评估一次,可以通过环境变量 `LIGHTLLM_EPLB_STEP_INTERVAL` 调整。该值必须大于 0。 + +只有当 rank 0 的平均样本量达到每个“层 × 逻辑专家”128 个 token 时才开始规划。各 rank 的路由分布高度相似,因此 rank 0 足以作为是否值得发起全量通信的低成本判断。样本不足不会清空环形缓冲区,低流量服务可以跨多个评估周期继续采样;缓冲区写满后只保留最近 24 次 prefill dispatch。 + +## 7. 专家分布分析 + +基于 DeepSeek-R1(EP8 + DP8,单机 8xH200)加载 ShareGPT 语料(6000 条,input 512-2048 token) +实测的路由负载分布,用于校准布局规划(第 8 节)的设计假设。 + +### 7.1 各 rank 分布相似性(实测) + +在 `EPLBManager` 全局汇总负载(后台原始 load `all_gather` 完成之后)时, +把各 rank 的 `[layer][logical_expert]` 负载逐层归一化为概率分布,以 rank0 的分布为基准, +与其余 rank 逐层计算 cosine 相似度。观测窗口为 warmup 阶段一次完整采样 +(58 个 MoE 层 × 7 对,共 406 对): + +| 指标 | 数值 | +| --- | --- | +| cosine 全体 min / 中位 / max | 0.9432 / 0.9728 / 0.9972 | +| 各 rank 对 rank0 的均值 | 0.971 ~ 0.977 | +| 最不相似层(层均值) | layer 43 (0.954)、32 (0.958)、41 (0.959) | +| 最相似层(层均值) | layer 0 (0.997)、1 (0.994)、4 (0.994) | + +探针日志示例: + +```text +eplb load prob cosine layer=0 vs_rank1..7: 0.9967 0.9965 0.9969 0.9962 0.9961 0.9972 0.9963 +eplb load prob cosine summary mean_by_rank(0..7): 1.0000 0.9773 0.9711 0.9745 0.9710 0.9762 0.9749 0.9735 +``` + +结论:**所有层上各 rank 的专家负载分布形状高度一致**。0.94~0.99 之间的小幅差异 +主要来自每 rank 仅承载约 1/8 流量的采样噪声,而非分布本身存在 rank 间异质性。 +DP 随机分流下每个 rank 的路由统计都是对全局路由分布的无偏采样, +分布形状(倾斜度、热点名单、长尾形态)在 rank 之间同源。 + +### 7.2 设计选择:异步汇集所有 rank 的原始样本 + +各 rank 的负载分布虽然高度相似,但单 rank 仍带有可观测的采样噪声。当前实现汇集所有 +rank 的 `[layer, sample, logical_expert]` 原始快照,并在通信完成后统一求和: + +1. **降低采样噪声**:规划器使用整个 world 的累计流量,热点排序和副本预算不依赖某个 + rank 的随机流量分片; +2. **保留分析信息**:通信结果在聚合前保留 rank 和 sample 维,后续可以直接增加跨 rank + 差异或采样稳定性指标,不需要重新设计采集路径; +3. **隔离关键路径**:all-gather 在后台线程和专用 Gloo 通信组中运行,主推理线程只在 + `WAIT_LOAD_GATHER_FINISHED` 状态轮询完成标记,不会被大块负载通信直接阻塞。 + +planner 接口直接接收 `[rank, layer, sample, logical_expert]` CPU Tensor。当前 Greedy +实现进入算法主体前沿 rank 和 sample 维求和为 `[layer, logical_expert]`,再转换成嵌套 +list;因此原始维度在 planner 边界仍然可用,而后续贪心逻辑保持简单的纯 Python 实现。 + +## 8. 布局规划 + +### 8.1 规划器接口与选择 + +所有布局算法实现统一的 `EPLBPlanner.plan(logical_expert_load_samples, current_placement)` 接口,返回: + +```text +logical_expert_load_samples: CPU Tensor[rank, layer, sample, logical_expert] +[layer][rank][local physical slot] -> logical expert ID +``` + +`--eplb_plan_mode` 只负责选择布局算法。`create_eplb_planner` 将字符串模式转换成具体实例,使状态机不依赖某个算法类。当前唯一模式为 `greedy`。 + +规划只在 rank 0 的后台线程执行。完成后,目标布局通过控制通信组广播给所有 rank。相同输入必须产生确定结果,便于所有 rank 生成一致的传输计划。 + +### 8.2 Greedy 规划算法 + +默认算法按层独立规划,主要步骤如下: + +1. **选择全卡冗余专家**:选取负载最高的 `R` 个逻辑专家,在每个 rank 上各放置一份; +2. **确定额外副本数**:其余专家先各保留一个副本,再把剩余 `R` 个副本逐次分给当前 `load / replica_count` 最大的专家; +3. **平铺多副本专家**:使用循环 rank 游标,把同一专家的副本放到不同 rank; +4. **放置单副本专家**:按专家负载从高到低处理,每次放到当前估算负载最低且仍有空槽的 rank; +5. **复用当前布局**:先把候选 rank 行匹配到共同专家最多的当前 rank,再让共同专家尽量保留原物理槽位,以减少跨 rank 传输和 rank 内覆盖。 + +规划负载按 128 token 对齐,降低很小的计数波动对布局的影响。专家 ID 和 rank ID 用作稳定的平局规则,因此结果是确定性的。 + +## 9. 权重迁移与安全提交 + +### 9.1 传输计划 + +传输规划器逐层比较当前布局和目标布局,为每个变化的目标槽绑定一个确定的源槽。选择源槽时优先使用不会被覆盖的稳定副本;没有稳定副本时,循环使用当前已有副本,避免把读取集中在同一个 rank。 + +随后根据“目标槽是否仍是其他任务的源槽”建立覆盖依赖: + +- **安全任务**:目标槽不再承担待处理任务的源,可以先传输并提交; +- **依赖环**:所有目标槽同时也是源槽,必须先把整个环的权重读入 pinned memory,再统一覆盖; +- **rank 冲突拆批**:普通批次中每个 rank 最多参与一条任务,在限制 pinned memory 峰值的同时保留跨 rank 并行性。 + +每层单独生成批次,再按层顺序拼接。这样完成一层的提交后就能立即发布该层的新路由 metadata。 + +### 9.2 数据路径 + +远程专家的传输路径为: + +```text +源 GPU 权重行 + -> 源 rank pinned CPU row + -> Gloo point-to-point + -> 目标 rank pinned CPU row + -> 目标 GPU live 权重行 +``` + +如果源和目标属于同一个 rank,则只执行 GPU 到 pinned CPU 的本地暂存,不经过网络。一次专家传输会覆盖实际推理需要的全部张量,包括量化权重及其 scale、zero point 等配套状态。 + +控制面和权重传输分别使用独立的 Gloo process group,避免两类通信相互干扰。后台线程只负责把数据传入 pinned memory,不直接修改 live 权重。 + +### 9.3 提交边界 + +只有当所有 rank 都确认当前批次传输完成后,主推理线程才会在 overlap stream 上统一: + +1. 把目标 rank 的 pinned row 写入 live GPU 权重槽; +2. 更新 `current_placement` 和本地槽位的 logical expert ID; +3. 为发生变化的层重建 `logical_to_physical_map`; +4. 将新 metadata 异步复制到 GPU。 + +权重和路由 metadata 在同一条 stream 上更新,后续 forward 只能看到完整提交后的状态,不会观察到“新路由指向旧权重”或“旧路由指向新权重”的中间状态。 + +全部批次完成后,manager 发布目标布局、清空 prefill 路由采样及设备端 sample index、增加完成次数,并回到 `COLLECTING`。 + +## 10. 布局持久化 + +成功完成重排后,rank 0 会把最新完整布局写回 `--eplb_config_path`。配置内容包括: + +```json +{ + "version": 1, + "num_logical_experts": 8, + "world_size": 4, + "num_redundant_experts_per_rank": 2, + "layers": { + "3": [[0, 1, 2, 3], [2, 3, 4, 5], [4, 5, 6, 7], [6, 7, 0, 1]] + } +} +``` + +写入前会再次校验全部层。实现使用独占创建的 `.lock` 文件避免多个服务同时写同一路径,并在成功写入后清除读取缓存。保存失败只记录 warning,不会中断在线推理。 + +## 11. 指标与运行行为 + +rank 0 周期性上报 logical expert 路由分布: + +```text +lightllm_eplb_topk_expert_imbalance_ratio_p25 +lightllm_eplb_topk_expert_imbalance_ratio_p50 +lightllm_eplb_topk_expert_imbalance_ratio_p100 +``` + +manager 先对每个有效层计算 `max(expert_load) / mean(expert_load)`,过滤没有采样负载的层,再对所有层的比值排序并使用 nearest-rank 位置取值: + +- P25:较均衡的四分之一位置,可观察大多数浅层或稳定层的基线; +- P50:中位层,用于描述典型 MoE 层的路由倾斜程度; +- P100:最大值,即当前窗口中最不均衡的层。 + +这三个值都以 `1` 表示完全均衡。例如 P50 为 `1.8`,表示中位层最热 logical expert 的 token 数是该层专家平均值的 1.8 倍。 + +当样本量达到规划阈值且 planner 产生目标布局后,rank 0 固定选取原始四维快照的第 0 个 sample 行,并仅沿 rank 维汇总为 `[layer, logical_expert]` 负载后上报: + +```text +lightllm_prefill_ep_compute_critical_overhead_ratio_before_rebalance +lightllm_prefill_ep_compute_critical_overhead_ratio_after_rebalance +``` + +这两个指标参考 `eplb2` 的关键路径计算开销定义,但不维护独立的 compute counter、后台 monitor 线程和额外通信组。manager 假设同一 logical expert 的流量由 hash 均匀分配给全部 physical 副本,并按 128 token 对每个副本的估算负载向上对齐。`before_rebalance` 使用当前布局,`after_rebalance` 使用 planner 给出的目标布局;二者的输入负载完全相同,可以直接衡量预计的重排收益。不会沿 sample 维累加,因为不同 sample 行来自不同 prefill 批次,累加后并不对应任何一次真实计算。如果 planner 判断布局无需改变,两个值应相同。 + +每层先计算最繁忙 rank 相对平均 rank 的额外负载,最后跨层汇总: + +```text +overhead_ratio = sum(max_rank_load - mean_rank_load) / sum(mean_rank_load) +``` + +因此不同层的热点 rank 不会互相抵消。指标为 `0` 表示估算的 EP rank 负载完全均衡,`0.3` 表示最慢 rank 造成的关键路径计算量比理想均衡状态高约 30%。它是基于聚合 logical route 和均匀副本分发假设的布局质量估算值,不是 DeepEP 接收缓冲区的实测 compute load,也不表示每次 prefill 的瞬时开销。只有实际执行 placement 规划时这两个 gauge 才会更新;其余时间保留最近一次规划结果。 + +`--eplb_rebalance_count` 的行为如下: + +- `-1`:持续允许动态规划和重排; +- `0`:使用初始或配置文件布局,不进行动态重排; +- 正整数:只统计实际完成且发生布局变化的重排;样本不足和布局不变不计数。 + +达到次数上限后,manager 仍会周期性采集和上报负载指标,但不再执行全局负载汇总和布局规划。 + +## 12. 当前限制 + +启用 EPLB 时需要满足: + +- 同时设置 `--enable_ep_moe`; +- `world_size > 1`; +- 逻辑专家数可以被 world size 整除; +- 冗余专家数大于 0,且不能超过本 rank 之外可复制的逻辑专家数; +- 同一 EP 通信组使用一致的 EPLB 参数; +- 不能与 `--enable_prefill_cudagraph` 同时使用; +- 当前不支持 `--enable_rl` 组合; +- 当前不支持 SM100 GPU; +- 当前动态 EPLB 执行路径依赖 EP DeepGEMM MoE 实现。 + +这些限制会在参数校验或 EPLB 初始化阶段尽早失败,避免服务带着不一致的布局进入推理。 + +## 13. 扩展新的规划算法 + +新增 planner 时建议遵循以下步骤: + +1. 在 `eplb/placement/` 中实现 `EPLBPlanner`; +2. 接收逻辑专家负载和当前完整布局,返回同 shape 的合法目标布局; +3. 在 `placement/factory.py` 的 builder 表中注册新的 `plan_mode`; +4. 在 CLI 和 `StartArgs` 的 `eplb_plan_mode` choices 中加入新名称; +5. 更新中英文参数文档; +6. 增加布局合法性、确定性、热点负载和迁移量测试; +7. 分别评估 prefill、decode 和混合流量,不要假设一种算法对所有部署形态都最优。 + +规划器必须保持以下边界: + +- 不直接修改 GPU 权重或路由 metadata; +- 不执行分布式通信; +- 不改变每个 rank 的物理槽位数; +- 不遗漏逻辑专家,不在同一 rank 重复放置同一专家; +- 相同输入返回确定结果; +- 尽量复用当前 rank 和物理槽位,避免均衡收益被迁移成本抵消。 + +`EPLBManager`、传输规划器和提交逻辑只依赖抽象的完整布局,因此新增算法不需要修改状态机。 + +## 14. 测试覆盖 + +EPLB 单元测试主要位于 `unit_tests/common/fused_moe/test_eplb.py`,覆盖: + +- 初始布局、路由 metadata 和拓扑排序; +- planner 输入校验、布局合法性、负载均衡和槽位复用; +- 状态机各分支及后台任务轮询; +- 链式依赖、环形依赖和并发传输批次; +- 权重与 metadata 的安全提交; +- 配置文件加载、校验和持久化; +- logical-to-physical kernel 的精确映射; +- prefill 环形采样的逐行计数、循环覆盖、sample index 推进和同步状态复位; +- 非 prefill 路径不修改采样状态,以及空输入不推进 sample index; +- token index `0..4096`、expert ID `0..255`,以及 `2、3、4、5、101、127、128、251` 个副本时的哈希分布。 + +多 GPU pinned-memory 传输测试位于 `unit_tests/common/fused_moe/test_eplb_transfer_gpu.py`。 diff --git a/docs/CN/source/index.rst b/docs/CN/source/index.rst index 8f79e5126f..2ca6e962c7 100755 --- a/docs/CN/source/index.rst +++ b/docs/CN/source/index.rst @@ -80,6 +80,7 @@ Lightllm 整合了众多的开源方案的优点,包括但不限于 FasterTran 架构介绍 token attention介绍 峰值显存调度器介绍 + EPLB 专家负载均衡实现 .. Indices and tables .. ================== diff --git a/docs/CN/source/tutorial/api_server_args.rst b/docs/CN/source/tutorial/api_server_args.rst index eee60082ca..2c291f94d0 100644 --- a/docs/CN/source/tutorial/api_server_args.rst +++ b/docs/CN/source/tutorial/api_server_args.rst @@ -705,6 +705,66 @@ PD 分离模式参数 使用 tgi 输入和输出格式 +专家并行与 EPLB 参数 +-------------------- + +.. option:: --enable_ep_moe + + 为支持的 MoE 模型启用专家并行。使用 EPLB 时必须开启此参数。 + +.. option:: --eplb_num_redundant_experts_per_rank + + 每个 MoE 层在每个 EP rank 上分配的冗余物理专家数量,默认值为 ``0``,表示关闭 EPLB。 + 设置为正数时启用 EPLB,并且必须同时设置 ``--enable_ep_moe``;负数会在启动阶段被拒绝。 + + 每个 rank 会额外分配指定数量的专家权重行。EPLB 将逻辑专家映射到主副本或冗余物理副本, + 统计路由负载,并可在线迁移冗余副本以改善专家负载均衡。增大此值可以提供更多布局选择, + 但也会占用更多 GPU 显存并增加专家迁移流量。 + + EPLB 当前不能与 ``--enable_prefill_cudagraph`` 同时使用,也不支持 SM100 GPU。 + 同一部署中的所有 rank 和节点必须使用相同的配置值。 + +.. option:: --eplb_plan_mode + + EPLB 动态重排使用的专家布局规划算法,默认值为 ``greedy``。当前支持: + + * ``greedy``:根据各层逻辑专家的全局路由负载生成近似均衡的完整布局, + 并尽量复用当前 rank 和物理槽位以减少专家迁移。 + + 此参数只选择布局规划算法,不改变 token 到已有专家副本的运行时分发策略。 + 同一个 EP 通信组内的所有 rank 必须使用相同的值。PD 分离部署中的 prefill + 和 decode 进程拥有各自独立的 EPLB manager,因此可以分别设置适合各自流量 + 特征的规划算法;非 PD 部署则使用一个算法处理该进程采集到的全部路由负载。 + +.. option:: --eplb_rebalance_count + + 动态 EPLB 最多成功执行的重排次数,默认值为 ``1``。只有新布局实际发生 + 专家权重迁移并完成提交后才计数;样本不足或规划布局不变不会消耗次数。 + + * ``-1``:不限制次数,持续进行动态重排; + * ``0``:不进行动态重排,仅使用初始化时的冗余布局; + * 正整数:完成指定次数的重排后停止规划。 + +.. option:: --eplb_config_path + + EPLB 布局 JSON 文件路径,默认值为 ``None``。指定后,LightLLM 会在初始化专家权重之前校验并 + 读取各层保存的布局,使服务启动后立即使用上一次优化得到的专家分配。同一个文件也作为输出: + 每次成功完成动态重排后,rank 0 会写回当前最新布局。 + + 如果文件不存在、JSON 无法解析、缺少模型层,或保存的专家拓扑与当前部署不匹配,LightLLM 会 + 记录 warning,并对受影响的层使用默认初始化布局。只有后续动态重排成功完成时,rank 0 才会把 + 新布局写入该路径。 + + 以下示例为每个 EP rank 配置两个冗余专家:: + + python -m lightllm.server.api_server \ + --model_dir /path/to/model \ + --enable_ep_moe \ + --eplb_num_redundant_experts_per_rank 2 \ + --eplb_plan_mode greedy \ + --eplb_rebalance_count 1 \ + --eplb_config_path /path/to/eplb-placement.json + MTP 多预测参数 -------------- @@ -732,17 +792,6 @@ MTP 多预测参数 增加此值允许更多预测,但确保模型与指定的步数兼容。 目前 deepseekv3/r1 模型仅支持 1 步 -DeepSeek 冗余专家参数 ---------------------- - -.. option:: --ep_redundancy_expert_config_path - - 冗余专家配置的路径。可用于 deepseekv3 模型。 - -.. option:: --auto_update_redundancy_expert - - 是否通过在线专家使用计数器为 deepseekv3 模型更新冗余专家。 - 监控和日志参数 -------------- diff --git a/docs/EN/source/tutorial/api_server_args.rst b/docs/EN/source/tutorial/api_server_args.rst index e4c9151c0d..6f3a07fe1a 100644 --- a/docs/EN/source/tutorial/api_server_args.rst +++ b/docs/EN/source/tutorial/api_server_args.rst @@ -721,6 +721,71 @@ Sampling and Generation Parameters Use tgi input and output format +Expert Parallelism and EPLB Parameters +-------------------------------------- + +.. option:: --enable_ep_moe + + Enable expert parallelism for supported MoE models. EPLB requires this option. + +.. option:: --eplb_num_redundant_experts_per_rank + + Number of redundant physical experts allocated on each EP rank for every MoE layer. The default is ``0``, + which disables EPLB. A positive value enables EPLB and must be used together with ``--enable_ep_moe``; + negative values are rejected during startup. + + Each rank allocates the configured number of additional expert weight rows. EPLB maps logical experts to + primary or redundant physical copies, records routing load, and can migrate redundant copies online to + improve expert load balance. Larger values provide more placement flexibility but consume more GPU memory + and increase expert migration traffic. + + EPLB currently cannot be combined with ``--enable_prefill_cudagraph`` and is not supported on SM100 GPUs. + Use the same value on every rank and node in one deployment. + +.. option:: --eplb_plan_mode + + Expert placement planning algorithm used for dynamic EPLB rebalances. The default is ``greedy``. The + currently supported value is: + + * ``greedy``: builds an approximately balanced full placement from the global logical-expert load of each + layer and attempts to reuse the current ranks and physical slots to reduce expert migration. + + This option selects the placement planner; it does not change how tokens are dispatched among replicas in + an existing placement. Every rank in one EP communication group must use the same value. Prefill and decode + processes in a PD-disaggregated deployment have independent EPLB managers and may select planners suited to + their respective traffic. A non-PD process uses one planner for all routing load collected by that process. + +.. option:: --eplb_rebalance_count + + Maximum number of successfully completed dynamic EPLB rebalances. The default is ``1``. A count is consumed + only after a new placement has transferred and committed its expert weights; insufficient samples and unchanged + placements do not consume the limit. + + * ``-1`` keeps dynamic rebalancing enabled indefinitely. + * ``0`` disables dynamic rebalancing, leaving only the initial redundant placement active. + * A positive value stops planning after that many completed rebalances. + +.. option:: --eplb_config_path + + Path to an EPLB placement JSON file. The default is ``None``. When specified, LightLLM validates and loads + the saved per-layer placement before expert weights are initialized, so the service starts directly with the + previous optimized layout. The same file is updated with the latest layout after every successfully completed + rebalance. + + If the file does not exist, cannot be decoded, is missing a model layer, or does not match the current expert + topology, LightLLM logs a warning and uses the default initial placement for the affected layer. Rank 0 writes + a new layout to this path only after a dynamic rebalance completes successfully. + + Example: enable EPLB with two redundant experts per EP rank:: + + python -m lightllm.server.api_server \ + --model_dir /path/to/model \ + --enable_ep_moe \ + --eplb_num_redundant_experts_per_rank 2 \ + --eplb_plan_mode greedy \ + --eplb_rebalance_count 1 \ + --eplb_config_path /path/to/eplb-placement.json + MTP Multi-Prediction Parameters ------------------------------- @@ -748,17 +813,6 @@ MTP Multi-Prediction Parameters Increasing this value allows more predictions, but ensure the model is compatible with the specified number of steps. Currently deepseekv3/r1 models only support 1 step -DeepSeek Redundant Expert Parameters ------------------------------------- - -.. option:: --ep_redundancy_expert_config_path - - Path to redundant expert configuration. Can be used for deepseekv3 models. - -.. option:: --auto_update_redundancy_expert - - Whether to update redundant experts for deepseekv3 models through online expert usage counters. - Monitoring and Logging Parameters --------------------------------- diff --git a/lightllm/common/basemodel/layer_infer/cache_tensor_manager.py b/lightllm/common/basemodel/layer_infer/cache_tensor_manager.py index 8bcf99b992..13906d0d8a 100644 --- a/lightllm/common/basemodel/layer_infer/cache_tensor_manager.py +++ b/lightllm/common/basemodel/layer_infer/cache_tensor_manager.py @@ -22,7 +22,12 @@ def custom_del(self: torch.Tensor): if hasattr(self, "storage_weak_ptr"): storage_weak_ptr = self.storage_weak_ptr else: - storage_weak_ptr = self.untyped_storage()._weak_ref() + try: + storage_weak_ptr = self.untyped_storage()._weak_ref() + except RuntimeError: + # Some tensor implementations, including UndefinedTensorImpl, + # have no backing storage. Their destructor must stay silent. + return UntypedStorage._free_weak_ref(storage_weak_ptr) if storage_weak_ptr in g_cache_manager.ptr_to_bufnode: g_cache_manager.changed_ptr.add(storage_weak_ptr) diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/ep_redundancy.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/ep_redundancy.py deleted file mode 100644 index 749400c8d8..0000000000 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/ep_redundancy.py +++ /dev/null @@ -1,195 +0,0 @@ -import numpy as np -import torch -from .fused_moe_weight import FusedMoeWeight -from lightllm.utils.log_utils import init_logger -from typing import Dict - -logger = init_logger(__name__) - - -class FusedMoeWeightEPAutoRedundancy: - def __init__( - self, - ep_fused_moe_weight: FusedMoeWeight, - ) -> None: - super().__init__() - self._ep_w = ep_fused_moe_weight - self.redundancy_expert_num = self._ep_w.redundancy_expert_num - - def clear_counter(self): - self._ep_w.routed_expert_counter_tensor.fill_(0) - return - - def prepare_redundancy_experts( - self, - ): - expert_counter = self._ep_w.routed_expert_counter_tensor.detach().cpu().numpy() - logger.info( - f"layer_index {self._ep_w.layer_num_} global_rank {self._ep_w.global_rank_}" - f" expert_counter: {expert_counter}" - ) - self._ep_w.routed_expert_counter_tensor.fill_(0) - ep_n_routed_experts = self._ep_w.n_routed_experts // self._ep_w.global_world_size - start_expert_id = ep_n_routed_experts * self._ep_w.global_rank_ - no_redundancy_expert_ids = list(range(start_expert_id, start_expert_id + ep_n_routed_experts)) - - # 统计 0 rank 上的全局 topk 冗余信息,帮助导出一份全局可用的静态使用的冗余专家静态配置。 - if self._ep_w.global_rank_ == 0: - # int(e) for serialization, int64 can not be serialized by json.dump. - topk_redundancy_expert_ids = list(int(e) for e in np.argsort(expert_counter)[-self.redundancy_expert_num :]) - else: - topk_redundancy_expert_ids = None - - # 不要选中当前已经存在的非冗余专家作为冗余专家 - expert_counter[no_redundancy_expert_ids] = 0 - - self.redundancy_expert_ids = list(np.argsort(expert_counter)[-self.redundancy_expert_num :]) - logger.info( - f"layer_index {self._ep_w.layer_num_} global_rank {self._ep_w.global_rank_}" - f" new select redundancy_expert_ids : {self.redundancy_expert_ids}" - ) - - # 准备加载过度变量。 - self.experts_up_projs = [None] * self.redundancy_expert_num - self.experts_gate_projs = [None] * self.redundancy_expert_num - self.experts_up_proj_scales = [None] * self.redundancy_expert_num - self.experts_gate_proj_scales = [None] * self.redundancy_expert_num - self.w2_list = [None] * self.redundancy_expert_num - self.w2_scale_list = [None] * self.redundancy_expert_num - self.w13 = [None, None] # weight, weight_scale - self.w2 = [None, None] # weight, weight_scale - return topk_redundancy_expert_ids - - def load_hf_weights(self, weights): - # 加载冗余专家的权重参数 - for i, redundant_expert_id in enumerate(self.redundancy_expert_ids): - i_experts = redundant_expert_id - w1_weight = f"{self._ep_w.weight_prefix}.{i_experts}.{self._ep_w.w1_weight_name}.weight" - w2_weight = f"{self._ep_w.weight_prefix}.{i_experts}.{self._ep_w.w2_weight_name}.weight" - w3_weight = f"{self._ep_w.weight_prefix}.{i_experts}.{self._ep_w.w3_weight_name}.weight" - if w1_weight in weights: - self.experts_gate_projs[i] = weights[w1_weight] - if w3_weight in weights: - self.experts_up_projs[i] = weights[w3_weight] - if w2_weight in weights: - self.w2_list[i] = weights[w2_weight] - - self._load_weight_scale(weights) - self._fuse() - - def _fuse(self): - self._fuse_weight_scale() - - with self._ep_w.lock: - if ( - hasattr(self, "experts_up_projs") - and None not in self.experts_up_projs - and None not in self.experts_gate_projs - and None not in self.w2_list - ): - gate_out_dim, gate_in_dim = self.experts_gate_projs[0].shape - up_out_dim, up_in_dim = self.experts_up_projs[0].shape - assert gate_in_dim == up_in_dim - dtype = self.experts_gate_projs[0].dtype - total_expert_num = self.redundancy_expert_num - - w13 = torch.empty((total_expert_num, gate_out_dim + up_out_dim, gate_in_dim), dtype=dtype, device="cpu") - - for i_experts in range(self.redundancy_expert_num): - w13[i_experts, 0:gate_out_dim:, :] = self.experts_gate_projs[i_experts] - w13[i_experts, gate_out_dim:, :] = self.experts_up_projs[i_experts] - - inter_shape, hidden_size = self.w2_list[0].shape[0], self.w2_list[0].shape[1] - w2 = torch._utils._flatten_dense_tensors(self.w2_list).view(len(self.w2_list), inter_shape, hidden_size) - if self._ep_w.quant_method._check_weight_need_quanted(weight=w13): - w13_pack, _ = self._ep_w.quant_method.create_moe_weight( - out_dims=[gate_out_dim + up_out_dim], - in_dim=1, - dtype=self._ep_w.data_type_, - device_id=self._ep_w.device_id_, - num_experts=self.redundancy_expert_num, - ) - self._ep_w.quant_method.quantize(w13, w13_pack) - w2_pack, _ = self._ep_w.quant_method.create_moe_weight( - out_dims=[inter_shape], - in_dim=hidden_size, - dtype=self._ep_w.data_type_, - device_id=self._ep_w.device_id_, - num_experts=self.redundancy_expert_num, - ) - self._ep_w.quant_method.quantize(w2, w2_pack) - - self.w13[0] = w13_pack.weight - self.w13[1] = w13_pack.weight_scale - self.w2[0] = w2_pack.weight - self.w2[1] = w2_pack.weight_scale - else: - self.w13[0] = w13 - self.w2[0] = w2 - delattr(self, "w2_list") - delattr(self, "experts_up_projs") - delattr(self, "experts_gate_projs") - - def _fuse_weight_scale(self): - with self._ep_w.lock: - if ( - hasattr(self, "experts_up_proj_scales") - and None not in self.experts_up_proj_scales - and None not in self.experts_gate_proj_scales - and None not in self.w2_scale_list - ): - gate_out_dim, gate_in_dim = self.experts_gate_proj_scales[0].shape - up_out_dim, up_in_dim = self.experts_up_proj_scales[0].shape - assert gate_in_dim == up_in_dim - dtype = self.experts_gate_proj_scales[0].dtype - total_expert_num = self.redundancy_expert_num - w13_scale = torch.empty( - (total_expert_num, gate_out_dim + up_out_dim, gate_in_dim), dtype=dtype, device="cpu" - ) - for i_experts in range(self.redundancy_expert_num): - w13_scale[i_experts, 0:gate_out_dim:, :] = self.experts_gate_proj_scales[i_experts] - w13_scale[i_experts, gate_out_dim:, :] = self.experts_up_proj_scales[i_experts] - - inter_shape, hidden_size = self.w2_scale_list[0].shape[0], self.w2_scale_list[0].shape[1] - w2_scale = torch._utils._flatten_dense_tensors(self.w2_scale_list).view( - len(self.w2_scale_list), inter_shape, hidden_size - ) - self.w13[1] = w13_scale - self.w2[1] = w2_scale - delattr(self, "w2_scale_list") - delattr(self, "experts_up_proj_scales") - delattr(self, "experts_gate_proj_scales") - - def _load_weight_scale(self, weights: Dict[str, torch.Tensor]) -> None: - # 加载冗余专家的scale参数 - for i, redundant_expert_id in enumerate(self.redundancy_expert_ids): - i_experts = redundant_expert_id - weight_scale_suffix = self._ep_w.quant_method.weight_scale_suffix - w1_scale = f"{self._ep_w.weight_prefix}.{i_experts}.{self._ep_w.w1_weight_name}.{weight_scale_suffix}" - w2_scale = f"{self._ep_w.weight_prefix}.{i_experts}.{self._ep_w.w2_weight_name}.{weight_scale_suffix}" - w3_scale = f"{self._ep_w.weight_prefix}.{i_experts}.{self._ep_w.w3_weight_name}.{weight_scale_suffix}" - if w1_scale in weights: - self.experts_gate_proj_scales[i] = weights[w1_scale] - if w3_scale in weights: - self.experts_up_proj_scales[i] = weights[w3_scale] - if w2_scale in weights: - self.w2_scale_list[i] = weights[w2_scale] - - def commit(self): - for index, dest_tensor in enumerate([self._ep_w.w13.weight, self._ep_w.w13.weight_scale]): - if dest_tensor is not None: - assert isinstance( - dest_tensor, torch.Tensor - ), f"dest_tensor should be a torch.Tensor, but got {type(dest_tensor)}" - dest_tensor[-self.redundancy_expert_num :, :, :] = self.w13[index][:, :, :] - - for index, dest_tensor in enumerate([self._ep_w.w2.weight, self._ep_w.w2.weight_scale]): - if dest_tensor is not None: - assert isinstance( - dest_tensor, torch.Tensor - ), f"dest_tensor should be a torch.Tensor, but got {type(dest_tensor)}" - dest_tensor[-self.redundancy_expert_num :, :, :] = self.w2[index][:, :, :] - - self._ep_w.redundancy_expert_ids_tensor.copy_( - torch.tensor(self.redundancy_expert_ids, dtype=torch.int64, device="cpu") - ) diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py index 7f369c4fd8..82546acf2e 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py @@ -8,10 +8,10 @@ get_col_slice_mixin, SliceMixinTpl, ) -from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl import select_fuse_moe_impl +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl import create_fuse_moe_impl from lightllm.common.basemodel.moe_route_info_manager import get_moe_capture_callback from lightllm.common.quantization.quantize_method import QuantizationMethod -from lightllm.utils.envs_utils import get_redundancy_expert_ids, get_redundancy_expert_num, get_env_start_args +from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.dist_utils import get_global_world_size, get_global_rank from lightllm.utils.log_utils import init_logger @@ -56,18 +56,15 @@ def __init__( self.n_routed_experts = n_routed_experts self.num_fused_shared_experts = num_fused_shared_experts self._init_config(network_config) - self._init_redundancy_expert_params() - self._init_parallel_params() - self.fuse_moe_impl = select_fuse_moe_impl(self.quant_method, self.enable_ep_moe)( + self.fuse_moe_impl = create_fuse_moe_impl( n_routed_experts=self.n_routed_experts, num_fused_shared_experts=self.num_fused_shared_experts, routed_scaling_factor=self.routed_scaling_factor, quant_method=self.quant_method, - redundancy_expert_num=self.redundancy_expert_num, - redundancy_expert_ids_tensor=self.redundancy_expert_ids_tensor, - routed_expert_counter_tensor=self.routed_expert_counter_tensor, - auto_update_redundancy_expert=self.auto_update_redundancy_expert, + enable_ep_moe=self.enable_ep_moe, + layer_index=self.layer_num_, ) + self._init_weight_partition() self.lock = threading.Lock() self._create_weight() @@ -80,16 +77,7 @@ def _init_config(self, network_config: Dict[str, Any]): self.routed_scaling_factor = network_config.get("routed_scaling_factor", 1.0) self.scoring_func = network_config.get("scoring_func", "softmax") - def _init_redundancy_expert_params(self): - self.redundancy_expert_num = get_redundancy_expert_num() - self.redundancy_expert_ids = get_redundancy_expert_ids(self.layer_num_) - self.auto_update_redundancy_expert: bool = get_env_start_args().auto_update_redundancy_expert - self.redundancy_expert_ids_tensor = torch.tensor(self.redundancy_expert_ids, dtype=torch.int64, device="cuda") - self.routed_expert_counter_tensor = torch.zeros((self.n_routed_experts,), dtype=torch.int64, device="cuda") - # TODO: find out the reason of failure of deepep when redundancy_expert_num is 1. - assert self.redundancy_expert_num != 1, "redundancy_expert_num can not be 1 for some unknown hang of deepep." - - def _init_parallel_params(self): + def _init_weight_partition(self): if self.enable_ep_moe: self.tp_rank_ = 0 self.tp_world_size_ = 1 @@ -103,27 +91,14 @@ def _init_parallel_params(self): self.split_inter_size = self.moe_intermediate_size // self.tp_world_size_ if self.enable_ep_moe: assert self.num_fused_shared_experts == 0, "num_fused_shared_experts must be 0 when enable_ep_moe" + self.local_logic_expert_ids_list = self.fuse_moe_impl.local_logics_expert_ids_list logger.debug( f"global_rank {self.global_rank_} layerindex {self.layer_num_} " - f"redundancy_expertids: {self.redundancy_expert_ids}" - ) - self.local_n_routed_experts = self.n_routed_experts // self.global_world_size + self.redundancy_expert_num - n_experts_per_rank = self.n_routed_experts // self.global_world_size - start_expert_id = self.global_rank_ * n_experts_per_rank - self.local_expert_ids = ( - list(range(start_expert_id, start_expert_id + n_experts_per_rank)) + self.redundancy_expert_ids + f"local_logic_expert_ids_list: {self.local_logic_expert_ids_list}" ) - self.expert_idx_to_local_idx = { - expert_idx: expert_idx - start_expert_id for expert_idx in self.local_expert_ids[:n_experts_per_rank] - } - self.redundancy_expert_idx_to_local_idx = { - redundancy_expert_idx: n_experts_per_rank + i - for (i, redundancy_expert_idx) in enumerate(self.redundancy_expert_ids) - } + self.local_n_routed_experts = len(self.local_logic_expert_ids_list) else: - self.local_expert_ids = list(range(self.n_routed_experts + self.num_fused_shared_experts)) - self.expert_idx_to_local_idx = {expert_idx: i for (i, expert_idx) in enumerate(self.local_expert_ids)} - self.rexpert_idx_to_local_idx = {} + self.local_logic_expert_ids_list = list(range(self.n_routed_experts + self.num_fused_shared_experts)) def experts( self, @@ -134,10 +109,11 @@ def experts( use_grouped_topk: bool, topk_group: int, num_expert_group: int, - is_prefill: Optional[bool] = None, + is_prefill: bool, infer_state=None, shared_expert_gate: Optional[torch.Tensor] = None, ) -> torch.Tensor: + assert is_prefill is not None, "is_prefill must be explicitly specified for fused MoE execution" # Captures MoE topk expert ids for routed-experts metadata when enabled. moe_capture_callback = get_moe_capture_callback(infer_state, self.layer_num_) return self.fuse_moe_impl( @@ -280,9 +256,7 @@ def load_hf_weights(self, weights): # Load bias self._load_e_score_correction_bias(weights) self._load_per_expert_scale(weights) - self._load_weight(self.expert_idx_to_local_idx, weights) - if self.redundancy_expert_num > 0: - self._load_weight(self.redundancy_expert_idx_to_local_idx, weights) + self._load_weight(self.local_logic_expert_ids_list, weights) def verify_load(self): weight_load_ok = all(all(_weight_pack.load_ok) for _weight_pack in self.w1_list + self.w2_list + self.w3_list) @@ -351,8 +325,8 @@ def _get_expert_weight_list(self, weight_pack: WeightPack): weight_list.append(expert_weight) return weight_list - def _load_weight(self, expert_idx_to_local_idx: Dict[int, int], weights: Dict[str, torch.Tensor]): - for expert_idx, local_expert_idx in expert_idx_to_local_idx.items(): + def _load_weight(self, local_logic_expert_ids_list: List[int], weights: Dict[str, torch.Tensor]): + for local_expert_idx, expert_idx in enumerate(local_logic_expert_ids_list): with self.lock: self._load_expert(expert_idx, local_expert_idx, weights) self._load_expert_scale( diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/gpt_oss_fused_moe_weight_tp.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/gpt_oss_fused_moe_weight_tp.py index 240bc726ca..8bc37d74c7 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/gpt_oss_fused_moe_weight_tp.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/gpt_oss_fused_moe_weight_tp.py @@ -144,10 +144,11 @@ def experts( use_grouped_topk: bool, topk_group: int, num_expert_group: int, - is_prefill: Optional[bool] = None, + is_prefill: bool, infer_state=None, shared_expert_gate: Optional[torch.Tensor] = None, ): + assert is_prefill is not None, "is_prefill must be explicitly specified for fused MoE execution" assert shared_expert_gate is None, "shared_expert_gate is not supported by GPT-OSS fused MoE" topk_weights, topk_ids = self._router(router_logits, top_k) diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/__init__.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/__init__.py index 67bb90e4ef..80a320cefa 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/__init__.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/__init__.py @@ -1,14 +1,36 @@ +from typing import Optional + from lightllm.common.quantization.quantize_method import QuantizationMethod from .triton_impl import FuseMoeTriton from .marlin_impl import FuseMoeMarlin from .deepgemm_impl import FuseMoeDeepGEMM -def select_fuse_moe_impl(quant_method: QuantizationMethod, enable_ep_moe: bool): - if enable_ep_moe: - return FuseMoeDeepGEMM +def create_fuse_moe_impl( + *, + n_routed_experts: int, + num_fused_shared_experts: int, + routed_scaling_factor: float, + quant_method: QuantizationMethod, + enable_ep_moe: bool = False, + layer_index: Optional[int] = None, +): + """创建持有自身路由运行态的 MoE 执行实现。 - if quant_method.method_name == "awq_marlin": - return FuseMoeMarlin + 这里直接返回完成初始化的对象,而不是仅返回实现类,使 EPLB 布局、路由 + 计数器等后端专属状态与使用它们的 kernel 保持在同一个实现对象中。 + """ + if enable_ep_moe: + impl_cls = FuseMoeDeepGEMM + elif quant_method.method_name == "awq_marlin": + impl_cls = FuseMoeMarlin else: - return FuseMoeTriton + impl_cls = FuseMoeTriton + + return impl_cls( + n_routed_experts=n_routed_experts, + num_fused_shared_experts=num_fused_shared_experts, + routed_scaling_factor=routed_scaling_factor, + quant_method=quant_method, + layer_index=layer_index, + ) diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py index 1e3ad4b196..dc6660925c 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py @@ -1,53 +1,53 @@ import torch -from abc import abstractmethod -from typing import Callable, Optional +from abc import ABC, abstractmethod +from typing import Callable, Optional, Tuple from lightllm.common.quantization.quantize_method import ( WeightPack, QuantizationMethod, ) -from lightllm.utils.dist_utils import ( - get_global_rank, - get_global_world_size, -) -class FuseMoeBaseImpl: +class FuseMoeBaseImpl(ABC): + """将逻辑专家路由与实际执行布局分离。 + + 融合 MoE 的调用流程如下:: + + _select_experts + -> topk_weights + logical_topk_ids + -> moe_capture_callback(logical_topk_ids) + -> _prepare_expert_execution + -> 追加 shared expert,或者 + -> 将 EPLB logical ID 映射为 physical expert ID + -> _fused_experts(topk_weights, execution_topk_ids) + + ``_select_experts`` 对所有实现都遵循同一套稳定接口:只在模型的逻辑专家 + 空间中选择专家,并返回原始 logical ID。该阶段不能应用任何与实际执行相关 + 的布局转换,例如追加 shared expert 行,或者映射到 EPLB 冗余 physical 行。 + + capture callback 在任何布局转换之前执行,因此采集到的路由元数据始终描述 + 模型的 logical expert。所有实现都必须将 logical ID 视为只读数据。 + ``_prepare_expert_execution`` 可以为实际执行布局分配新的 ID tensor,但不能 + 原地修改 logical ID tensor。 + + 因此,``_fused_experts`` 接收的是 execution ID:普通路径中仍是未经修改的 + logical ID;融合 shared expert 路径中是追加了 shared expert 行的 ID;EPLB + 路径中则是完成映射后的 physical ID。 + """ + def __init__( self, n_routed_experts: int, num_fused_shared_experts: int, routed_scaling_factor: float, quant_method: QuantizationMethod, - redundancy_expert_num: int, - redundancy_expert_ids_tensor: torch.Tensor, - routed_expert_counter_tensor: torch.Tensor, - auto_update_redundancy_expert: bool, + layer_index: Optional[int] = None, ): self.n_routed_experts = n_routed_experts self.num_fused_shared_experts = num_fused_shared_experts self.routed_scaling_factor = routed_scaling_factor self.quant_method = quant_method - self.global_rank_ = get_global_rank() - self.global_world_size_ = get_global_world_size() - self.ep_n_routed_experts = self.n_routed_experts // self.global_world_size_ - self.total_expert_num_contain_redundancy = ( - self.n_routed_experts + redundancy_expert_num * self.global_world_size_ - ) - - # redundancy expert related - self.redundancy_expert_num = redundancy_expert_num - self.redundancy_expert_ids_tensor = redundancy_expert_ids_tensor - self.routed_expert_counter_tensor = routed_expert_counter_tensor - self.auto_update_redundancy_expert = auto_update_redundancy_expert - - # workspace for kernel optimization - self.workspace = self.create_workspace() - - @abstractmethod - def create_workspace(self): - pass + self.layer_index = layer_index - @abstractmethod def __call__( self, input_tensor: torch.Tensor, @@ -61,11 +61,113 @@ def __call__( use_grouped_topk: bool, topk_group: int, num_expert_group: int, - is_prefill: Optional[bool] = None, + is_prefill: bool, # Callback to capture MoE topk expert ids (routed experts metadata). moe_capture_callback: Optional[Callable[[torch.Tensor], None]] = None, per_expert_scale: Optional[torch.Tensor] = None, - # Qwen3.5 uses this gate to control fused shared expert aggregation weights. + # Qwen3Next/Qwen3.5-MoE 在 TP 模式下将 shared expert 融合进 routed MoE。 + # 该参数是 shared_expert_gate(hidden_states) 产生的逐 token 门控 logit; + # 追加 shared expert 时使用 sigmoid(logit) 作为其聚合权重。 + shared_expert_gate: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + assert is_prefill is not None, "is_prefill must be explicitly specified for fused MoE execution" + topk_weights, topk_ids = self._select_experts( + input_tensor=input_tensor, + router_logits=router_logits, + correction_bias=correction_bias, + top_k=top_k, + renormalize=renormalize, + use_grouped_topk=use_grouped_topk, + topk_group=topk_group, + num_expert_group=num_expert_group, + scoring_func=scoring_func, + per_expert_scale=per_expert_scale, + ) + if moe_capture_callback is not None: + moe_capture_callback(topk_ids) + topk_weights, topk_ids = self._prepare_expert_execution( + topk_weights=topk_weights, + topk_ids=topk_ids, + shared_expert_gate=shared_expert_gate, + is_prefill=is_prefill, + ) + return self._fused_experts( + input_tensor=input_tensor, + w13=w13, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + router_logits=router_logits, + is_prefill=is_prefill, + ) + + @abstractmethod + def _select_experts( + self, + input_tensor: torch.Tensor, + router_logits: torch.Tensor, + correction_bias: Optional[torch.Tensor], + top_k: int, + renormalize: bool, + use_grouped_topk: bool, + topk_group: int, + num_expert_group: int, + scoring_func: str, + per_expert_scale: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """在模型的逻辑专家空间中完成 top-k 路由选择。 + + 返回形状一致的 ``topk_weights`` 和 ``logical_topk_ids``。这里的 + ``topk_ids`` 必须始终表示模型配置中的原始 logical expert,不能追加 + shared expert,也不能映射到 EPLB physical expert。返回的 ID tensor 会 + 先交给 capture callback,后续实现必须将其视为只读数据。 + """ + pass + + @abstractmethod + def _prepare_expert_execution( + self, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + is_prefill: bool, shared_expert_gate: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """将逻辑路由结果转换为 MoE kernel 实际需要的执行布局。 + + 该方法在 capture callback 之后执行。每个实现都必须明确处理自己的执行 + 布局:普通路径保持权重和 logical ID 不变,shared-expert 路径追加对应 + expert,EPLB 路径将 logical ID 修复为 physical ID。发生布局转换时应 + 返回新的 ID tensor,不能原地修改传入的 logical ``topk_ids``;如果专家 + 数量发生变化,``topk_weights`` 必须同步调整。 + + ``shared_expert_gate`` 当前由 Qwen3Next/Qwen3.5-MoE 使用,其内容是 + ``shared_expert_gate(hidden_states)`` 计算出的逐 token 门控 logit。在 TP + fused shared-expert 路径中,``sigmoid(logit)`` 会作为追加 shared expert + 的权重,使其输出按 token 动态参与 routed expert 输出的聚合;传入 ``None`` + 时,普通 fused shared expert 的追加权重为 1。EP 路径不使用该参数,而是 + 单独计算 shared expert,并在应用相同门控后与 routed MoE 输出相加。 + + ``is_prefill`` 必须由上层入口显式传入 ``True`` 或 ``False``。即使当前实现 + 尚未使用该信息,也不允许用 ``None`` 隐式表示执行阶段。 + """ + pass + + @abstractmethod + def _fused_experts( + self, + input_tensor: torch.Tensor, + w13: WeightPack, + w2: WeightPack, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + is_prefill: bool, + router_logits: Optional[torch.Tensor] = None, ) -> torch.Tensor: + """根据准备完成的路由结果执行融合 MoE 计算。 + + 这里的 ``topk_ids`` 已处于实际执行所需的 ID 空间:普通路径为 logical + ID,shared expert 路径包含追加的专家 ID,EPLB 路径则为 physical ID。 + 实现只能读取路由 ID,不能原地修改其内容。``is_prefill`` 只用于选择底层 + 执行策略,不应再影响专家选择或 logical-to-physical 映射语义。 + """ pass diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py index 024be9f55c..82b7309967 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py @@ -1,12 +1,18 @@ import torch from typing import Optional, Tuple, Any -from .triton_impl import FuseMoeTriton +from .base_impl import FuseMoeBaseImpl from lightllm.distributed import dist_group_manager from lightllm.common.quantization.quantize_method import WeightPack from lightllm.utils.envs_utils import ( + get_env_start_args, get_deepep_num_max_dispatch_tokens_per_rank_prefill, get_deepep_num_max_dispatch_tokens_per_rank_decode, ) +from lightllm.utils.dist_utils import ( + get_global_rank, + get_global_world_size, + get_node_world_size, +) from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import ( fused_experts, get_ep_num_sms, @@ -15,11 +21,98 @@ quantize_fused_experts_input, ) from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd +from lightllm.common.basemodel.triton_kernel.fused_moe.eplb_topk_ids import ( + eplb_repair_topk_ids, +) from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType -from lightllm.common.basemodel.triton_kernel.redundancy_topk_ids_repair import redundancy_topk_ids_repair -class FuseMoeDeepGEMM(FuseMoeTriton): +class FuseMoeDeepGEMM(FuseMoeBaseImpl): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._init_eplb_runtime() + + def _init_eplb_runtime(self): + """初始化本地物理槽位以及可更新的 EPLB 路由运行态。 + + ``local_logics_expert_ids_list`` 始终描述全部本地物理行。初始化时主专家 + 在前、冗余专家在后;负载均衡运行后允许替换任意物理行,并在同一个 + 安全推理边界同时更新专家权重和 ``logical_to_physical_map``。 + """ + world_size = get_global_world_size() + assert self.n_routed_experts % world_size == 0 + global_rank = get_global_rank() + start_args = get_env_start_args() + self.num_redundant_experts_per_rank = start_args.eplb_num_redundant_experts_per_rank + + if self.num_redundant_experts_per_rank > 0: + # 延迟导入:顶层导入会经 mode_backend 包形成 meta_weights -> server 的循环依赖。 + from lightllm.server.router.model_infer.mode_backend.eplb.placement import ( + build_initial_local_expert_ids, + build_logical_to_physical_map, + load_layer_placement, + ) + + self.num_total_physical_experts = self.n_routed_experts + world_size * self.num_redundant_experts_per_rank + + # 阶段 1:先构造确定性的默认布局。未指定配置文件,或配置读取、校验失败时, + # 后续权重初始化会继续使用这份布局。 + initial_local_expert_ids_by_rank = build_initial_local_expert_ids( + self.n_routed_experts, + world_size, + self.num_redundant_experts_per_rank, + ) + + # 阶段 2:如果指定了配置文件,尝试读取与当前层及部署拓扑匹配的历史布局。 + # load_layer_placement 会负责记录 warning,并在任何异常或配置无效时返回 None。 + config_path = start_args.eplb_config_path + if config_path is not None: + saved_placement = load_layer_placement( + config_path, + layer_index=self.layer_index, + num_logical_experts=self.n_routed_experts, + world_size=world_size, + num_redundant_experts_per_rank=self.num_redundant_experts_per_rank, + ) + + # 阶段 3:只有完整校验通过的历史布局才会替换默认布局,使专家权重在 + # 初始化时直接加载到上一次优化后的物理槽位中。 + if saved_placement is not None: + initial_local_expert_ids_by_rank = saved_placement + self.local_logics_expert_ids_list = initial_local_expert_ids_by_rank[global_rank] + self.logical_to_physical_map = torch.tensor( + build_logical_to_physical_map( + initial_local_expert_ids_by_rank, + self.n_routed_experts, + current_rank=global_rank, + node_world_size=get_node_world_size(), + ), + dtype=torch.int32, + ).cuda() + # 环形缓冲区保留最近 24 次 prefill 路由采样,每次采样写入独立的一行; + # 始终按 logical expert 统计,冗余副本不会拆散规划器观察到的负载信号。 + self.prefill_route_counter = torch.zeros( + (24, self.n_routed_experts), + dtype=torch.int64, + device="cuda", + ) + # [0] 是单调递增的 sample index;[1] 用于在同一个 kernel 内协调 + # 目标行清零,并从所有 program 中选出最后完成者。 + self.prefill_route_sample_index = torch.zeros(2, dtype=torch.int64, device="cuda") + # 动态 EPLB 默认采集路由负载;以后使用配置文件固定专家布局时, + # 可以关闭该开关,避免执行不再需要的 atomic counter 更新。 + self.recording = True + else: + self.num_total_physical_experts = self.n_routed_experts + num_local_experts = self.n_routed_experts // world_size + first_local_expert_id = global_rank * num_local_experts + self.local_logics_expert_ids_list = list( + range( + first_local_expert_id, + first_local_expert_id + num_local_experts, + ) + ) + def _select_experts( self, input_tensor: torch.Tensor, @@ -32,10 +125,8 @@ def _select_experts( num_expert_group: int, scoring_func: str, per_expert_scale: Optional[torch.Tensor] = None, - shared_expert_gate: Optional[torch.Tensor] = None, ): - """Select experts and return topk weights and ids.""" - assert shared_expert_gate is None, "fused shared expert as MoE is not supported by DeepGEMM fused MoE" + """只选择逻辑专家,不在此阶段应用 EPLB 物理布局。""" from lightllm.common.basemodel.triton_kernel.fused_moe.topk_select import select_experts topk_weights, topk_ids = select_experts( @@ -53,19 +144,27 @@ def _select_experts( topk_weights.mul_(self.routed_scaling_factor) if per_expert_scale is not None: topk_weights = topk_weights * per_expert_scale[topk_ids.to(torch.long)].to(topk_weights.dtype) - origin_topk_ids = topk_ids - if self.redundancy_expert_num > 0: - # 因为 redundancy_topk_ids_repair 会修改 topk_ids,所以需要先复制一份 - origin_topk_ids = topk_ids.clone() - redundancy_topk_ids_repair( - topk_ids=topk_ids, - redundancy_expert_ids=self.redundancy_expert_ids_tensor, - ep_expert_num=self.ep_n_routed_experts, - global_rank=self.global_rank_, - expert_counter=self.routed_expert_counter_tensor, - enable_counter=self.auto_update_redundancy_expert, + return topk_weights, topk_ids + + def _prepare_expert_execution( + self, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + is_prefill: bool, + shared_expert_gate: Optional[torch.Tensor] = None, + ): + assert is_prefill is not None, "is_prefill must be explicitly specified for fused MoE execution" + assert shared_expert_gate is None, "fused shared expert as MoE is not supported by DeepGEMM fused MoE" + if self.num_redundant_experts_per_rank > 0: + topk_ids = eplb_repair_topk_ids( + logical_topk_ids=topk_ids, + logical_to_physical_map=self.logical_to_physical_map, + prefill_route_counter=self.prefill_route_counter, + prefill_route_sample_index=self.prefill_route_sample_index, + update_prefill_route_counter=self.recording and is_prefill is True, + mode="global_first", ) - return topk_weights, topk_ids, origin_topk_ids + return topk_weights, topk_ids def _fused_experts( self, @@ -74,8 +173,8 @@ def _fused_experts( w2: WeightPack, topk_weights: torch.Tensor, topk_ids: torch.Tensor, + is_prefill: bool, router_logits: Optional[torch.Tensor] = None, - is_prefill: Optional[bool] = None, ): output = fused_experts( hidden_states=input_tensor, @@ -83,7 +182,7 @@ def _fused_experts( w2=w2, topk_weights=topk_weights, topk_idx=topk_ids.to(torch.long), - num_experts=self.total_expert_num_contain_redundancy, # number of all experts contain redundancy + num_experts=self.num_total_physical_experts, quant_method=self.quant_method, is_prefill=is_prefill, previous_event=None, # for overlap @@ -102,7 +201,7 @@ def low_latency_dispatch( n_group: int, scoring_func: str, ): - topk_weights, topk_idx, _ = self._select_experts( + topk_weights, topk_idx = self._select_experts( input_tensor=hidden_states, router_logits=router_logits, correction_bias=e_score_correction_bias, @@ -113,6 +212,7 @@ def low_latency_dispatch( num_expert_group=n_group, scoring_func=scoring_func, ) + topk_weights, topk_idx = self._prepare_expert_execution(topk_weights, topk_idx, is_prefill=False) topk_idx = topk_idx.to(torch.long) num_max_dispatch_tokens_per_rank = get_deepep_num_max_dispatch_tokens_per_rank_decode() @@ -121,7 +221,7 @@ def low_latency_dispatch( topk_idx=topk_idx, x=hidden_states, num_max_dispatch_tokens_per_rank=num_max_dispatch_tokens_per_rank, - num_experts=self.total_expert_num_contain_redundancy, + num_experts=self.num_total_physical_experts, use_fp8=use_fp8_w8a8, async_finish=False, return_recv_hook=True, @@ -141,7 +241,7 @@ def select_experts_and_quant_input( n_group: int, scoring_func: str, ): - topk_weights, topk_idx, _ = self._select_experts( + topk_weights, topk_idx = self._select_experts( input_tensor=hidden_states, router_logits=router_logits, correction_bias=e_score_correction_bias, @@ -152,6 +252,7 @@ def select_experts_and_quant_input( num_expert_group=n_group, scoring_func=scoring_func, ) + topk_weights, topk_idx = self._prepare_expert_execution(topk_weights, topk_idx, is_prefill=True) qinput_tensor = quantize_fused_experts_input(hidden_states, w13, self.quant_method) return topk_weights, topk_idx.to(torch.long), qinput_tensor @@ -168,7 +269,7 @@ def dispatch( qinput_tensor, topk_idx=topk_idx, topk_weights=topk_weights, - num_experts=self.total_expert_num_contain_redundancy, + num_experts=self.num_total_physical_experts, num_max_tokens_per_rank=num_max_tokens_per_rank, expert_alignment=128, num_sms=get_ep_num_sms(), diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/marlin_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/marlin_impl.py index 0094b09b1c..c4937c681d 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/marlin_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/marlin_impl.py @@ -11,6 +11,10 @@ class FuseMoeMarlin(FuseMoeTriton): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.workspace = self.create_workspace() + def create_workspace(self): from lightllm.utils.vllm_utils import HAS_VLLM @@ -28,8 +32,8 @@ def _fused_experts( w2: WeightPack, topk_weights: torch.Tensor, topk_ids: torch.Tensor, + is_prefill: bool, router_logits: Optional[torch.Tensor] = None, - is_prefill: Optional[bool] = None, ): w1_weight, w1_scale, w1_zero_point = w13.weight, w13.weight_scale, w13.weight_zero_point diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py index 1d6a38c069..886217420e 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py @@ -1,36 +1,10 @@ import torch -from typing import Callable, Optional +from typing import Optional from lightllm.common.quantization.no_quant import WeightPack -from lightllm.common.quantization.quantize_method import QuantizationMethod from .base_impl import FuseMoeBaseImpl class FuseMoeTriton(FuseMoeBaseImpl): - def __init__( - self, - n_routed_experts: int, - num_fused_shared_experts: int, - routed_scaling_factor: float, - quant_method: QuantizationMethod, - redundancy_expert_num: int, - redundancy_expert_ids_tensor: torch.Tensor, - routed_expert_counter_tensor: torch.Tensor, - auto_update_redundancy_expert: bool, - ): - super().__init__( - n_routed_experts=n_routed_experts, - num_fused_shared_experts=num_fused_shared_experts, - routed_scaling_factor=routed_scaling_factor, - quant_method=quant_method, - redundancy_expert_num=redundancy_expert_num, - redundancy_expert_ids_tensor=redundancy_expert_ids_tensor, - routed_expert_counter_tensor=routed_expert_counter_tensor, - auto_update_redundancy_expert=auto_update_redundancy_expert, - ) - - def create_workspace(self): - return None - def _select_experts( self, input_tensor: torch.Tensor, @@ -43,7 +17,6 @@ def _select_experts( num_expert_group: int, scoring_func: str, per_expert_scale: Optional[torch.Tensor] = None, - shared_expert_gate: Optional[torch.Tensor] = None, ): """Select experts and return topk weights and ids.""" from lightllm.common.basemodel.triton_kernel.fused_moe.topk_select import select_experts @@ -63,7 +36,16 @@ def _select_experts( topk_weights.mul_(self.routed_scaling_factor) if per_expert_scale is not None: topk_weights = topk_weights * per_expert_scale[topk_ids.to(torch.long)].to(topk_weights.dtype) - origin_topk_ids = topk_ids + return topk_weights, topk_ids + + def _prepare_expert_execution( + self, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + is_prefill: bool, + shared_expert_gate: Optional[torch.Tensor] = None, + ): + assert is_prefill is not None, "is_prefill must be explicitly specified for fused MoE execution" if self.num_fused_shared_experts > 0: from lightllm.common.basemodel.triton_kernel.fused_moe.append_shared_expert_topk import ( append_fused_shared_experts, @@ -76,7 +58,7 @@ def _select_experts( num_fused_shared_experts=self.num_fused_shared_experts, shared_expert_gate=shared_expert_gate, ) - return topk_weights, topk_ids, origin_topk_ids + return topk_weights, topk_ids def _fused_experts( self, @@ -85,8 +67,8 @@ def _fused_experts( w2: WeightPack, topk_weights: torch.Tensor, topk_ids: torch.Tensor, + is_prefill: bool, router_logits: Optional[torch.Tensor] = None, - is_prefill: bool = False, ): w13_weight, w13_scale = w13.weight, w13.weight_scale w2_weight, w2_scale = w2.weight, w2.weight_scale @@ -106,50 +88,3 @@ def _fused_experts( w2_scale=w2_scale, ) return input_tensor - - def __call__( - self, - input_tensor: torch.Tensor, - router_logits: torch.Tensor, - w13: WeightPack, - w2: WeightPack, - correction_bias: Optional[torch.Tensor], - scoring_func: str, - top_k: int, - renormalize: bool, - use_grouped_topk: bool, - topk_group: int, - num_expert_group: int, - is_prefill: Optional[bool] = None, - # Callback to capture MoE topk expert ids (routed experts metadata). - moe_capture_callback: Optional[Callable[[torch.Tensor], None]] = None, - per_expert_scale: Optional[torch.Tensor] = None, - shared_expert_gate: Optional[torch.Tensor] = None, - ): - topk_weights, topk_ids, origin_topk_ids = self._select_experts( - input_tensor=input_tensor, - router_logits=router_logits, - correction_bias=correction_bias, - top_k=top_k, - renormalize=renormalize, - use_grouped_topk=use_grouped_topk, - topk_group=topk_group, - num_expert_group=num_expert_group, - scoring_func=scoring_func, - per_expert_scale=per_expert_scale, - shared_expert_gate=shared_expert_gate, - ) - - if moe_capture_callback is not None: - moe_capture_callback(origin_topk_ids) - - output = self._fused_experts( - input_tensor=input_tensor, - w13=w13, - w2=w2, - topk_weights=topk_weights, - topk_ids=topk_ids, - router_logits=router_logits, - is_prefill=is_prefill, - ) - return output diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/eplb_topk_ids.py b/lightllm/common/basemodel/triton_kernel/fused_moe/eplb_topk_ids.py new file mode 100644 index 0000000000..cd328e6eaa --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/eplb_topk_ids.py @@ -0,0 +1,307 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _replica_index(token_index, logical_expert_id, num_valid_replicas): + # 先用 logical expert ID 给 token index 加盐,再用 32-bit avalanche + # finalizer 打散规律性 token 间隔,避免低位周期与副本数产生相关性。 + value = token_index.to(tl.uint32) + value ^= (logical_expert_id.to(tl.uint32) + 1) * 0x9E3779B9 + value ^= value >> 16 + value *= 0x7FEB352D + value ^= value >> 15 + value *= 0x846CA68B + value ^= value >> 16 + value = value.to(tl.uint32) + return value % num_valid_replicas.to(tl.uint32) + + +@triton.jit +def _record_prefill_route_sample( + logical_expert_ids, + valid_mask, + prefill_route_counter_ptr, + prefill_route_counter_row_stride, + prefill_route_sample_index_ptr, + NUM_LOGICAL_EXPERTS: tl.constexpr, + PREFILL_ROUTE_COUNTER_CAPACITY: tl.constexpr, + COUNTER_BLOCK_SIZE: tl.constexpr, +): + """在主路由 kernel 内完成一次 prefill 采样事务。""" + # 同步槽 ``prefill_route_sample_index[1]`` 的状态机: + # + # 0 : 本次采样尚未清零; + # 1 : program 0 已清零,ready 标记已经发布; + # 1 + completed : ready 标记加上已经完成的 program 数; + # 1 + num_programs : 本次所有 program 均已完成。 + # + # 所有调用必须排在同一 CUDA stream 上,不能并发复用同一组 counter + # 和同步状态。 + program_id = tl.program_id(0) + + # 1. 所有 program 读取相同的 sample index,并通过取余定位本次写入的 + # 环形行。sample index 只有在全部 program 完成后才会推进,因此本次 + # kernel 生命周期内,各 program 看到的目标行保持不变。 + sample_index = tl.load(prefill_route_sample_index_ptr) + sample_row = sample_index % PREFILL_ROUTE_COUNTER_CAPACITY + + if program_id == 0: + # 2. program 0 负责初始化目标行。正常入口处同步槽必须为 0;清零 + # 完成后,通过 release 原子加一发布 ready=1,使此前的 store 对 + # 随后获得 ready 标记的其他 program 可见。 + sync_state = tl.atomic_add( + prefill_route_sample_index_ptr + 1, + 0, + sem="acquire", + scope="gpu", + ) + if sync_state == 0: + expert_offsets = tl.arange(0, COUNTER_BLOCK_SIZE) + tl.store( + prefill_route_counter_ptr + sample_row * prefill_route_counter_row_stride + expert_offsets, + 0, + mask=expert_offsets < NUM_LOGICAL_EXPERTS, + ) + tl.atomic_add( + prefill_route_sample_index_ptr + 1, + 1, + sem="release", + scope="gpu", + ) + else: + # 3. 其他 program 使用 acquire 原子读等待 ready 标记。只有观察到 + # sync >= 1 后才能离开循环,从而保证不会与 program 0 的清零 store + # 并发访问同一个 counter 行。 + sync_state = tl.atomic_add( + prefill_route_sample_index_ptr + 1, + 0, + sem="acquire", + scope="gpu", + ) + while sync_state < 1: + sync_state = tl.atomic_add( + prefill_route_sample_index_ptr + 1, + 0, + sem="acquire", + scope="gpu", + ) + + # 4. 清零屏障通过后,各 program 将自己的 logical expert 路由结果原子 + # 累加到同一采样行。这里统计 logical expert,冗余 physical 副本不会 + # 拆散规划器观察到的负载信号。 + tl.atomic_add( + prefill_route_counter_ptr + sample_row * prefill_route_counter_row_stride + logical_expert_ids, + 1, + mask=valid_mask, + sem="relaxed", + ) + + # 5. 本 program 完成计数后,向同步槽提交一个完成信号。调用发生在主 + # kernel 的 physical ID 写回之后,因此该信号同时表示两部分工作均完成。 + # atomic_add 返回旧值,故完成后的新值需要显式加一。当新值等于 + # ``num_programs + 1`` 时,ready 标记和全部 program 的完成信号均已到达。 + completed_programs = tl.atomic_add( + prefill_route_sample_index_ptr + 1, + 1, + sem="acq_rel", + scope="gpu", + ) + sync_after_completion = completed_programs + 1 + is_last_program = sync_after_completion == tl.num_programs(0) + 1 + if is_last_program: + # 最后完成者提交本次事务:先推进 sample index,使下一次采样指向 + # 后续环形行;再把同步槽复位为 0,供下一次 program 0 执行清零。 + tl.atomic_add( + prefill_route_sample_index_ptr, + 1, + sem="release", + scope="gpu", + ) + tl.atomic_xchg( + prefill_route_sample_index_ptr + 1, + 0, + sem="release", + scope="gpu", + ) + + +@triton.jit +def _eplb_repair_topk_ids_kernel( + logical_topk_ids_ptr, + physical_topk_ids_ptr, + num_topk_ids, + top_k, + logical_to_physical_map_ptr, + logical_to_physical_map_row_stride, + prefill_route_counter_ptr, + prefill_route_counter_row_stride, + prefill_route_sample_index_ptr, + DISPATCH_MODE: tl.constexpr, + UPDATE_PREFILL_ROUTE_COUNTER: tl.constexpr, + NUM_LOGICAL_EXPERTS: tl.constexpr, + PREFILL_ROUTE_COUNTER_CAPACITY: tl.constexpr, + COUNTER_BLOCK_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + # 阶段 1:将二维 [num_tokens, top_k] 路由结果展平后分块处理。 + # topk_id_offsets 同时用于访问输入、输出,并可恢复它所属的 token 下标。 + program_id = tl.program_id(0) + topk_id_offsets = program_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + valid_mask = topk_id_offsets < num_topk_ids + logical_expert_ids = tl.load(logical_topk_ids_ptr + topk_id_offsets, mask=valid_mask, other=0) + + # 阶段 2:定位每个 logical expert 的打包映射行。固定头部的布局为: + # + # [0] 所有 rank 上的有效副本总数 + # [1] 当前节点上的有效副本数,包含本卡副本 + # [2] 当前 GPU 上的有效副本数 + # [3:] 按本卡、本节点其他卡、其他节点排列的 physical expert IDs + # + # 因为三层候选在 physical ID 列表中都是连续前缀,所以选择对应层级的 + # count 后,可以直接对 [3:] 的前 count 项做 hash。 + map_row_offsets = logical_expert_ids * logical_to_physical_map_row_stride + num_global_replicas = tl.load(logical_to_physical_map_ptr + map_row_offsets, mask=valid_mask, other=1) + num_node_replicas = tl.load(logical_to_physical_map_ptr + map_row_offsets + 1, mask=valid_mask, other=0) + num_current_gpu_replicas = tl.load( + logical_to_physical_map_ptr + map_row_offsets + 2, + mask=valid_mask, + other=0, + ) + + # 阶段 3:根据调用方显式指定的分发模式选择参与 hash 的候选前缀。 + if DISPATCH_MODE == 0: + # current_gpu_first: 本卡 -> 全局。 + # TODO: 等 EPLB 布局算法支持节点拓扑感知后,再考虑增加 + # 本卡 -> 本节点 -> 全局的分层回退行为。 + num_preferred_replicas = tl.where( + num_current_gpu_replicas > 0, + num_current_gpu_replicas, + num_global_replicas, + ) + elif DISPATCH_MODE == 1: + # current_node_first: 本节点 -> 全局;不单独优先本卡。 + num_preferred_replicas = tl.where( + num_node_replicas > 0, + num_node_replicas, + num_global_replicas, + ) + else: + # global_first: 直接在所有 rank 的有效副本间分发。 + num_preferred_replicas = num_global_replicas + token_indices = topk_id_offsets // top_k + selected_replica_indices = _replica_index(token_indices, logical_expert_ids, num_preferred_replicas) + + # 阶段 4:读取选中槽位的 physical expert ID 并写入新的输出 tensor。 + # logical_topk_ids 只读,后续 callback 仍可安全观察原始逻辑路由结果。 + physical_expert_ids = tl.load( + logical_to_physical_map_ptr + map_row_offsets + 3 + selected_replica_indices, + mask=valid_mask, + other=-1, + ) + tl.store(physical_topk_ids_ptr + topk_id_offsets, physical_expert_ids, mask=valid_mask) + + # 阶段 5:记录本次 prefill 路由采样。子函数负责目标行清零、program 间 + # ready 同步、logical expert 计数,以及最后完成者对 sample index 的提交。 + if UPDATE_PREFILL_ROUTE_COUNTER: + _record_prefill_route_sample( + logical_expert_ids, + valid_mask, + prefill_route_counter_ptr, + prefill_route_counter_row_stride, + prefill_route_sample_index_ptr, + NUM_LOGICAL_EXPERTS, + PREFILL_ROUTE_COUNTER_CAPACITY, + COUNTER_BLOCK_SIZE, + ) + + +@torch.no_grad() +def eplb_repair_topk_ids( + logical_topk_ids: torch.Tensor, + logical_to_physical_map: torch.Tensor, + prefill_route_counter: torch.Tensor, + prefill_route_sample_index: torch.Tensor, + update_prefill_route_counter: bool, + mode: str, +) -> torch.Tensor: + """将 logical top-k ID 转换为当前 EPLB 布局中的 physical expert ID。 + + 参数: + logical_topk_ids: logical expert ID,shape 为 ``[num_tokens, top_k]``。 + logical_to_physical_map: 打包路由表,shape 为 + ``[num_logical_experts, 3 + routing_slots]``,单行布局为: + + ``[global_count, node_count, current_gpu_count, physical_ids..., padding...]`` + + 三个计数依次表示全局、当前节点和当前 GPU 上的有效副本数量。 + physical IDs 按本卡、本节点其他卡、其他节点排列;具体参与分发的 + 候选前缀由 ``mode`` 决定。 + 有效副本之后未使用的 padding 槽位为 -1,kernel 不会读取它们。 + prefill_route_counter: prefill 路由采样的环形缓冲区,shape 为 + ``[sample_capacity, num_logical_experts]``。一次 kernel 调用只写 + ``prefill_route_sample_index[0] % sample_capacity`` 对应的一行。 + prefill_route_sample_index: shape 为 ``[2]`` 的设备端同步状态。第 0 项 + 是单调递增的 sample index;第 1 项用于核内清零与完成同步:0 + 表示尚未清零,正数为 ready 标记 1 加上已完成的 program 数量; + 达到 ``num_programs + 1`` 后,最后完成者推进 sample index 并清零。 + update_prefill_route_counter: 是否记录本次 prefill 路由采样并推进 sample + index。decode 或固定布局不需要采样时可以关闭。 + mode: 必须显式指定的副本分发模式,不提供默认值: + + * ``current_gpu_first``:本卡优先,没有本卡副本时回退到全局; + * ``current_node_first``:本节点优先,没有节点内副本时回退到全局; + * ``global_first``:直接在全局全部有效副本间分发。 + + 返回: + physical expert ID,shape 为 ``[num_tokens, top_k]``。 + """ + dispatch_mode_ids = { + "current_gpu_first": 0, + "current_node_first": 1, + "global_first": 2, + } + assert ( + mode in dispatch_mode_ids + ), f"unsupported EPLB dispatch mode {mode!r}; expected one of {tuple(dispatch_mode_ids)}" + dispatch_mode = dispatch_mode_ids[mode] + + assert logical_topk_ids.is_contiguous() + assert logical_topk_ids.ndim == 2 + assert logical_to_physical_map.ndim == 2 + assert logical_to_physical_map.shape[1] > 3 + assert logical_to_physical_map.stride(1) == 1 + assert prefill_route_counter.ndim == 2 + assert prefill_route_counter.shape[0] > 1 + assert prefill_route_counter.shape[1] == logical_to_physical_map.shape[0] + assert prefill_route_counter.is_contiguous() + assert prefill_route_counter.dtype is torch.int64 + assert prefill_route_sample_index.shape == (2,) + assert prefill_route_sample_index.dtype is torch.int64 + assert prefill_route_sample_index.device == prefill_route_counter.device + physical_topk_ids = torch.empty_like(logical_topk_ids) + if logical_topk_ids.numel() == 0: + return physical_topk_ids + + block_size = 512 + _eplb_repair_topk_ids_kernel[(triton.cdiv(logical_topk_ids.numel(), block_size),)]( + logical_topk_ids_ptr=logical_topk_ids, + physical_topk_ids_ptr=physical_topk_ids, + num_topk_ids=logical_topk_ids.numel(), + top_k=logical_topk_ids.shape[1], + logical_to_physical_map_ptr=logical_to_physical_map, + logical_to_physical_map_row_stride=logical_to_physical_map.stride(0), + prefill_route_counter_ptr=prefill_route_counter, + prefill_route_counter_row_stride=prefill_route_counter.stride(0), + prefill_route_sample_index_ptr=prefill_route_sample_index, + DISPATCH_MODE=dispatch_mode, + UPDATE_PREFILL_ROUTE_COUNTER=update_prefill_route_counter, + NUM_LOGICAL_EXPERTS=prefill_route_counter.shape[1], + PREFILL_ROUTE_COUNTER_CAPACITY=prefill_route_counter.shape[0], + COUNTER_BLOCK_SIZE=triton.next_power_of_2(prefill_route_counter.shape[1]), + BLOCK_SIZE=block_size, + num_warps=4, + num_stages=1, + ) + return physical_topk_ids diff --git a/lightllm/common/basemodel/triton_kernel/redundancy_topk_ids_repair.py b/lightllm/common/basemodel/triton_kernel/redundancy_topk_ids_repair.py deleted file mode 100644 index ba48f414db..0000000000 --- a/lightllm/common/basemodel/triton_kernel/redundancy_topk_ids_repair.py +++ /dev/null @@ -1,111 +0,0 @@ -import torch -import triton -import triton.language as tl - - -@triton.jit -def _redundancy_topk_ids_repair_kernel( - topk_ids_ptr, - topk_total_num, - ep_expert_num, - redundancy_expert_num, - global_rank, - redundancy_expert_ids_ptr, - expert_counter_ptr, - BLOCK_SIZE: tl.constexpr, - ENABLE_COUNTER: tl.constexpr, -): - block_index = tl.program_id(0) - offs_d = block_index * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = offs_d < topk_total_num - current_topk_ids = tl.load(topk_ids_ptr + offs_d, mask=mask, other=0) - - if ENABLE_COUNTER: - tl.atomic_add(expert_counter_ptr + current_topk_ids, 1, mask=mask) - - # Remap original expert IDs to a new space that accounts for redundant expert slots. - new_current_topk_ids = (current_topk_ids // ep_expert_num) * redundancy_expert_num + current_topk_ids - - for i in tl.range(0, redundancy_expert_num, step=1, num_stages=3): - cur_redundancy_expert_id = tl.load(redundancy_expert_ids_ptr + i) - cur_redundancy_expert_id = ( - cur_redundancy_expert_id // ep_expert_num - ) * redundancy_expert_num + cur_redundancy_expert_id - new_current_topk_ids = tl.where( - new_current_topk_ids == cur_redundancy_expert_id, - (ep_expert_num + redundancy_expert_num) * (global_rank) + ep_expert_num + i, - new_current_topk_ids, - ) - - tl.store(topk_ids_ptr + offs_d, new_current_topk_ids, mask=mask) - return - - -@torch.no_grad() -def redundancy_topk_ids_repair( - topk_ids: torch.Tensor, - redundancy_expert_ids: torch.Tensor, - ep_expert_num: int, - global_rank: int, - expert_counter: torch.Tensor = None, - enable_counter: bool = False, -): - assert topk_ids.is_contiguous() - assert len(topk_ids.shape) == 2 - assert redundancy_expert_ids is not None - redundancy_expert_num = redundancy_expert_ids.shape[0] - BLOCK_SIZE = 512 - grid = (triton.cdiv(topk_ids.numel(), BLOCK_SIZE),) - num_warps = 4 - - _redundancy_topk_ids_repair_kernel[grid]( - topk_ids_ptr=topk_ids, - topk_total_num=topk_ids.numel(), - ep_expert_num=ep_expert_num, - redundancy_expert_num=redundancy_expert_num, - global_rank=global_rank, - redundancy_expert_ids_ptr=redundancy_expert_ids, - expert_counter_ptr=expert_counter, - BLOCK_SIZE=BLOCK_SIZE, - ENABLE_COUNTER=enable_counter, - num_warps=num_warps, - num_stages=3, - ) - return - - -@triton.jit -def _expert_id_counter_kernel( - topk_ids_ptr, - topk_total_num, - expert_counter_ptr, - BLOCK_SIZE: tl.constexpr, -): - block_index = tl.program_id(0) - offs_d = block_index * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = offs_d < topk_total_num - current_topk_ids = tl.load(topk_ids_ptr + offs_d, mask=mask, other=0) - tl.atomic_add(expert_counter_ptr + current_topk_ids, 1, mask=mask) - return - - -@torch.no_grad() -def expert_id_counter( - topk_ids: torch.Tensor, - expert_counter: torch.Tensor, -): - assert topk_ids.is_contiguous() - assert len(topk_ids.shape) == 2 - BLOCK_SIZE = 512 - grid = (triton.cdiv(topk_ids.numel(), BLOCK_SIZE),) - num_warps = 4 - - _expert_id_counter_kernel[grid]( - topk_ids_ptr=topk_ids, - topk_total_num=topk_ids.numel(), - expert_counter_ptr=expert_counter, - BLOCK_SIZE=BLOCK_SIZE, - num_warps=num_warps, - num_stages=1, - ) - return diff --git a/lightllm/distributed/communication_op.py b/lightllm/distributed/communication_op.py index 93c603212d..375df2d6f2 100644 --- a/lightllm/distributed/communication_op.py +++ b/lightllm/distributed/communication_op.py @@ -30,7 +30,6 @@ get_env_start_args, get_deepep_num_max_dispatch_tokens_per_rank_prefill, get_deepep_num_max_dispatch_tokens_per_rank_decode, - get_redundancy_expert_num, ) from lightllm.utils.dist_utils import ( get_global_world_size, @@ -193,7 +192,8 @@ def new_deepep_group( self.ll_num_tokens = prefill_num_max_dispatch_tokens_per_rank self.ll_decode_num_tokens = decode_num_max_dispatch_tokens_per_rank self.ll_hidden = hidden_size - self.ll_num_experts = n_routed_experts + get_redundancy_expert_num() * global_world_size + total_redundant_experts = get_env_start_args().eplb_num_redundant_experts_per_rank * global_world_size + self.ll_num_experts = n_routed_experts + total_redundant_experts self.ep_buffer = deep_ep.ElasticBuffer( deepep_group, num_max_tokens_per_rank=self.ll_num_tokens, @@ -274,9 +274,11 @@ def new_deepep_group( moe_intermediate_size, ) logger.info( - "Initialize DeepEP MoE buffers: low_latency=%s, mega_moe=%s, expert_quant_method_names=%s", + "Initialize DeepEP MoE buffers: low_latency=%s, mega_moe=%s, " + "ll_num_experts=%s, expert_quant_method_names=%s", enable_low_latency_buffer, enable_mega_moe_buffer, + self.ll_num_experts, sorted(expert_quant_method_names), ) theoretical_sms = self.ep_buffer.get_theoretical_num_sms(self.ll_num_experts, num_experts_per_tok) diff --git a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py index 3254031056..92ad86d21f 100644 --- a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py @@ -234,6 +234,7 @@ def _moe_ffn_tp( use_grouped_topk=self.n_group, topk_group=self.topk_group, num_expert_group=self.n_group, + is_prefill=infer_state.is_prefill, infer_state=infer_state, ) diff --git a/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py b/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py index 490d2dc4c5..1861b61848 100644 --- a/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py @@ -52,6 +52,7 @@ def _ffn(self, input, infer_state, layer_weight: GptOssTransformerLayerWeight) - use_grouped_topk=False, topk_group=None, num_expert_group=None, + is_prefill=infer_state.is_prefill, infer_state=infer_state, ) hidden_states = hidden_states.view(num_tokens, hidden_dim) diff --git a/lightllm/models/mixtral/layer_infer/transformer_layer_infer.py b/lightllm/models/mixtral/layer_infer/transformer_layer_infer.py index 8134dc266d..4e00659e80 100644 --- a/lightllm/models/mixtral/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/mixtral/layer_infer/transformer_layer_infer.py @@ -26,6 +26,7 @@ def _ffn(self, input, infer_state: InferStateInfo, layer_weight: MixtralTransfor use_grouped_topk=False, topk_group=None, num_expert_group=None, + is_prefill=infer_state.is_prefill, infer_state=infer_state, ) return hidden_states.view(num_tokens, hidden_dim) diff --git a/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py index 7311c4d141..d33367ab28 100644 --- a/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py @@ -88,6 +88,7 @@ def _moe_ffn_tp( use_grouped_topk=False, topk_group=None, num_expert_group=None, + is_prefill=infer_state.is_prefill, infer_state=infer_state, ) return hidden_states.view(num_tokens, hidden_dim) diff --git a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py index 92d68c9fd2..1249b8ff34 100644 --- a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py @@ -95,6 +95,7 @@ def _moe_ffn_tp( use_grouped_topk=False, topk_group=None, num_expert_group=None, + is_prefill=infer_state.is_prefill, infer_state=infer_state, shared_expert_gate=shared_expert_gate, ) diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index e01789d946..646cb4d5b3 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -767,15 +767,33 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: help="""Whether to enable ep moe for deepseekv3 model.""", ) parser.add_argument( - "--ep_redundancy_expert_config_path", + "--eplb_num_redundant_experts_per_rank", + type=int, + default=0, + help="""Number of redundant physical experts per EP rank for each MoE layer. + Set to 0 to disable EPLB.""", + ) + parser.add_argument( + "--eplb_plan_mode", type=str, - default=None, - help="""Path of the redundant expert config. It can be used for deepseekv3 model.""", + choices=["greedy"], + default="greedy", + help="""EPLB placement planning algorithm used by this inference process. + Prefill and decode processes may select their planner independently in PD deployments.""", ) parser.add_argument( - "--auto_update_redundancy_expert", - action="store_true", - help="""Whether to update the redundant expert for deepseekv3 model by online expert used counter.""", + "--eplb_rebalance_count", + type=int, + default=1, + help="""Maximum number of completed EPLB rebalances. -1 means unlimited, + 0 disables dynamic rebalancing, and the default is 1.""", + ) + parser.add_argument( + "--eplb_config_path", + type=str, + default=None, + help="""Path to an EPLB placement JSON file. A valid saved layout is loaded during weight + initialization, and the latest runtime layout is written back to the same path.""", ) parser.add_argument( "--enable_fused_shared_experts", diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index 1c433ec60e..7c72a7cbbc 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -158,6 +158,16 @@ def _launch_subprocesses(args: StartArgs): if args.enable_dp_prefill_balance: assert args.enable_tpsp_mix_mode and args.dp > 1, "need set --enable_tpsp_mix_mode firstly and --dp > 1" + assert ( + args.eplb_num_redundant_experts_per_rank >= 0 + ), "--eplb_num_redundant_experts_per_rank must be greater than or equal to 0" + assert args.eplb_rebalance_count >= -1, "--eplb_rebalance_count must be greater than or equal to -1" + if args.eplb_num_redundant_experts_per_rank > 0: + assert args.enable_ep_moe, "EPLB requires --enable_ep_moe" + assert not args.enable_prefill_cudagraph, "EPLB does not support --enable_prefill_cudagraph" + # TODO: Support EPLB redundant experts together with RL after their runtime state updates are coordinated. + assert not args.enable_rl, "EPLB redundant experts do not support --enable_rl" + if args.enable_ep_moe: allowed_ep_prefill_att_backends = {"auto", "fa3", "triton", "flashqla"} for backend in args.llm_prefill_att_backend: diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 254099d5bc..a3ff8061df 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -186,8 +186,10 @@ class StartArgs: default="gpu_counter", metadata={"choices": ["cpu_counter", "pin_mem_counter", "gpu_counter"]} ) enable_ep_moe: bool = field(default=False) - ep_redundancy_expert_config_path: Optional[str] = field(default=None) - auto_update_redundancy_expert: bool = field(default=False) + eplb_num_redundant_experts_per_rank: int = field(default=0) + eplb_plan_mode: str = field(default="greedy", metadata={"choices": ["greedy"]}) + eplb_rebalance_count: int = field(default=1) + eplb_config_path: Optional[str] = field(default=None) enable_fused_shared_experts: bool = field(default=False) mtp_mode: Optional[str] = field( default=None, diff --git a/lightllm/server/metrics/metrics.py b/lightllm/server/metrics/metrics.py index 0d42462c3f..a0db50c92c 100644 --- a/lightllm/server/metrics/metrics.py +++ b/lightllm/server/metrics/metrics.py @@ -32,6 +32,23 @@ "lightllm_cache_hit_rate": "Prefix cache hit rate of latest completed request", "lightllm_gen_throughput": "Generation throughput of latest completed request (tokens/s)", "lightllm_num_running_reqs": "Number of running requests", + "lightllm_prefill_ep_compute_critical_overhead_ratio_before_rebalance": ( + "Estimated excess critical EP-rank compute divided by balanced compute for the global load used by the latest " + "EPLB placement plan, evaluated before rebalance; 0.3 means 30% overhead" + ), + "lightllm_prefill_ep_compute_critical_overhead_ratio_after_rebalance": ( + "Estimated excess critical EP-rank compute divided by balanced compute for the global load used by the latest " + "EPLB placement plan, evaluated on the planned placement; 0.3 means 30% overhead" + ), + "lightllm_eplb_topk_expert_imbalance_ratio_p25": ( + "P25 across MoE layers of maximum-to-mean logical expert routed-token load" + ), + "lightllm_eplb_topk_expert_imbalance_ratio_p50": ( + "P50 across MoE layers of maximum-to-mean logical expert routed-token load" + ), + "lightllm_eplb_topk_expert_imbalance_ratio_p100": ( + "P100 across MoE layers of maximum-to-mean logical expert routed-token load" + ), } @@ -111,6 +128,11 @@ def init_metrics(self, args): self.create_gauge("lightllm_cache_hit_rate") self.create_gauge("lightllm_gen_throughput") self.create_gauge("lightllm_num_running_reqs") + self.create_gauge("lightllm_prefill_ep_compute_critical_overhead_ratio_before_rebalance") + self.create_gauge("lightllm_prefill_ep_compute_critical_overhead_ratio_after_rebalance") + self.create_gauge("lightllm_eplb_topk_expert_imbalance_ratio_p25") + self.create_gauge("lightllm_eplb_topk_expert_imbalance_ratio_p50") + self.create_gauge("lightllm_eplb_topk_expert_imbalance_ratio_p100") def create_histogram(self, name, buckets, labelnames=None): all_labels = ["model_name"] + (labelnames or []) diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index c4317df652..48de845967 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -70,6 +70,7 @@ def __init__(self) -> None: self.enable_decode_microbatch_overlap = get_env_start_args().enable_decode_microbatch_overlap self.enable_prefill_microbatch_overlap = get_env_start_args().enable_prefill_microbatch_overlap self.spec_engine = None + self.eplb_manager = None # 控制 _get_classed_reqs 分类的参数变量,不同的 backend 具有可能需要不同的分类运行条件。 self.classed_req_no_decode = False @@ -256,6 +257,15 @@ def init_model(self, kvargs): prof_name = f"lightllm-model_backend-node{self.node_rank}_dev{get_current_device_id()}" prof_mode = self.args.enable_profiling self.profiler = ProcessProfiler(mode=prof_mode, name=prof_name, use_multi_thread=True) if prof_mode else None + if self.args.eplb_num_redundant_experts_per_rank > 0: + from lightllm.server.router.model_infer.mode_backend.eplb.runtime_manager import EPLBManager + + self.eplb_manager = EPLBManager( + self.model, + max_rebalance_count=self.args.eplb_rebalance_count, + config_path=self.args.eplb_config_path, + plan_mode=self.args.eplb_plan_mode, + ) # 启动infer_loop_thread, 启动两个线程进行推理,对于具备双batch推理折叠得场景 # 可以降低 cpu overhead,大幅提升gpu得使用率。 diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index 2326c8c515..50fc80cd54 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py @@ -63,6 +63,12 @@ def infer_loop(self): self._try_read_new_reqs() + # EPLB step 可能发起所有 rank 都必须按相同顺序参与的控制面 + # collective。固定放在请求读取之后、常规通信和 forward 之前, + # 即使本 rank 当前没有请求,也不会与后续 collective 交错。 + if self.eplb_manager is not None: + self.eplb_manager.step() + prefill_reqs, decode_reqs = self._get_classed_reqs( no_decode=self.classed_req_no_decode, strict_prefill=self.classed_req_strict_prefill, diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py index ce21d01987..ace0d2b617 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py @@ -124,6 +124,12 @@ def infer_loop(self): self._try_read_new_reqs() + # EPLB step 可能发起所有 rank 都必须按相同顺序参与的控制面 + # collective。固定放在请求读取之后、常规通信和 forward 之前, + # 即使本 rank 当前没有请求,也不会与后续 collective 交错。 + if self.eplb_manager is not None: + self.eplb_manager.step() + prefill_reqs, decode_reqs = self._get_classed_reqs( no_decode=self.classed_req_no_decode, strict_prefill=self.classed_req_strict_prefill, diff --git a/lightllm/server/router/model_infer/mode_backend/eplb/__init__.py b/lightllm/server/router/model_infer/mode_backend/eplb/__init__.py new file mode 100644 index 0000000000..ca9366913f --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/eplb/__init__.py @@ -0,0 +1 @@ +"""EPLB 布局规划、传输规划和运行时管理。""" diff --git a/lightllm/server/router/model_infer/mode_backend/eplb/async_expert_transfer.py b/lightllm/server/router/model_infer/mode_backend/eplb/async_expert_transfer.py new file mode 100644 index 0000000000..fe1c79ec4f --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/eplb/async_expert_transfer.py @@ -0,0 +1,373 @@ +"""EPLB 专家权重的异步逐层迁移。""" + +import zlib +from dataclasses import dataclass +from typing import List, Sequence + +import torch +import torch.distributed as dist + +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.fused_moe_weight import FusedMoeWeight + +from .async_task import EPLBAsyncTask +from .eplb_utils import NamedTensor, extract_eplb_expert_tensors + + +@dataclass(frozen=True) +class ExpertTensorBuffer: + """一项 live 专家张量及其单专家 pinned memory 缓冲行。""" + + name: str + live_tensor: torch.Tensor + pinned_row: torch.Tensor + + +@dataclass(frozen=True) +class EPLBTransferInfo: + """单个逻辑专家的一次传输描述。 + + ``layer_index`` 和 ``source_logical_expert_id`` 标识需要传输的专家; + ``source_rank``、``source_local_expert_index`` 描述当前物理槽, + ``dest_rank``、``dest_local_expert_index`` 描述目标物理槽。 + """ + + layer_index: int + source_logical_expert_id: int + source_rank: int + source_local_expert_index: int + dest_rank: int + dest_local_expert_index: int + + +class PinnedMemoryEPLBTransfer(EPLBAsyncTask): + """在后台线程中传输一个逻辑专家的全部权重张量。 + + 只有源 rank 和目标 rank 参与 Gloo 点对点通信,具体的数据路径是: + + ``源 GPU 权重行 -> 源 rank 的 pinned CPU 行 -> 目标 rank 的 pinned CPU 行`` + + 如果源和目标是同一个 rank,则只执行 GPU 到 pinned CPU 的本地复制,不产生 + 网络通信。其他 rank 不分配 pinned row,也不参与该专家的数据传输。 + + 本类只负责异步传输,不修改 live 权重,也不更新路由 metadata。传输成功后, + ``status`` 会变为 ``"succeeded"``,收到的数据保存在 ``tensor_buffers``。 + EPLBManager 在主循环的安全边界同步提交这些数据。 + + 每个对象只表示构造函数中 ``transfer_info`` 指定的一次传输。源 rank 直接 + 读取 ``source_local_expert_index`` 指定的物理行;目标物理槽位不属于传输 + 职责,由 manager 根据目标 placement 决定。 + """ + + def __init__( + self, + *, + weights: Sequence[FusedMoeWeight], + transfer_group: dist.ProcessGroup, + current_global_rank: int, + transfer_info: EPLBTransferInfo, + ) -> None: + self._p2p_group: dist.ProcessGroup = transfer_group + self._is_source_rank: bool = current_global_rank == transfer_info.source_rank + self._is_destination_rank: bool = current_global_rank == transfer_info.dest_rank + self.transfer_info: EPLBTransferInfo = transfer_info + + layer_weight: FusedMoeWeight = weights[transfer_info.layer_index] + # 提取该 MoE 层实际参与推理的专家张量,包括 w13/w2 的量化后权重 + # (或非量化权重),以及配套的 weight_scale、weight_zero_point 等量化 + # 信息。后续会为每项张量创建对应的 pinned row,确保专家状态完整迁移。 + named_live_tensors: List[NamedTensor] = extract_eplb_expert_tensors(layer_weight) + self._device: torch.device = named_live_tensors[0][1].device + + # 只有源和目标 rank 需要保存该专家的 pinned row。源 rank 用它作为 + # send buffer,目标 rank 用它作为 recv buffer,并在传输完成后直接交给 + # manager 提交到 live 权重,避免无关 rank 分配同样大小的 pinned memory。 + self.tensor_buffers: List[ExpertTensorBuffer] = [] + if self._is_source_rank or self._is_destination_rank: + for tensor_name, live_tensor in named_live_tensors: + pinned_row = torch.empty( + tuple(live_tensor.shape[1:]), + dtype=live_tensor.dtype, + device="cpu", + pin_memory=True, + ) + self.tensor_buffers.append( + ExpertTensorBuffer( + name=tensor_name, + live_tensor=live_tensor, + pinned_row=pinned_row, + ) + ) + self._device_to_host_stream: torch.cuda.Stream = torch.cuda.Stream(device=self._device) + + super().__init__( + thread_name=( + f"eplb-transfer-layer-{transfer_info.layer_index}-expert-{transfer_info.source_logical_expert_id}" + ) + ) + + def execute(self) -> None: + """把指定专家的全部权重行传输到各 rank 的 pinned memory。""" + transfer_info: EPLBTransferInfo = self.transfer_info + if self._is_source_rank: + torch.cuda.set_device(self._device) + with torch.cuda.stream(self._device_to_host_stream): + for tensor_buffer in self.tensor_buffers: + tensor_buffer.pinned_row.copy_( + tensor_buffer.live_tensor[transfer_info.source_local_expert_index], + non_blocking=True, + ) + # Gloo 读取 pinned row 前,源 rank 必须等待 GPU -> CPU 拷贝完成。 + self._device_to_host_stream.synchronize() + + if transfer_info.source_rank != transfer_info.dest_rank: + if self._is_source_rank: + for tensor_buffer in self.tensor_buffers: + message_tag = self._build_p2p_message_tag(tensor_buffer.name) + dist.send( + tensor_buffer.pinned_row, + dst=transfer_info.dest_rank, + group=self._p2p_group, + tag=message_tag, + ) + elif self._is_destination_rank: + for tensor_buffer in self.tensor_buffers: + message_tag = self._build_p2p_message_tag(tensor_buffer.name) + dist.recv( + tensor_buffer.pinned_row, + src=transfer_info.source_rank, + group=self._p2p_group, + tag=message_tag, + ) + + def _build_p2p_message_tag(self, tensor_name: str) -> int: + """为当前专家张量生成 source 和 destination 一致的 Gloo 整数 tag。 + + Python ``hash`` 会因进程随机种子不同而产生不同结果,因此这里使用稳定的 + CRC32,并限制到 Gloo 可安全使用的有符号 31 位整数范围。标识中包含层、 + 逻辑专家、源物理槽、目标物理槽和张量名称,避免依赖张量列表的隐式顺序。 + """ + transfer_info: EPLBTransferInfo = self.transfer_info + message_identity = ( + f"{transfer_info.layer_index}:" + f"{transfer_info.source_logical_expert_id}:" + f"{transfer_info.source_rank}:" + f"{transfer_info.source_local_expert_index}:" + f"{transfer_info.dest_rank}:" + f"{transfer_info.dest_local_expert_index}:" + f"{tensor_name}" + ) + return zlib.crc32(message_identity.encode("utf-8")) & 0x7FFFFFFF + + +def build_transfer_plan( + current_placement: Sequence[Sequence[int]], + target_placement: Sequence[Sequence[int]], + layer_index: int, + num_logical_experts: int, + world_size: int, +) -> List[List[EPLBTransferInfo]]: + """生成一层中所有发生变化的专家传输批次。 + + ``current_placement`` 和 ``target_placement`` 的形状均为 + ``[world_size, num_local_experts]``。所有物理槽位都允许变化,因此传输源 + 必须从当前实际存在的副本中选择。 + + 规划分为三个阶段: + + 1. 建立当前槽位索引,同时找出布局调整前后专家不变的稳定槽位; + 2. 为每个变化的目标槽位绑定一个确定的源槽位。优先使用稳定副本,因为 + 这种源槽位永远不会被本轮迁移覆盖;没有稳定副本时,循环使用当前已有 + 的各个副本,避免把全部读取集中到同一个 rank; + 3. 根据槽位覆盖依赖生成两类执行批次。目标槽位不再作为任何待处理任务源 + 的任务属于“安全任务”,彼此不要求原子提交,可以继续拆成较小批次并 + 提前 commit;若不存在安全任务,剩余依赖必然由一个或多个环组成,环内 + 任务不可拆分,必须全部传入 pinned memory 后统一 commit。 + + 图中使用 ``[槽位:当前专家] --传输专家--> [目标槽位]`` 表示一条任务。 + + 链式依赖示例 + ------------ + 当前布局和目标布局分别为: + + ``current: [A:e0] [B:e1] [C:e2] [D:e2]`` + ``target: [A:e0] [B:e0] [C:e1] [D:e2]`` + + 需要执行的传输形成一条依赖链: + + ``[A:e0] --e0--> [B:e1] --e1--> [C:e2]`` + + 初始 ``source_slots={A, B}``,因此只有 C 可以覆盖。虽然 C 中的 e2 被 + 覆盖,但 D 中仍有稳定的 e2;第一批执行 ``B --e1--> C`` 后,e1 已经在 + C 中建立新副本,B 才不再作为源。第二批再执行 ``A --e0--> B``,最终 + 得到目标布局。 + + 这个过程同时保护传入和被覆盖的专家:如果 B 保存的是 e1 的最后一个在线 + 副本,而目标布局仍要求保留 e1,那么一定存在一条以 B 为源的待处理任务。 + 此时 ``B in source_slots``,任何以 B 为目标的任务都不会进入安全批次。只有 + e1 已经存在于其他稳定槽位,或者前一批已经为 e1 建立新位置后,B 才允许 + 被覆盖。 + + 环形依赖示例 + ------------ + 当前布局为 ``[A:e0] [B:e1] [C:e2]``,目标布局为 + ``[A:e1] [B:e2] [C:e0]``,依赖关系为: + + ``[A:e0] --e0--> [C:e2] --e2--> [B:e1] --e1--> [A:e0]`` + + A、B、C 都既是源又是目标,不存在安全目标槽位。三条任务必须组成同一个 + 批次:先将 e0、e1、e2 全部传入 pinned memory,等全部传输完成后再统一 + 覆盖 A、B、C,最后发布新 metadata。 + + 返回值按执行顺序保存最终的 commit 批次:同一安全波次会按参与 rank 拆成 + 若干小批次,每个 rank 在一个小批次中最多参与一条任务;不冲突的 rank 仍 + 可并行传输,已完成的小批次也可以立即提交。环批次则始终包含完整环。这样 + 普通批次在每个 rank 上最多缓存一个专家,只有环形依赖才需要同时缓存多个 + 专家,同时仍保证任何源专家都不会在最后一次读取之前被覆盖。 + """ + assert world_size > 0 + assert num_logical_experts % world_size == 0 + assert len(current_placement) == len(target_placement) == world_size + num_local_experts_per_rank = len(current_placement[0]) + assert all(len(row) == num_local_experts_per_rank for row in current_placement) + assert all(len(row) == num_local_experts_per_rank for row in target_placement) + assert all(0 <= expert < num_logical_experts for row in current_placement for expert in row) + assert all(0 <= expert < num_logical_experts for row in target_placement for expert in row) + + # 阶段 1:记录每个逻辑专家当前位于哪些物理槽位,并单独记录不会变化的 + # 稳定槽位。槽位统一表示为 (rank, local_expert_index)。 + Slot = tuple[int, int] + current_slots_by_expert: List[List[Slot]] = [[] for _ in range(num_logical_experts)] + stable_slots_by_expert: List[List[Slot]] = [[] for _ in range(num_logical_experts)] + for rank, (current_row, target_row) in enumerate(zip(current_placement, target_placement)): + for local_expert_index, (current_expert, target_expert) in enumerate(zip(current_row, target_row)): + slot = (rank, local_expert_index) + current_slots_by_expert[current_expert].append(slot) + if current_expert == target_expert: + stable_slots_by_expert[current_expert].append(slot) + assert all(current_slots_by_expert), "current placement must contain every logical expert" + assert set(expert for row in target_placement for expert in row) == set(range(num_logical_experts)) + + # 阶段 2:为每个变化的目标槽位绑定一个确定的当前源槽位。 + # + # 稳定副本不会出现在任何任务的目标位置,因此可以反复读取而没有覆盖 + # 风险。只有不存在稳定副本时,才循环使用该专家当前已有的所有副本。 + # 源槽位完整保存在 transfer_info 中,后续依赖分析和实际传输共用同一份信息。 + source_use_count = [0] * num_logical_experts + pending_transfers: List[EPLBTransferInfo] = [] + for destination_rank, (current_row, target_row) in enumerate(zip(current_placement, target_placement)): + for destination_local_expert_index in range(num_local_experts_per_rank): + current_expert_id = current_row[destination_local_expert_index] + target_expert_id = target_row[destination_local_expert_index] + if target_expert_id == current_expert_id: + continue + source_slots = stable_slots_by_expert[target_expert_id] or current_slots_by_expert[target_expert_id] + source_slot = source_slots[source_use_count[target_expert_id] % len(source_slots)] + source_use_count[target_expert_id] += 1 + transfer_info = EPLBTransferInfo( + layer_index=layer_index, + source_logical_expert_id=target_expert_id, + source_rank=source_slot[0], + source_local_expert_index=source_slot[1], + dest_rank=destination_rank, + dest_local_expert_index=destination_local_expert_index, + ) + pending_transfers.append(transfer_info) + + # 阶段 3:按照槽位覆盖依赖,将任务拆成可安全提交的执行批次。 + transfer_batches: List[List[EPLBTransferInfo]] = [] + while pending_transfers: + source_slots = { + (transfer_info.source_rank, transfer_info.source_local_expert_index) for transfer_info in pending_transfers + } + + # 3.1 收集当前拓扑层次的全部安全任务。source_slots 是当前仍需保护的 + # 槽位集合:只要某个槽位中的专家尚未完成最后一次读取,该槽位就仍在 + # 集合中,任何以它为目标的任务都不能提交。这同时保护了目标槽位里即将 + # 被覆盖的旧专家,避免其最后一个在线副本被提前删除。 + # + # 目标槽位不在 source_slots 的任务可以并行传输,并在整批完成后统一 + # 提交。若旧专家仍需迁往其他位置,该目标槽位必然也是相应任务的源, + # 因而不会在本轮被选中;若它不是源,则旧专家已经有其他可用副本。 + # + # 这里不能在找到第一个任务后立即修改 source_slots。只有整批提交并从 + # pending 中移除后,下一层目标槽位才真正变得安全。 + safe_transfer_batch: List[EPLBTransferInfo] = [] + remaining_transfers: List[EPLBTransferInfo] = [] + for transfer_info in pending_transfers: + destination_slot = (transfer_info.dest_rank, transfer_info.dest_local_expert_index) + if destination_slot not in source_slots: + safe_transfer_batch.append(transfer_info) + else: + remaining_transfers.append(transfer_info) + + if safe_transfer_batch: + # 安全任务之间没有原子提交要求,但若同一 rank 在一个批次中参与 + # 多条任务,就会同时创建多份专家 pinned buffer。这里按 rank 冲突 + # 继续拆分:每个 rank 在一个小批次中最多参与一条任务,不冲突的 + # rank 仍可并行传输,从而兼顾吞吐和 pinned memory 峰值。 + unbatched_transfers = safe_transfer_batch + while unbatched_transfers: + current_batch: List[EPLBTransferInfo] = [] + occupied_ranks: set[int] = set() + deferred_transfers: List[EPLBTransferInfo] = [] + + # 顺序扫描尚未分组的任务:rank 不冲突的任务进入当前批次, + # 冲突任务留到下一轮。每轮至少取出一个任务,因此一定结束。 + for transfer_info in unbatched_transfers: + participant_ranks = {transfer_info.source_rank, transfer_info.dest_rank} + if participant_ranks & occupied_ranks: + deferred_transfers.append(transfer_info) + else: + current_batch.append(transfer_info) + occupied_ranks.update(participant_ranks) + + transfer_batches.append(current_batch) + unbatched_transfers = deferred_transfers + + pending_transfers = remaining_transfers + else: + # 3.2 没有叶子时,每个目标槽位也一定是某条任务的源槽位。每个目标 + # 槽位只有一条写入任务,因此此时源槽位也不会重复,剩余依赖图必然 + # 分解为若干互不相交的简单环。任选第一条任务的源槽位,沿着 + # source_slot -> destination_slot 追踪,回到起点便得到一个完整环。 + # + # 这里有 N 个互不重复的目标槽位,并且没有安全任务意味着这 N 个 + # 目标都包含在 source_slots 中。source_slots 最多也只有 N 项,因此 + # 它必然恰好有 N 项,即每个源槽位只对应一个目标;先显式校验这个 + # 条件,再构造字典,不会因重复 key 丢失任务。 + # + # 例如 ``S -> A、S -> B、A -> S`` 中,源集合只有 ``{S, A}``, + # B 不在源集合中,所以 ``S -> B`` 会先作为安全任务移除;剩余的 + # ``S -> A、A -> S`` 才会进入这里,并且每个源都只对应一个目标。 + assert len(source_slots) == len(pending_transfers) + transfer_by_source_slot = { + (transfer_info.source_rank, transfer_info.source_local_expert_index): transfer_info + for transfer_info in pending_transfers + } + + first_transfer = pending_transfers[0] + cycle_start_slot = (first_transfer.source_rank, first_transfer.source_local_expert_index) + source_slot = cycle_start_slot + cycle_batch: List[EPLBTransferInfo] = [] + cycle_source_slots: set[Slot] = set() + + # 从任意源槽位出发,当前任务的目标槽位就是下一条任务的源槽位; + # 目标重新回到起点时,一个完整环便已经收集完成。 + while True: + assert source_slot not in cycle_source_slots + cycle_source_slots.add(source_slot) + transfer_info = transfer_by_source_slot[source_slot] + cycle_batch.append(transfer_info) + destination_slot = (transfer_info.dest_rank, transfer_info.dest_local_expert_index) + if destination_slot == cycle_start_slot: + break + source_slot = destination_slot + + transfer_batches.append(cycle_batch) + pending_transfers = [ + transfer_info + for transfer_info in pending_transfers + if (transfer_info.source_rank, transfer_info.source_local_expert_index) not in cycle_source_slots + ] + + return transfer_batches diff --git a/lightllm/server/router/model_infer/mode_backend/eplb/async_load_gather_task.py b/lightllm/server/router/model_infer/mode_backend/eplb/async_load_gather_task.py new file mode 100644 index 0000000000..48e18efea3 --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/eplb/async_load_gather_task.py @@ -0,0 +1,42 @@ +"""EPLB 原始路由负载的异步汇集任务。""" + +from typing import Optional + +import torch +import torch.distributed as dist + +from .async_task import EPLBAsyncTask + + +class EPLBLoadGatherTask(EPLBAsyncTask): + """在独立 Gloo 通信组中汇集每个 rank 的逐样本路由负载。""" + + def __init__( + self, + *, + local_load_samples: torch.Tensor, + load_gather_group: dist.ProcessGroup, + ) -> None: + assert local_load_samples.device.type == "cpu" + assert local_load_samples.ndim == 3 + self.local_load_samples = local_load_samples.contiguous() + self.load_gather_group = load_gather_group + self.world_size = dist.get_world_size(group=load_gather_group) + assert self.world_size > 0 + self.result: Optional[torch.Tensor] = None + super().__init__(thread_name="eplb-load-gather") + + def execute(self) -> None: + """生成 ``[rank, layer, sample, logical_expert]`` 的连续结果。""" + gathered_load = torch.empty( + (self.world_size, *self.local_load_samples.shape), + dtype=self.local_load_samples.dtype, + device=self.local_load_samples.device, + ) + load_by_rank = list(gathered_load.unbind(dim=0)) + dist.all_gather( + load_by_rank, + self.local_load_samples, + group=self.load_gather_group, + ) + self.result = gathered_load diff --git a/lightllm/server/router/model_infer/mode_backend/eplb/async_placement_plan_task.py b/lightllm/server/router/model_infer/mode_backend/eplb/async_placement_plan_task.py new file mode 100644 index 0000000000..8d380fd63b --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/eplb/async_placement_plan_task.py @@ -0,0 +1,32 @@ +"""EPLB 专家布局的异步规划任务。""" + +from typing import Optional + +import torch + +from .async_task import EPLBAsyncTask +from .placement import EPLBPlanner, ExpertPlacement + + +class EPLBPlanTask(EPLBAsyncTask): + """在后台线程中根据全局专家负载生成新布局。""" + + def __init__( + self, + *, + planner: EPLBPlanner, + logical_expert_load_samples: torch.Tensor, + current_placement: ExpertPlacement, + ) -> None: + self.planner = planner + self.logical_expert_load_samples = logical_expert_load_samples + self.current_placement = current_placement + self.result: Optional[ExpertPlacement] = None + super().__init__(thread_name="eplb-plan") + + def execute(self) -> None: + """根据全局 logical expert 负载生成目标布局。""" + self.result = self.planner.plan( + logical_expert_load_samples=self.logical_expert_load_samples, + current_placement=self.current_placement, + ) diff --git a/lightllm/server/router/model_infer/mode_backend/eplb/async_task.py b/lightllm/server/router/model_infer/mode_backend/eplb/async_task.py new file mode 100644 index 0000000000..b7c901180f --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/eplb/async_task.py @@ -0,0 +1,43 @@ +"""EPLB 后台线程任务的公共生命周期。""" + +import os +import threading + +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) + + +class EPLBAsyncTask: + """提供单次后台任务统一的启动、完成和失败处理。""" + + def __init__(self, *, thread_name: str) -> None: + self.status = "idle" + self._thread = threading.Thread( + target=self._run, + name=thread_name, + daemon=True, + ) + + def start(self) -> None: + """启动后台任务;同一个任务对象只能启动一次。""" + assert self.status == "idle", f"{type(self).__name__} has already been started" + self.status = "running" + self._thread.start() + + def is_finished(self) -> bool: + """返回后台任务是否已经成功完成。""" + return self.status == "succeeded" + + def _run(self) -> None: + try: + self.execute() + except BaseException: + logger.exception(f"{type(self).__name__} failed") + os._exit(1) + else: + self.status = "succeeded" + + def execute(self) -> None: + """执行子类定义的具体后台任务。""" + raise NotImplementedError diff --git a/lightllm/server/router/model_infer/mode_backend/eplb/async_transfer_planner.py b/lightllm/server/router/model_infer/mode_backend/eplb/async_transfer_planner.py new file mode 100644 index 0000000000..c8646dbbe0 --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/eplb/async_transfer_planner.py @@ -0,0 +1,46 @@ +"""EPLB 专家传输计划的异步生成器。""" + +from typing import List, Optional + +from .async_task import EPLBAsyncTask +from .async_expert_transfer import EPLBTransferInfo, build_transfer_plan +from .placement import ExpertPlacement + + +class EPLBTransferPlanner(EPLBAsyncTask): + """在后台线程中逐层生成并按 layer 顺序拼接专家传输批次。 + + 输入布局在规划期间保持只读。每层独立调用 ``build_transfer_plan``,因此 + 一个批次只包含同一层的任务,manager 可以在提交后立即发布该层的路由 + metadata。 + """ + + def __init__( + self, + *, + current_placement: ExpertPlacement, + target_placement: ExpertPlacement, + num_logical_experts: int, + world_size: int, + ) -> None: + self.current_placement = current_placement + self.target_placement = target_placement + self.num_logical_experts = num_logical_experts + self.world_size = world_size + self.result: Optional[List[List[EPLBTransferInfo]]] = None + super().__init__(thread_name="eplb-transfer-plan") + + def execute(self) -> None: + """逐层生成传输批次,并按 layer 顺序保存完整结果。""" + transfer_batches: List[List[EPLBTransferInfo]] = [] + layer_placements = zip(self.current_placement, self.target_placement) + for layer_index, (current_layer, target_layer) in enumerate(layer_placements): + layer_transfer_batches = build_transfer_plan( + current_placement=current_layer, + target_placement=target_layer, + layer_index=layer_index, + num_logical_experts=self.num_logical_experts, + world_size=self.world_size, + ) + transfer_batches.extend(layer_transfer_batches) + self.result = transfer_batches diff --git a/lightllm/server/router/model_infer/mode_backend/eplb/eplb_utils.py b/lightllm/server/router/model_infer/mode_backend/eplb/eplb_utils.py new file mode 100644 index 0000000000..b84afcb20d --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/eplb/eplb_utils.py @@ -0,0 +1,35 @@ +"""EPLB 专家权重提取工具。""" + +from typing import List, Optional, Protocol, Tuple + +import torch + +NamedTensor = Tuple[str, torch.Tensor] + + +class ExpertWeightPack(Protocol): + """EPLB 需要迁移的单组专家权重。""" + + weight: torch.Tensor + weight_scale: Optional[torch.Tensor] + weight_zero_point: Optional[torch.Tensor] + + +class EPLBExpertWeight(Protocol): + """包含门控投影和下投影专家权重的 MoE 层。""" + + w13: ExpertWeightPack + w2: ExpertWeightPack + + +def extract_eplb_expert_tensors(weight: EPLBExpertWeight) -> List[NamedTensor]: + """按固定顺序返回 EPLB 必须迁移的权重及量化参数。""" + named_tensors: List[NamedTensor] = [] + for pack_name in ("w13", "w2"): + weight_pack: ExpertWeightPack = getattr(weight, pack_name) + for value_name in ("weight", "weight_scale", "weight_zero_point"): + tensor: Optional[torch.Tensor] = getattr(weight_pack, value_name, None) + if tensor is not None: + assert tensor.ndim >= 1 and tensor.is_contiguous(), f"{pack_name}.{value_name} must be contiguous" + named_tensors.append((f"{pack_name}.{value_name}", tensor)) + return named_tensors diff --git a/lightllm/server/router/model_infer/mode_backend/eplb/metrics.py b/lightllm/server/router/model_infer/mode_backend/eplb/metrics.py new file mode 100644 index 0000000000..c94ff2cabb --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/eplb/metrics.py @@ -0,0 +1,142 @@ +"""计算并上报 EPLB 的 logical expert 与 EP rank 负载指标。""" + +import numpy as np +import torch + +from lightllm.server.metrics.manager import MetricClient + +from .placement import ExpertPlacement + + +COMPUTE_CRITICAL_OVERHEAD_RATIO_BEFORE_REBALANCE_METRIC = ( + "lightllm_prefill_ep_compute_critical_overhead_ratio_before_rebalance" +) +COMPUTE_CRITICAL_OVERHEAD_RATIO_AFTER_REBALANCE_METRIC = ( + "lightllm_prefill_ep_compute_critical_overhead_ratio_after_rebalance" +) +EXPERT_IMBALANCE_RATIO_METRICS = { + 25: "lightllm_eplb_topk_expert_imbalance_ratio_p25", + 50: "lightllm_eplb_topk_expert_imbalance_ratio_p50", + 100: "lightllm_eplb_topk_expert_imbalance_ratio_p100", +} + + +def logical_expert_imbalance_percentiles(*, expert_load: torch.Tensor) -> dict[int, float]: + """统计各层 ``最热 logical expert / 本层平均负载`` 的分位数。""" + assert expert_load.ndim == 2 and expert_load.numel() > 0 + + expert_load = expert_load.to(torch.float64) + mean_load_by_layer = expert_load.mean(dim=1) + + # 没有 token 的层不参与分位数计算,避免产生 0 / 0。 + has_route_load = mean_load_by_layer > 0 + if not torch.any(has_route_load): + return {percentile: 0.0 for percentile in EXPERT_IMBALANCE_RATIO_METRICS} + + hottest_expert_load = expert_load.max(dim=1).values + imbalance_ratio_by_layer = hottest_expert_load[has_route_load] / mean_load_by_layer[has_route_load] + percentiles = tuple(EXPERT_IMBALANCE_RATIO_METRICS) + + # inverted_cdf 就是 nearest-rank 定义;默认 linear 或 nearest 插值都会改变 + # 层数较少时的 P25/P50 语义。 + percentile_values = np.percentile( + a=imbalance_ratio_by_layer.numpy(), + q=percentiles, + method="inverted_cdf", + ) + return dict(zip(percentiles, percentile_values.tolist())) + + +def publish_expert_load_metrics(*, metric_client: MetricClient, expert_load: torch.Tensor) -> None: + """上报 logical expert 层间不均衡分位数。""" + imbalance_percentiles = logical_expert_imbalance_percentiles(expert_load=expert_load) + for percentile, metric_name in EXPERT_IMBALANCE_RATIO_METRICS.items(): + metric_client.gauge_set( + name=metric_name, + value=imbalance_percentiles[percentile], + ) + + +def compute_critical_overhead_ratio( + *, + logical_expert_load: torch.Tensor, + placement: ExpertPlacement, + expert_alignment: int, +) -> float: + """估算 placement 相对理想 rank 均衡状态的关键路径额外计算比例。 + + 假设同一 logical expert 的流量由 hash 均匀分配给所有 physical 副本, + 并按 DeepEP 的 expert alignment 对每个副本负载向上取整。每层单独选择 + 最繁忙 rank,避免不同层的热点 rank 在汇总时相互抵消。 + """ + assert logical_expert_load.ndim == 2 and logical_expert_load.numel() > 0 + assert expert_alignment > 0 + + num_layers, num_logical_experts = logical_expert_load.shape + + # placement: [layer, rank, local physical expert] + placement_tensor = torch.tensor(placement, dtype=torch.int64) + assert placement_tensor.ndim == 3 and placement_tensor.shape[0] == num_layers + + # 将每个 physical slot 中保存的 logical expert ID 转成 one-hot: + # [layer, rank, physical slot] -> [layer, rank, physical slot, logical expert]。 + # 例如 slot 中保存 expert 2,就会在 logical expert 维得到 [0, 0, 1, ...]。 + expert_mask_by_physical_slot = torch.nn.functional.one_hot( + placement_tensor, + num_classes=num_logical_experts, + ) + + # 沿 physical slot 维求和,得到每个 rank 持有的专家副本数:[layer, rank, logical expert]。 + replicas_per_rank = expert_mask_by_physical_slot.sum(dim=2).to(torch.float64) + + # 再沿 rank 维求和,得到每个 logical expert 的全局副本数:[layer, logical expert]。 + replica_count_per_expert = replicas_per_rank.sum(dim=1) + assert torch.all(replica_count_per_expert > 0) + + # logical_expert_load: [layer, logical expert] + # hash 均匀分流后,每个 physical 副本承担 logical expert 总负载的 1/N。 + load_per_replica = logical_expert_load.to(torch.float64) / replica_count_per_expert + aligned_load_per_replica = torch.ceil(load_per_replica / expert_alignment) * expert_alignment + + # 广播为 [layer, rank, logical expert] 后沿 expert 维求和,得到每层各 rank + # 的估算计算量:[layer, rank]。 + estimated_rank_load = (replicas_per_rank * aligned_load_per_replica.unsqueeze(dim=1)).sum(dim=2) + + mean_rank_load_by_layer = estimated_rank_load.mean(dim=1) + critical_rank_load_by_layer = estimated_rank_load.max(dim=1).values + total_balanced_compute = mean_rank_load_by_layer.sum() + if total_balanced_compute == 0: + return 0.0 + + total_critical_overhead = (critical_rank_load_by_layer - mean_rank_load_by_layer).sum() + return float((total_critical_overhead / total_balanced_compute).item()) + + +def publish_rebalance_compute_metrics( + *, + metric_client: MetricClient, + sample_load: torch.Tensor, + current_placement: ExpertPlacement, + target_placement: ExpertPlacement, + expert_alignment: int, +) -> None: + """使用同一个 prefill 样本上报重排前后的关键路径开销。""" + before_rebalance_ratio = compute_critical_overhead_ratio( + logical_expert_load=sample_load, + placement=current_placement, + expert_alignment=expert_alignment, + ) + after_rebalance_ratio = compute_critical_overhead_ratio( + logical_expert_load=sample_load, + placement=target_placement, + expert_alignment=expert_alignment, + ) + + metric_client.gauge_set( + name=COMPUTE_CRITICAL_OVERHEAD_RATIO_BEFORE_REBALANCE_METRIC, + value=before_rebalance_ratio, + ) + metric_client.gauge_set( + name=COMPUTE_CRITICAL_OVERHEAD_RATIO_AFTER_REBALANCE_METRIC, + value=after_rebalance_ratio, + ) diff --git a/lightllm/server/router/model_infer/mode_backend/eplb/placement/__init__.py b/lightllm/server/router/model_infer/mode_backend/eplb/placement/__init__.py new file mode 100644 index 0000000000..4d4d664fa0 --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/eplb/placement/__init__.py @@ -0,0 +1,22 @@ +"""Expert placement construction, routing metadata, and planning APIs.""" + +from .types import ExpertPlacement, LayerPlacement, LogicalToPhysicalMap +from .planner import EPLBPlanner +from .initial import build_initial_local_expert_ids +from .routing import build_logical_to_physical_map +from .greedy import GreedyEPLBPlanner +from .factory import create_eplb_planner +from .config import load_layer_placement, save_placement_config + +__all__ = [ + "EPLBPlanner", + "ExpertPlacement", + "GreedyEPLBPlanner", + "LayerPlacement", + "LogicalToPhysicalMap", + "build_initial_local_expert_ids", + "build_logical_to_physical_map", + "create_eplb_planner", + "load_layer_placement", + "save_placement_config", +] diff --git a/lightllm/server/router/model_infer/mode_backend/eplb/placement/config.py b/lightllm/server/router/model_infer/mode_backend/eplb/placement/config.py new file mode 100644 index 0000000000..0a7238e983 --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/eplb/placement/config.py @@ -0,0 +1,230 @@ +"""EPLB expert placement JSON loading and persistence.""" + +import json +import os +from functools import lru_cache +from typing import Any, Dict, Optional, Sequence + +from lightllm.utils.log_utils import init_logger + +from .types import ExpertPlacement, LayerPlacement + +logger = init_logger(__name__) + +EPLB_PLACEMENT_CONFIG_VERSION = 1 + + +def load_layer_placement( + config_path: str, + layer_index: int, + num_logical_experts: int, + world_size: int, + num_redundant_experts_per_rank: int, +) -> Optional[LayerPlacement]: + """读取并校验一个模型层的布局;任何错误都返回 ``None`` 以触发默认流程。""" + config = _read_config(config_path) + if config is None: + return None + + # 阶段 1:构造当前部署期望的元数据,后续各项校验和 warning 都以此为准。 + expected_metadata = { + "version": EPLB_PLACEMENT_CONFIG_VERSION, + "num_logical_experts": num_logical_experts, + "world_size": world_size, + "num_redundant_experts_per_rank": num_redundant_experts_per_rank, + } + + # 阶段 2:逐项显式校验配置元数据,便于直接定位具体的不匹配字段。 + version = config.get("version") + if version != expected_metadata["version"]: + logger.warning( + "EPLB placement config %s does not match the current deployment: version=%r, expected %r; " + "using the default initial placement", + config_path, + version, + expected_metadata["version"], + ) + return None + + config_num_logical_experts = config.get("num_logical_experts") + if config_num_logical_experts != expected_metadata["num_logical_experts"]: + logger.warning( + "EPLB placement config %s does not match the current deployment: num_logical_experts=%r, " + "expected %r; using the default initial placement", + config_path, + config_num_logical_experts, + expected_metadata["num_logical_experts"], + ) + return None + + config_world_size = config.get("world_size") + if config_world_size != expected_metadata["world_size"]: + logger.warning( + "EPLB placement config %s does not match the current deployment: world_size=%r, expected %r; " + "using the default initial placement", + config_path, + config_world_size, + expected_metadata["world_size"], + ) + return None + + config_num_redundant_experts_per_rank = config.get("num_redundant_experts_per_rank") + if config_num_redundant_experts_per_rank != expected_metadata["num_redundant_experts_per_rank"]: + logger.warning( + "EPLB placement config %s does not match the current deployment: " + "num_redundant_experts_per_rank=%r, expected %r; using the default initial placement", + config_path, + config_num_redundant_experts_per_rank, + expected_metadata["num_redundant_experts_per_rank"], + ) + return None + + # 阶段 3:读取当前模型层的物理槽布局,并校验形状、ID 范围及专家覆盖关系。 + try: + layers = config.get("layers") + assert isinstance(layers, dict), "the layers field must be a JSON object" + layer_placement = layers.get(str(layer_index)) + assert layer_placement is not None, f"layer {layer_index} is missing" + _validate_layer_placement( + layer_placement, + num_logical_experts=num_logical_experts, + world_size=world_size, + num_redundant_experts_per_rank=num_redundant_experts_per_rank, + ) + except Exception as exc: + logger.warning( + "Layer %s in EPLB placement config %s is invalid (%s); using the default initial placement", + layer_index, + config_path, + exc, + ) + return None + + # 校验后复制一份,避免缓存中的原始 JSON 对象被运行态修改。 + return [list(rank_placement) for rank_placement in layer_placement] + + +def save_placement_config( + config_path: str, + layer_indexes: Sequence[int], + placement: ExpertPlacement, + num_logical_experts: int, + world_size: int, + num_redundant_experts_per_rank: int, +) -> bool: + """将完整 EPLB 布局原子写入配置路径,失败时仅记录 warning。""" + if len(layer_indexes) != len(placement) or len(set(layer_indexes)) != len(layer_indexes): + logger.warning( + "Failed to save EPLB placement config %s: layer indexes do not match the placement", + config_path, + ) + return False + + for layer_index, layer_placement in zip(layer_indexes, placement): + try: + _validate_layer_placement( + layer_placement, + num_logical_experts=num_logical_experts, + world_size=world_size, + num_redundant_experts_per_rank=num_redundant_experts_per_rank, + ) + except Exception as exc: + logger.warning( + "Failed to save EPLB placement config %s: layer %s is invalid (%s)", + config_path, + layer_index, + exc, + ) + return False + + config = { + "version": EPLB_PLACEMENT_CONFIG_VERSION, + "num_logical_experts": num_logical_experts, + "world_size": world_size, + "num_redundant_experts_per_rank": num_redundant_experts_per_rank, + "layers": { + str(layer_index): [list(rank_placement) for rank_placement in layer_placement] + for layer_index, layer_placement in zip(layer_indexes, placement) + }, + } + + absolute_path = os.path.abspath(config_path) + parent_dir = os.path.dirname(absolute_path) + lock_path = f"{absolute_path}.lock" + lock_acquired = False + try: + os.makedirs(parent_dir, exist_ok=True) + # 通过 x 模式原子创建锁文件,避免多个服务进程同时写入同一个布局文件。 + with open(lock_path, "x", encoding="utf-8") as lock_file: + lock_acquired = True + lock_file.write(str(os.getpid())) + with open(absolute_path, "w", encoding="utf-8") as config_file: + json.dump(config, config_file, ensure_ascii=False, indent=2) + config_file.write("\n") + _read_config.cache_clear() + return True + except OSError as exc: + logger.warning("Failed to save EPLB placement config %s: %s", config_path, exc) + return False + finally: + if lock_acquired: + try: + os.unlink(lock_path) + except OSError as exc: + logger.warning("Failed to remove EPLB placement lock file %s: %s", lock_path, exc) + + +def _validate_layer_placement( + layer_placement: Any, + num_logical_experts: int, + world_size: int, + num_redundant_experts_per_rank: int, +) -> None: + """按顺序断言单层布局满足当前部署的全部约束。""" + # 阶段 1:校验 rank 维度和每个 rank 应持有的物理槽数量。 + assert isinstance(layer_placement, list), "the layer is not an array" + assert len(layer_placement) == world_size, f"the number of ranks is {len(layer_placement)}, expected {world_size}" + + num_physical_experts_per_rank = num_logical_experts // world_size + num_redundant_experts_per_rank + covered_experts = set() + for rank, rank_placement in enumerate(layer_placement): + # 阶段 2:依次校验每个 rank 的布局形状和 expert ID。 + assert isinstance(rank_placement, list), f"the placement for rank {rank} is not an array" + assert ( + len(rank_placement) == num_physical_experts_per_rank + ), f"rank {rank} has {len(rank_placement)} physical slots, expected {num_physical_experts_per_rank}" + assert all( + not isinstance(expert_id, bool) and isinstance(expert_id, int) for expert_id in rank_placement + ), f"rank {rank} contains a non-integer expert ID" + assert all( + 0 <= expert_id < num_logical_experts for expert_id in rank_placement + ), f"rank {rank} contains an out-of-range expert ID" + assert len(set(rank_placement)) == len(rank_placement), f"rank {rank} contains duplicate expert IDs" + covered_experts.update(rank_placement) + + # 阶段 3:确认所有 logical expert 至少存在一个物理副本。 + missing_experts = sorted(set(range(num_logical_experts)) - covered_experts) + assert not missing_experts, f"logical experts are missing: {missing_experts}" + + +@lru_cache(maxsize=None) +def _read_config(config_path: str) -> Optional[Dict[str, Any]]: + """读取并缓存配置,避免模型的每个 MoE 层重复解析同一个文件。""" + try: + with open(config_path, "r", encoding="utf-8") as config_file: + config = json.load(config_file) + except (OSError, UnicodeError, json.JSONDecodeError) as exc: + logger.warning( + "Failed to read EPLB placement config %s; using the default initial placement: %s", + config_path, + exc, + ) + return None + + if not isinstance(config, dict): + logger.warning( + "The root of EPLB placement config %s must be a JSON object; using the default initial placement", + config_path, + ) + return None + return config diff --git a/lightllm/server/router/model_infer/mode_backend/eplb/placement/factory.py b/lightllm/server/router/model_infer/mode_backend/eplb/placement/factory.py new file mode 100644 index 0000000000..97c9e24dcd --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/eplb/placement/factory.py @@ -0,0 +1,25 @@ +"""EPLB placement planner selection.""" + +from typing import Callable, Dict + +from .greedy import GreedyEPLBPlanner +from .planner import EPLBPlanner + + +def create_eplb_planner( + plan_mode: str, + num_ranks: int, + num_redundant_experts_per_rank: int, + expert_alignment: int, +) -> EPLBPlanner: + """根据启动参数为当前推理进程创建布局规划器。""" + planner_builders: Dict[str, Callable[[], EPLBPlanner]] = { + "greedy": lambda: GreedyEPLBPlanner( + num_ranks, + num_redundant_experts_per_rank, + expert_alignment=expert_alignment, + ), + } + if plan_mode not in planner_builders: + raise ValueError(f"unsupported EPLB plan mode {plan_mode!r}; expected one of {tuple(planner_builders)}") + return planner_builders[plan_mode]() diff --git a/lightllm/server/router/model_infer/mode_backend/eplb/placement/greedy.py b/lightllm/server/router/model_infer/mode_backend/eplb/placement/greedy.py new file mode 100644 index 0000000000..5796c5b8ea --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/eplb/placement/greedy.py @@ -0,0 +1,367 @@ +"""使用 CPU Tensor 输入和纯 Python 核心逻辑实现贪心 EPLB 布局规划。 + +入口保留 all-gather 产生的 rank、layer、sample 和 logical expert 维度; +Greedy planner 先完成必要的聚合,再转换为嵌套 list。实际贪心分析仍只使用 +Python 数值和容器,更易于阅读、测试和替换算法。 +""" + +import heapq +from math import ceil +from typing import List + +import torch + +from .planner import EPLBPlanner +from .types import ExpertPlacement, ExpertReplicaGroup, LayerPlacement + + +class GreedyEPLBPlanner(EPLBPlanner): + """使用两阶段启发式算法生成完整的专家布局。 + + 设计目标 + -------- + 对每一层的全局逻辑专家负载进行快速近似均衡,同时满足以下约束: + + * 每个 rank 的物理专家槽位数量必须固定; + * 每个逻辑专家至少有一个副本; + * 同一个逻辑专家不能在同一个 rank 上出现两次; + * 尽量减少最繁忙 rank 的负载,并保持结果确定,方便规划和迁移测试。 + + ``num_redundant_experts_per_rank`` 个槽位用于复制专家。这里的“冗余专家” + 是布局中的副本槽位,不代表某些主专家槽位不可移动;当前实现允许所有 + 物理槽位重新排列。 + + 单层规划流程 + ------------ + 1. 选择全卡冗余专家:取负载最高的 ``R`` 个逻辑专家,并将它们各放一份 + 到所有 rank。它们在每个 rank 上的负载贡献完全相同,因此在后续比较 + rank 之间的相对负载时可以暂时忽略。 + 2. 决定额外副本数:全卡冗余专家占用了 ``R * world_size`` 个槽位,剩余 + 槽位总数正好是逻辑专家数 ``E``。非冗余专家先各保留一个副本,再将 + 多出的 ``R`` 个副本逐次加给当前 ``load / replica_count`` 最高的专家。 + 这样拆分后的每个副本负载尽量接近。 + 3. 平铺多副本专家:副本数大于 1 的专家按顺序使用循环 rank 游标放置。 + 一组副本数最多为 ``world_size``,所以它会落在互不相同的 rank;连续 + 使用游标还使第一阶段各 rank 的槽位数最多相差一个。 + 4. 放置单副本专家:根据第三步已经形成的 ``rank_load`` 建立最小堆,按 + 单副本负载从高到低取专家,每次分给当前负载最低且仍有空槽的 rank。 + 此阶段只有一个副本的专家,不存在同专家重复约束;rank 填满后从堆中 + 移除。这里负载是第一优先级,剩余槽位数量只用于判断 rank 是否已满。 + 5. 复用当前布局:按当前 rank 顺序执行贪心匹配,每次从尚未使用的候选 + rank 中选择共同专家数量最多的一行,先减少跨 rank 的专家迁移;随后 + 让共同专家继续占用原物理槽位,再减少 rank 内的权重搬运。候选 rank + 重排和本地槽位复用都不会改变规划负载。 + + 示例 + ---- + 假设 ``E=8``、``world_size=4``、``R=1``、``expert_alignment=1``,逻辑 + 专家负载为 ``[40, 12, 9, 8, 7, 5, 4, 3]``。总物理槽位数为 + ``E + R * world_size = 12``,所以每个 rank 必须恰好放置三个专家。 + + 1. 选择全卡冗余专家 + + 专家 0 的负载 40 最高,因此每个 rank 都先放置专家 0。四个副本各 + 分担 ``40 / 4 = 10`` 的负载: + + ``placement = [[0], [0], [0], [0]]`` + + 此时每个 rank 的公共负载都是 10、剩余槽位都是 2。公共负载不会影响 + rank 之间的大小关系,所以后续平衡过程只记录非冗余专家的负载。 + + 2. 计算非冗余专家的副本数 + + 去掉专家 0 后,专家 1 到 7 先各保留一个副本,只能占用 7 个槽位;但 + 当前共有 8 个剩余槽位,因此还需要增加一个副本。专家 1 的当前单副本 + 负载 12 最高,所以将其拆成两个负载为 6 的副本。最终专家组为: + + ``[(expert=1, copies=2, load=6),`` + `` (expert=2..7, copies=1, load=9, 8, 7, 5, 4, 3)]`` + + 3. 第一阶段平铺多副本专家 + + 循环游标从 rank 0 开始,将专家 1 的两个副本依次放到 rank 0、1: + + ``placement = [[0, 1], [0, 1], [0], [0]]`` + ``rank_load = [6, 6, 0, 0]`` + ``remaining_slots = [1, 1, 2, 2]`` + + 4. 第二阶段分配单副本专家 + + 专家 2 到 7 已按负载从高到低排列。每次从最小堆中取当前负载最低的 + 未满 rank;负载相同时使用 rank ID 打破平局: + + * 专家 2,负载 9:放到 rank 2,负载变为 ``[6, 6, 9, 0]``; + * 专家 3,负载 8:放到 rank 3,负载变为 ``[6, 6, 9, 8]``; + * 专家 4,负载 7:放到 rank 0,负载变为 ``[13, 6, 9, 8]``,rank 0 填满; + * 专家 5,负载 5:放到 rank 1,负载变为 ``[13, 11, 9, 8]``,rank 1 填满; + * 专家 6,负载 4:放到 rank 3,负载变为 ``[13, 11, 9, 12]``,rank 3 填满; + * 专家 7,负载 3:放到 rank 2,负载变为 ``[13, 11, 12, 12]``,rank 2 填满。 + + 最终候选布局为: + + ``rank 0: [0, 1, 4]`` + ``rank 1: [0, 1, 5]`` + ``rank 2: [0, 2, 7]`` + ``rank 3: [0, 3, 6]`` + + 将专家 0 的公共负载 10 加回来后,完整 rank 负载为 + ``[23, 21, 22, 22]``,平均负载为 22,最大负载为 23。 + + 5. 复用当前布局 + + 上述 rank 编号只是负载分组结果。规划器会依次处理当前 rank 0 到 3, + 每次从尚未匹配的候选行中选择共同专家最多的一行;匹配完成后,再让 + 共同专家尽量保留原物理槽位。两次排序都不改变负载,完成后直接返回 + 新布局。 + + 合法性和确定性 + -------------- + * 专家覆盖:全卡冗余专家和非冗余专家组来自互斥集合,并且两者合起来 + 包含所有逻辑专家,所以不会遗漏任何专家。 + * 本地去重:一个多副本专家最多有 ``world_size`` 个副本,循环游标在一组 + 副本分配完成前不会第二次经过同一 rank;单副本专家只放置一次。因此 + 同一逻辑专家不会在一个 rank 上出现两次。 + * 槽位守恒:全卡冗余专家放置完成后,剩余副本总数严格等于剩余槽位总数。 + 第一阶段连续平铺使各 rank 的已用槽位数最多相差一个;第二阶段只从尚有 + 空位的 rank 中选择,并在 rank 填满后将其移出最小堆,最终所有槽位恰好 + 填满。 + * 结果确定:专家组排序和最小堆比较最终都使用专家 ID 或 rank ID 打破 + 平局,因此相同输入始终得到相同布局。 + """ + + def __init__( + self, + world_size: int, + num_redundant_experts_per_rank: int, + *, + expert_alignment: int = 1, + ): + if world_size <= 1: + raise ValueError("world_size must be greater than one") + if num_redundant_experts_per_rank <= 0: + raise ValueError("num_redundant_experts_per_rank must be positive") + if expert_alignment <= 0: + raise ValueError("expert_alignment must be positive") + self.world_size = world_size + self.num_redundant_experts_per_rank = num_redundant_experts_per_rank + self.expert_alignment = expert_alignment + + def plan( + self, + *, + logical_expert_load_samples: torch.Tensor, + current_placement: ExpertPlacement, + ) -> ExpertPlacement: + """聚合全局逐样本负载,逐层规划并组合成完整的多层布局。""" + # logical_expert_load_samples: [rank, layer, sample, logical_expert] + # CPU Tensor。 + # Greedy 算法只需要整个采样窗口内每层各 logical expert 的累计负载, + # 因此沿 rank 和 sample 维求和为 [layer, logical_expert],再转成 list + # 进入后续纯 Python 分析逻辑。 + assert logical_expert_load_samples.device.type == "cpu" + assert logical_expert_load_samples.ndim == 4 + assert logical_expert_load_samples.shape[0] == self.world_size + aggregated_load = logical_expert_load_samples.sum(dim=(0, 2)).to(torch.float64).tolist() + + # 一次性校验所有层的形状和布局约束。后续每层规划之间没有共享的 + # 可变状态。 + self._validate_inputs(aggregated_load, current_placement) + + # 每层只依赖自己的逻辑专家负载和当前布局。先完成单层规划,再将结果 + # 按原 layer 顺序组合,避免多层候选和负载数据交叉索引。 + return [ + self._plan_layer(layer_load, current_layer) + for layer_load, current_layer in zip(aggregated_load, current_placement) + ] + + def _plan_layer( + self, + logical_load: List[float], + current_placement: LayerPlacement, + ) -> LayerPlacement: + """完成单层副本分配、rank 排布和物理槽位复用。""" + # 阶段 1:选出最热的 R 个专家,并为每个 rank 固定预留它们的副本。 + redundant_experts = self._select_redundant_experts(logical_load) + + # 阶段 2:只在非冗余专家中增加副本,数量恰好填满所有剩余槽位。 + remaining_expert_groups = self._build_remaining_expert_groups(logical_load, redundant_experts) + + # 阶段 3:每个 rank 先放入相同的冗余专家,再通过 rank 优先队列 + # 分配其余专家。同一专家的一组副本会一次性放到不同 rank。 + candidate_placement = self._distribute_remaining_experts(redundant_experts, remaining_expert_groups) + + # 阶段 4:先将候选行贪心匹配到最相似的当前 rank,再复用原物理槽位, + # 依次减少跨 rank 迁移和 rank 内部的槽位搬运。 + return self._reuse_current_slots(candidate_placement, current_placement) + + def _select_redundant_experts(self, logical_load: List[float]) -> List[int]: + """选择需要在所有 rank 上固定放置的最热专家。""" + return sorted(range(len(logical_load)), key=lambda expert: (-logical_load[expert], expert))[ + : self.num_redundant_experts_per_rank + ] + + def _build_remaining_expert_groups( + self, + logical_load: List[float], + redundant_experts: List[int], + ) -> List[ExpertReplicaGroup]: + """确定非冗余专家的副本数,并按安全的分配顺序组成专家组。""" + redundant_expert_set = set(redundant_experts) + remaining_experts = [expert for expert in range(len(logical_load)) if expert not in redundant_expert_set] + replica_counts = {expert: 1 for expert in remaining_experts} + + # 每个 rank 的 R 个槽位已由全卡冗余专家占据。此时剩余槽位总数为 + # logical_expert_count,而未分配专家只有 logical_expert_count-R 个, + # 所以还需在非冗余专家中增加 R 个副本。 + for _ in range(self.num_redundant_experts_per_rank): + expert = min( + (expert for expert in remaining_experts if replica_counts[expert] < self.world_size), + key=lambda expert: ( + -logical_load[expert] / replica_counts[expert], + expert, + ), + ) + replica_counts[expert] += 1 + + # 分配器按专家组工作,而不是把同一专家拆成多个独立元素。这样分配 + # 一组副本时可以暂时取出多个不同 rank,从结构上避免本地重复专家。 + expert_groups = [] + for expert in remaining_experts: + replica_count = replica_counts[expert] + load_per_replica = logical_load[expert] / replica_count + aligned_load_per_replica = ceil(load_per_replica / self.expert_alignment) * self.expert_alignment + expert_groups.append((expert, replica_count, aligned_load_per_replica)) + + # 多副本专家需要在第一阶段先完成平铺,因此排在单副本专家之前。 + # 副本数相同时优先处理单副本负载较高的专家,最后用专家 ID 打破平局。 + expert_groups.sort(key=lambda group: (-group[1], -group[2], group[0])) + return expert_groups + + def _distribute_remaining_experts( + self, + redundant_experts: List[int], + expert_groups: List[ExpertReplicaGroup], + ) -> LayerPlacement: + """先平铺多副本专家,再按当前 rank 负载分配单副本专家。""" + placement = [list(redundant_experts) for _ in range(self.world_size)] + + total_replica_count = sum(replica_count for _, replica_count, _ in expert_groups) + assert total_replica_count % self.world_size == 0 + remaining_slots_per_rank = total_replica_count // self.world_size + remaining_slots = [remaining_slots_per_rank] * self.world_size + rank_load = [0.0] * self.world_size + + replicated_expert_groups = [group for group in expert_groups if group[1] > 1] + single_expert_groups = [group for group in expert_groups if group[1] == 1] + + # 阶段 1:用同一个循环游标依次平铺所有多副本专家。每个专家最多有 + # world_size 个副本,所以一组副本在游标绕回起点之前已经分配完毕, + # 同一 rank 不会出现该专家的两个副本。连续使用同一个游标还会让 + # 各 rank 在第一阶段获得的槽位数最多相差一个,不会提前填满某个 rank。 + next_rank = 0 + for expert, replica_count, load_per_replica in replicated_expert_groups: + for _ in range(replica_count): + assert remaining_slots[next_rank] > 0 + placement[next_rank].append(expert) + remaining_slots[next_rank] -= 1 + rank_load[next_rank] += load_per_replica + next_rank = (next_rank + 1) % self.world_size + + # 阶段 2:多副本专家的位置固定后,再把尚有空位的 rank 按当前负载 + # 放入最小堆。这里负载是第一优先级,剩余槽位数不再参与排序;每次 + # 都把当前最热的单副本专家交给最轻的未满 rank。 + # 全卡冗余专家对每个 rank 的贡献相同,因此无需计入 rank_load。 + rank_queue = [(rank_load[rank], rank) for rank in range(self.world_size) if remaining_slots[rank] > 0] + heapq.heapify(rank_queue) + + for expert, replica_count, load_per_replica in single_expert_groups: + assert replica_count == 1 + assert rank_queue, "not enough rank slots to place remaining experts" + + current_load, rank = heapq.heappop(rank_queue) + placement[rank].append(expert) + remaining_slots[rank] -= 1 + + # remaining_slots == 0 表示该 rank 已经刚好填满,不再放回队列。 + if remaining_slots[rank] > 0: + heapq.heappush(rank_queue, (current_load + load_per_replica, rank)) + + assert not rank_queue + assert all(slots == 0 for slots in remaining_slots) + + return placement + + def _reuse_current_slots( + self, + candidate_placement: LayerPlacement, + current_placement: LayerPlacement, + ) -> LayerPlacement: + """贪心匹配候选 rank,并让共同专家尽量复用当前物理槽位。""" + # 步骤 1:准备所有尚未匹配的候选行。 + # + # 候选布局中的 rank 编号只是负载规划阶段产生的临时编号。任意交换 + # 两个候选行都不会改变每行的专家组合和整体负载,因此可以重新排列 + # 候选行,使其尽量贴近当前运行布局。这里同时缓存专家集合,后续可以 + # 直接用集合交集计算两个 rank 之间的相似度。 + unmatched_candidates = [ + (candidate_rank, candidate_experts, set(candidate_experts)) + for candidate_rank, candidate_experts in enumerate(candidate_placement) + ] + placement = [] + + # 步骤 2:按当前 rank 0 -> N-1 的顺序贪心匹配候选行。 + # + # 相似度定义为两个 rank 共同持有的专家数量。共同专家越多,需要跨 + # rank 传输的专家权重就越少。一个候选行被选中后立即从待选列表移除, + # 从而建立当前 rank 和候选行之间的一一对应关系。 + for current_experts in current_placement: + current_expert_set = set(current_experts) + + # max() 首先选择共同专家数量最多的候选行。相似度相同时,负的 + # candidate_rank 让原候选 rank ID 更小的行优先,保证结果确定。 + best_candidate_index = max( + range(len(unmatched_candidates)), + key=lambda index: ( + len(current_expert_set & unmatched_candidates[index][2]), + -unmatched_candidates[index][0], + ), + ) + _, selected_experts, selected_expert_set = unmatched_candidates.pop(best_candidate_index) + + # 步骤 3:在已经匹配的 rank 内复用当前物理槽位。 + # + # new_experts 只包含候选行新引入的专家,并保持候选行中的原始顺序。 + # 它的数量必然等于当前行中需要被替换的专家数量。 + new_experts = [expert for expert in selected_experts if expert not in current_expert_set] + new_expert_index = 0 + rank_placement = [] + + # 依次检查当前物理槽位:如果槽位中的专家仍被候选行选中,就原地 + # 保留;否则用下一个新专家填充。这样只有真正变化的槽位需要搬运 + # 权重,共同专家不会因为候选行内部顺序不同而发生无意义移动。 + for current_expert in current_experts: + if current_expert in selected_expert_set: + rank_placement.append(current_expert) + continue + + rank_placement.append(new_experts[new_expert_index]) + new_expert_index += 1 + + # 所有需要替换的槽位都应恰好消费一个新专家。 + assert new_expert_index == len(new_experts) + placement.append(rank_placement) + + # 每个当前 rank 都必须匹配且只匹配一个候选行。 + assert not unmatched_candidates + return placement + + def _validate_inputs( + self, + logical_expert_load: List[List[float]], + placement: ExpertPlacement, + ) -> None: + """拒绝会导致逐层规划静默截断的输入。""" + if not logical_expert_load: + raise ValueError("logical_expert_load must contain at least one layer") + if len(placement) != len(logical_expert_load): + raise ValueError("load and placement must have the same number of layers") diff --git a/lightllm/server/router/model_infer/mode_backend/eplb/placement/initial.py b/lightllm/server/router/model_infer/mode_backend/eplb/placement/initial.py new file mode 100644 index 0000000000..626b7caa15 --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/eplb/placement/initial.py @@ -0,0 +1,39 @@ +"""Construct the deterministic expert placement used during model loading.""" + +from .types import LayerPlacement + + +def build_initial_local_expert_ids( + num_logical_experts: int, + num_ranks: int, + num_redundant_experts_per_rank: int, +) -> LayerPlacement: + """构建每个 rank 初始持有的完整 logical expert ID 列表。 + + 每个 rank 先持有连续划分得到的主专家,再按 rank 顺序选择不属于 + 本 rank 的专家作为默认冗余副本。这里仅负责生成 Python 列表;调用方如果要参与 tensor + 运算,需要自行转换为 ``torch.Tensor``。 + + 例如 ``num_logical_experts=8``、``num_ranks=4``、每个 rank 有 2 个 + 额外槽时,每个 rank 分到 2 个主专家,结果为: + + ``[[0, 1, 2, 3], [2, 3, 4, 5], [4, 5, 6, 7], [6, 7, 0, 1]]`` + + 其中每行前两个值是主专家,后两个值是已在初始加载阶段就可用的冗余副本。 + """ + assert num_logical_experts % num_ranks == 0 + num_experts_per_rank = num_logical_experts // num_ranks + assert 0 <= num_redundant_experts_per_rank <= num_logical_experts - num_experts_per_rank + + local_expert_ids_by_rank = [] + for rank in range(num_ranks): + first_expert_id = rank * num_experts_per_rank + local_expert_ids = list(range(first_expert_id, first_expert_id + num_experts_per_rank)) + first_redundant_expert_id = ((rank + 1) * num_experts_per_rank) % num_logical_experts + local_expert_ids.extend( + (first_redundant_expert_id + offset) % num_logical_experts + for offset in range(num_redundant_experts_per_rank) + ) + local_expert_ids_by_rank.append(local_expert_ids) + + return local_expert_ids_by_rank diff --git a/lightllm/server/router/model_infer/mode_backend/eplb/placement/planner.py b/lightllm/server/router/model_infer/mode_backend/eplb/placement/planner.py new file mode 100644 index 0000000000..2b9e66da38 --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/eplb/placement/planner.py @@ -0,0 +1,25 @@ +"""Abstract interface implemented by EPLB placement planners.""" + +from abc import ABC, abstractmethod + +import torch + +from .types import ExpertPlacement + + +class EPLBPlanner(ABC): + """根据逻辑专家负载生成完整物理布局。""" + + @abstractmethod + def plan( + self, + *, + logical_expert_load_samples: torch.Tensor, + current_placement: ExpertPlacement, + ) -> ExpertPlacement: + """根据 CPU 负载生成 ``[layer][rank][local physical expert]`` 布局。 + + ``logical_expert_load_samples`` 的 shape 为 + ``[rank, layer, sample, logical_expert]``。具体 planner 决定如何聚合 + rank 和 sample 维度。 + """ diff --git a/lightllm/server/router/model_infer/mode_backend/eplb/placement/routing.py b/lightllm/server/router/model_infer/mode_backend/eplb/placement/routing.py new file mode 100644 index 0000000000..9fcb69c679 --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/eplb/placement/routing.py @@ -0,0 +1,171 @@ +"""Build compact logical-to-physical routing metadata for one MoE layer. + +The input layout uses ``[rank][local physical expert] -> logical expert``. +Initial placement construction lives in :mod:`.initial`; this module only +inverts an existing placement into the fixed-width rows consumed by the +EPLB routing kernel. +""" + +from .types import LayerPlacement, LogicalToPhysicalMap + + +def build_logical_to_physical_map( + rank_to_logic_expert_ids: LayerPlacement, + num_logical_experts: int, + current_rank: int, + node_world_size: int, +) -> LogicalToPhysicalMap: + """使用普通 CPU list 构建单层 logical 到 physical expert 的路由表。 + + ``rank_to_logic_expert_ids`` 的 shape 为 + ``[num_ranks, num_physical_experts_per_rank]``,每行包含该 rank 的全部 + 物理专家。 + + 返回值的 shape 为 ``[num_logical_experts, 3 + routing_slots]``。每一行的 + 可视化结构如下: + + ``[global_count, node_count, current_gpu_count, physical_ids..., -1 padding...]`` + + * ``global_count``:所有 rank 上的有效副本总数; + * ``node_count``:当前节点上的有效副本数,包含本卡副本; + * ``current_gpu_count``:当前 GPU 上的有效副本数; + * ``physical_ids``:依次按本卡、本节点其他卡、其他节点排列的副本 ID。 + + 路由时优先使用最靠近当前 GPU 的非空候选集合:先使用本卡副本,其次使用 + 本节点副本,当前节点没有副本时才使用所有 rank 的副本。有效副本之后未 + 使用的固定宽度槽位填充为 ``-1``。 + + 本函数只负责 CPU 元数据计算。调用方需要设备 Tensor 时,应在函数外 + 显式执行 ``torch.tensor(...)``。 + """ + # 阶段 1:校验输入布局,并根据每个 rank 的物理槽位数计算冗余容量。 + num_ranks = len(rank_to_logic_expert_ids) + assert num_ranks > 0 + assert num_logical_experts % num_ranks == 0 + num_physical_experts_per_rank = len(rank_to_logic_expert_ids[0]) + assert all(len(rank_expert_ids) == num_physical_experts_per_rank for rank_expert_ids in rank_to_logic_expert_ids) + num_primary_experts_per_rank = num_logical_experts // num_ranks + num_redundant_experts_per_rank = num_physical_experts_per_rank - num_primary_experts_per_rank + assert num_redundant_experts_per_rank >= 0 + # 阶段 2:使用整个 world 的物理槽位总数作为固定路由槽宽度。实际候选 + # 仍只写入有效副本,其余槽位统一 padding 为 -1。 + num_routing_slots = num_ranks * num_physical_experts_per_rank + assert 0 <= current_rank < num_ranks + assert 0 < node_world_size <= num_ranks + assert num_ranks % node_world_size == 0 + + # 阶段 3:把“物理槽 -> logical expert”的完整布局反转为 + # “logical expert -> 全部物理槽”,得到每个专家的候选副本列表。 + physical_ids_by_logical_expert = _collect_physical_ids_by_logical_expert( + rank_to_logic_expert_ids, + num_logical_experts, + ) + + # 阶段 4:对每个候选列表做稳定排序。当前 rank 的 physical ID 排在最前, + # 同节点其他 rank 次之,跨节点副本最后。 + _sort_physical_ids_by_locality( + physical_ids_by_logical_expert, + current_rank, + num_physical_experts_per_rank, + node_world_size, + ) + # 阶段 5:逐个 logical expert 打包固定宽度的路由行。实际副本不足固定 + # 宽度时,剩余槽位使用 -1 padding;kernel 只会索引有效副本范围。 + logical_to_physical_map = [] + current_node = current_rank // node_world_size + current_node_rank_start = current_node * node_world_size + current_node_rank_end = current_node_rank_start + node_world_size + for physical_expert_ids in physical_ids_by_logical_expert: + replica_ranks = [ + physical_expert_id // num_physical_experts_per_rank for physical_expert_id in physical_expert_ids + ] + num_node_replicas = sum( + current_node_rank_start <= replica_rank < current_node_rank_end for replica_rank in replica_ranks + ) + num_current_gpu_replicas = sum(replica_rank == current_rank for replica_rank in replica_ranks) + logical_to_physical_map.append( + _build_routing_row( + physical_expert_ids=physical_expert_ids, + num_node_replicas=num_node_replicas, + num_current_gpu_replicas=num_current_gpu_replicas, + num_routing_slots=num_routing_slots, + ) + ) + return logical_to_physical_map + + +def _collect_physical_ids_by_logical_expert( + rank_to_logic_expert_ids: LayerPlacement, + num_logical_experts: int, +) -> list[list[int]]: + """将完整物理布局反转为每个 logical expert 对应的物理槽位。 + + 例如输入 ``[[0, 1, 1], [2, 3, 0]]``,先按 rank 顺序拼成 + ``[0, 1, 1, 2, 3, 0]``。该列表的下标就是 physical expert ID,值就是 + logical expert ID,因此最终返回 ``[[0, 5], [1, 2], [3], [4]]``。 + """ + logical_expert_ids_by_physical_id = [ + logical_expert_id for rank_expert_ids in rank_to_logic_expert_ids for logical_expert_id in rank_expert_ids + ] + physical_ids_by_logical_expert = [[] for _ in range(num_logical_experts)] + for physical_expert_id, logical_expert_id in enumerate(logical_expert_ids_by_physical_id): + assert 0 <= logical_expert_id < num_logical_experts + physical_ids_by_logical_expert[logical_expert_id].append(physical_expert_id) + + return physical_ids_by_logical_expert + + +def _sort_physical_ids_by_locality( + physical_ids_by_logical_expert: list[list[int]], + current_rank: int, + num_physical_experts_per_rank: int, + node_world_size: int, +) -> None: + """按当前 rank、当前节点、其他节点的优先级稳定排序副本。 + + 当前 rank 的排序键为 0,同节点其他 rank 为 1,其他节点为 2。同一优先级 + 内保持原 physical ID 顺序不变。 + """ + current_node = current_rank // node_world_size + + def locality_priority(physical_expert_id: int) -> int: + physical_rank = physical_expert_id // num_physical_experts_per_rank + if physical_rank == current_rank: + return 0 + if physical_rank // node_world_size == current_node: + return 1 + return 2 + + for physical_expert_ids in physical_ids_by_logical_expert: + # list.sort 是稳定排序:排序键相同时,physical ID 的原始顺序不变。 + physical_expert_ids.sort(key=locality_priority) + + +def _build_routing_row( + physical_expert_ids: list[int], + num_node_replicas: int, + num_current_gpu_replicas: int, + num_routing_slots: int, +) -> list[int]: + """将一个 logical expert 的候选 physical IDs 打包为固定宽度路由行。 + + ``physical_expert_ids`` 已由调用方完成拓扑优先的稳定排序,所以本函数 + 不再依赖 ``current_rank``。列表长度就是该 logical expert 的有效物理 + 副本数,无需额外传入容易失配的副本数量。 + """ + # 阶段 1:候选列表包含该专家的全部物理副本,其长度就是有效副本数。 + num_global_replicas = len(physical_expert_ids) + assert 0 < num_global_replicas <= num_routing_slots + assert 0 <= num_current_gpu_replicas <= num_node_replicas <= num_global_replicas + + # 阶段 2:有效槽位直接保存稳定排序后的候选;固定宽度中未使用的尾部 + # 槽位统一填充 -1。kernel 的副本索引严格小于 num_global_replicas, + # 因而不会读取 padding。 + num_padding_slots = num_routing_slots - num_global_replicas + routing_slots = physical_expert_ids + [-1] * num_padding_slots + + # 阶段 3:将三层有效副本计数放在固定头部,后面拼接按拓扑优先级排序的 + # physical IDs 和 -1 padding: + # + # [global_count, node_count, current_gpu_count, physical_ids..., -1 padding...] + return [num_global_replicas, num_node_replicas, num_current_gpu_replicas, *routing_slots] diff --git a/lightllm/server/router/model_infer/mode_backend/eplb/placement/types.py b/lightllm/server/router/model_infer/mode_backend/eplb/placement/types.py new file mode 100644 index 0000000000..f6534ff85f --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/eplb/placement/types.py @@ -0,0 +1,13 @@ +"""Shared type aliases for EPLB placement planning.""" + +from typing import List, Tuple + + +# [rank][local physical expert] -> logical expert +LayerPlacement = List[List[int]] +# [layer][rank][local physical expert] -> logical expert +ExpertPlacement = List[LayerPlacement] +# [logical expert][replica metadata] +LogicalToPhysicalMap = List[List[int]] +# (logical expert, replica count, aligned load per replica) +ExpertReplicaGroup = Tuple[int, int, float] diff --git a/lightllm/server/router/model_infer/mode_backend/eplb/runtime_manager.py b/lightllm/server/router/model_infer/mode_backend/eplb/runtime_manager.py new file mode 100644 index 0000000000..6ff08ebb0f --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/eplb/runtime_manager.py @@ -0,0 +1,610 @@ +from enum import Enum +import time +from typing import List, Optional + +import torch +import torch.distributed as dist + +from lightllm.common.basemodel.basemodel import TpPartBaseModel +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.fused_moe_weight import ( + FusedMoeWeight, +) +from lightllm.server.metrics.manager import MetricClient +from lightllm.utils.dist_utils import ( + get_global_rank, + get_global_world_size, + get_node_world_size, +) +from lightllm.utils.device_utils import is_sm100_gpu +from lightllm.utils.envs_utils import get_eplb_step_interval +from lightllm.utils.log_utils import init_logger +from lightllm.utils.shm_port_args import get_shm_port_args + +from . import metrics as eplb_metrics +from .async_expert_transfer import ( + EPLBTransferInfo, + PinnedMemoryEPLBTransfer, +) +from .async_load_gather_task import EPLBLoadGatherTask +from .async_placement_plan_task import EPLBPlanTask +from .async_transfer_planner import EPLBTransferPlanner +from .placement import ( + EPLBPlanner, + ExpertPlacement, + build_logical_to_physical_map, + create_eplb_planner, + save_placement_config, +) + +logger = init_logger(__name__) +EPLB_EXPERT_ALIGNMENT = 128 +EPLB_MIN_AVERAGE_TOKENS_PER_EXPERT = 128 + + +class EPLBManagerState(Enum): + """EPLB 管理器在一次负载均衡循环中的阶段。""" + + COLLECTING = "collecting" + EVALUATING = "evaluating" + WAIT_LOAD_GATHER_FINISHED = "wait_load_gather_finished" + PLAN_PLACEMENT = "plan_placement" + WAIT_PLAN_PLACEMENT_FINISHED = "wait_plan_placement_finished" + PLAN_TRANSFER = "plan_transfer" + WAIT_PLAN_TRANSFER_FINISHED = "wait_plan_transfer_finished" + TRANSFERRING = "transferring" + + +class EPLBManager: + """由 :meth:`step` 驱动的 EPLB 状态机。 + + 每次调用 :meth:`step` 最多处理一个状态。主路径及各状态的职责如下:: + + [COLLECTING] + prefill kernel 将各层 logical expert 负载写入 24 行环形样本; + manager 只记录采样 step,等待下一个评估周期。 + | + | 评估周期到达 + v + [EVALUATING] + 将本地环形样本完整复制到 CPU 并上报负载指标;在独立 Gloo + 通信组中启动逐样本负载的后台 all-gather。rank 0 在通信前判断 + 本地样本量;样本不足或达到重排次数上限时不启动通信。 + | + v + [WAIT_LOAD_GATHER_FINISHED] + 等待所有 rank 完成负载汇集,保留 planner 需要的 + rank 和 sample 原始维度。 + | + | 样本充足且仍允许重排 + v + [PLAN_PLACEMENT] + rank 0 使用完整的全局专家负载启动后台布局规划任务。 + | + v + [WAIT_PLAN_PLACEMENT_FINISHED] + 轮询 rank 0 的规划任务,并向所有 rank 广播目标布局。 + | + | 目标布局发生变化 + v + [PLAN_TRANSFER] + 每个 rank 根据相同的当前/目标布局启动后台传输规划任务。 + | + v + [WAIT_PLAN_TRANSFER_FINISHED] + 等待所有 rank 生成一致的、按依赖关系分批的传输任务。 + | + v + [TRANSFERRING] + 分批启动并轮询后台权重传输;整批完成后,主推理线程在安全 + 边界统一提交权重和路由 metadata。全部批次完成后发布新布局、 + 清空 prefill 路由样本并回到 COLLECTING。 + + 以下分支会提前回到 ``COLLECTING``:: + + EVALUATING + |-- 平均 token 数不足 --------> 保留环形窗口,继续滚动采样 + `-- 已达到重排次数上限 ------> 清空样本,仅周期性上报指标 + + WAIT_PLAN_PLACEMENT_FINISHED + `-- 目标布局与当前布局相同 ---> 保留环形窗口,等待下次评估 + + 布局规划、传输规划和权重传输在后台执行;主推理线程只负责创建任务、 + 轮询状态,以及在安全边界提交已经完成的结果。 + """ + + def __init__( + self, + model: TpPartBaseModel, + max_rebalance_count: int = 1, + config_path: Optional[str] = None, + plan_mode: str = "greedy", + ) -> None: + # SM100 FP4 Mega-MoE 会将在线专家权重转换为独立的 kernel 布局,并使用源 tensor 的 data_ptr + # 作为 key 缓存这些转换后的副本。EPLB 通过原地 copy_ 替换专家行,只改变权重内容而不会改变 + # data_ptr,因此重平衡后 Mega-MoE 仍会读取旧的转换权重。在 EPLB 能够失效或更新该缓存前, + # 暂不支持 SM100。 + assert not is_sm100_gpu(), "EPLB does not support SM100" + + weights: List[FusedMoeWeight] = _find_fused_moe_weights(model) + assert weights, "EPLB requires at least one EP MoE layer" + assert max_rebalance_count >= -1 + + # 模型与专家拓扑:初始化后保持不变。 + self._weights: List[FusedMoeWeight] = weights + self.config_path = config_path + self.global_rank: int = get_global_rank() + self.world_size: int = get_global_world_size() + assert self.world_size > 1, "EPLB requires more than one rank" + self.node_world_size: int = get_node_world_size() + self.layer_indexes = [weight.layer_num_ for weight in weights] + self._eplb_impls = [weight.fuse_moe_impl for weight in weights] + + first_impl = self._eplb_impls[0] + self.num_logical_experts: int = first_impl.n_routed_experts + self.num_redundant_experts_per_rank: int = first_impl.num_redundant_experts_per_rank + self.plan_mode: str = plan_mode + self.planner: EPLBPlanner = create_eplb_planner( + self.plan_mode, + self.world_size, + self.num_redundant_experts_per_rank, + expert_alignment=EPLB_EXPERT_ALIGNMENT, + ) + + # 评估调度:steps 只在 COLLECTING 状态递增。prefill 路由样本从当前 + # 布局生效时开始写入,并在固定容量内保留最近的采样窗口。 + self.step_interval: int = get_eplb_step_interval() + self.steps: int = 0 + self.max_rebalance_count: int = max_rebalance_count + self.completed_rebalance_count: int = 0 + + # 一次重排周期中的短期状态。统一初始化为 None,避免各状态 + # 通过 hasattr()/del 隐式定义 EPLBManager 的属性结构。 + self._load_gather_task: Optional[EPLBLoadGatherTask] = None + self._pending_plan_load_samples: Optional[torch.Tensor] = None + self._metric_sample_load: Optional[torch.Tensor] = None + self._plan_task: Optional[EPLBPlanTask] = None + self.target_placement: Optional[ExpertPlacement] = None + self._transfer_planner: Optional[EPLBTransferPlanner] = None + self.pending_transfer_batches: Optional[List[List[EPLBTransferInfo]]] = None + self.rebalance_started_at: Optional[float] = None + self.active_transfer_batch: Optional[List[EPLBTransferInfo]] = None + self.active_transfers: Optional[List[PinnedMemoryEPLBTransfer]] = None + + # 分布式通信:后台负载汇集、主线程控制面和后台权重传输分别使用 + # 独立的 Gloo 通信组。负载 all-gather 可能跨越多个 manager step, + # 不能与主线程中按 step 排序的控制 collective 共用同一个 group。 + self.load_gather_group = dist.new_group(list(range(self.world_size)), backend="gloo") + self.control_group = dist.new_group(list(range(self.world_size)), backend="gloo") + self.transfer_group = dist.new_group(list(range(self.world_size)), backend="gloo") + + # 每层布局都保存完整的本地专家列表;完成初始化后,所有物理槽位 + # 都可以由 EPLB 重新分配,不再区分固定主专家槽和冗余专家槽。 + # 本 rank 的布局索引为 [layer][local_expert]。 + local_expert_ids_by_layer = [list(impl.local_logics_expert_ids_list) for impl in self._eplb_impls] + + # all_gather 后的布局索引为 [rank][layer][local_expert]。 + expert_ids_by_rank_and_layer: List[List[List[int]]] = [[] for _ in range(self.world_size)] + dist.all_gather_object( + expert_ids_by_rank_and_layer, + local_expert_ids_by_layer, + group=self.control_group, + ) + + # 转置为全局统一使用的 [layer][rank][local_expert]。 + self.current_placement: ExpertPlacement = [ + [expert_ids_by_rank_and_layer[rank][layer_index] for rank in range(self.world_size)] + for layer_index in range(len(weights)) + ] + + self.state = EPLBManagerState.COLLECTING + self.next_evaluation_step = self.step_interval + self._clear_prefill_route_samples() + + if self.global_rank == 0: + self.metric_client: MetricClient = MetricClient(get_shm_port_args().metric_port) + logger.info( + f"eplb enabled layers={len(weights)} num_logical_experts={self.num_logical_experts} " + f"num_redundant_experts_per_rank={self.num_redundant_experts_per_rank} " + f"step_interval={self.step_interval} max_rebalance_count={self.max_rebalance_count} " + f"plan_mode={self.plan_mode} planner={type(self.planner).__name__}" + ) + + def step(self) -> None: + """在一个安全的推理边界推进一次状态机。""" + if self.state is EPLBManagerState.COLLECTING: + self._step_collecting() + return + + if self.state is EPLBManagerState.EVALUATING: + self._step_evaluating() + return + + if self.state is EPLBManagerState.WAIT_LOAD_GATHER_FINISHED: + self._step_wait_load_gather_finished() + return + + if self.state is EPLBManagerState.PLAN_PLACEMENT: + self._step_plan_placement() + return + + if self.state is EPLBManagerState.WAIT_PLAN_PLACEMENT_FINISHED: + self._step_wait_plan_placement_finished() + return + + if self.state is EPLBManagerState.PLAN_TRANSFER: + self._step_plan_transfer() + return + + if self.state is EPLBManagerState.WAIT_PLAN_TRANSFER_FINISHED: + self._step_wait_plan_transfer_finished() + return + + if self.state is EPLBManagerState.TRANSFERRING: + self._step_transferring() + return + + raise RuntimeError(f"unknown EPLB manager state: {self.state!r}") + + # 状态处理:与 step() 的分发顺序保持一致。 + + def _step_collecting(self) -> None: + """记录一个采样步,并在当前评估周期结束后进入评估状态。""" + self.steps += 1 + if self.steps < self.next_evaluation_step: + return + else: + self.next_evaluation_step += self.step_interval + self.state = EPLBManagerState.EVALUATING + + def _step_evaluating(self) -> None: + """快照本地逐样本负载,并在独立通信组中启动后台汇集。""" + counters = [impl.prefill_route_counter for impl in self._eplb_impls] + if any(counter.ndim != 2 or counter.shape[1] != self.num_logical_experts for counter in counters): + raise RuntimeError("EPLB prefill route counter shape must be [sample_capacity, num_logical_experts]") + if len({counter.shape[0] for counter in counters}) != 1: + raise RuntimeError("EPLB prefill route counter capacities must match across layers") + + # 一次性堆叠各层环形样本并复制到 CPU,保留完整的 + # [layer, sample, logical_expert] 维度。后续后台 all-gather 会继续保留 + # rank 和 sample 维;通信完成后将原始四维快照交给 planner。 + # 此处先不清零 GPU 样本:如果样本不足或无需迁移,下一周期会继续 + # 滚动覆盖最旧行;达到重排上限或成功切换到新布局后才重置窗口。本轮 + # 异步规划使用独立的 CPU 快照,不会与后续的 atomic add 竞争。 + local_load_samples = torch.stack(counters).detach().cpu() + if self.global_rank == 0: + local_load = local_load_samples.sum(dim=1) + eplb_metrics.publish_expert_load_metrics( + metric_client=self.metric_client, + expert_load=local_load, + ) + + # 达到重排次数上限后仍保留周期性负载上报,但不再执行后续的跨 rank + # 通信和布局规划。清空本轮样本,使下一次指标对应新的采样窗口。 + reached_rebalance_limit = ( + self.max_rebalance_count != -1 and self.completed_rebalance_count >= self.max_rebalance_count + ) + if reached_rebalance_limit: + self._clear_prefill_route_samples() + self.state = EPLBManagerState.COLLECTING + else: + # rank 0 的路由分布足以代表全局分布,因此只使用 rank 0 的本地 + # 样本判断统计量是否充足,再广播布尔决策以保持所有 rank 的状态 + # 转移一致。样本不足时不启动大块原始负载 all-gather。 + has_enough_load = None + if self.global_rank == 0: + average_tokens_per_expert = local_load.sum().item() / local_load.numel() + has_enough_load = average_tokens_per_expert >= EPLB_MIN_AVERAGE_TOKENS_PER_EXPERT + if not has_enough_load: + logger.info( + "eplb continue collecting average_tokens_per_expert=%.2f threshold=%s", + average_tokens_per_expert, + EPLB_MIN_AVERAGE_TOKENS_PER_EXPERT, + ) + + values = [has_enough_load] + dist.broadcast_object_list( + values, + src=0, + group=self.control_group, + ) + has_enough_load = values[0] + assert has_enough_load is not None + + if has_enough_load: + self._load_gather_task = EPLBLoadGatherTask( + local_load_samples=local_load_samples, + load_gather_group=self.load_gather_group, + ) + self._load_gather_task.start() + self.state = EPLBManagerState.WAIT_LOAD_GATHER_FINISHED + else: + self.state = EPLBManagerState.COLLECTING + + def _step_wait_load_gather_finished(self) -> None: + """等待原始负载汇集完成,并准备 planner 和 metrics 输入。""" + assert self._load_gather_task is not None + if not self._all_ranks_finished(self._load_gather_task.is_finished()): + return + + gathered_load_samples = self._load_gather_task.result + assert gathered_load_samples is not None + assert gathered_load_samples.ndim == 4 + assert gathered_load_samples.shape[0] == self.world_size + assert gathered_load_samples.shape[1] == len(self._eplb_impls) + assert gathered_load_samples.shape[3] == self.num_logical_experts + self._load_gather_task = None + + # gathered_load_samples 保留 [rank, layer, sample, logical_expert] + # 原始结构并 + # 直接交给 planner。重排计算指标应表示一次真实 prefill 的 + # 关键路径开销,而不是多个不同批次累加后的虚拟大批次。 + # EP rank 以相同顺序执行 prefill,每次 dispatch 都将同一 sample + # index 推进一次;因此各 rank 的第 0 行属于同一采样位置。 + # metrics 固定取该行,只汇总其在各 rank 上的分片,得到 + # [layer, logical_expert]。 + if self.global_rank == 0: + self._pending_plan_load_samples = gathered_load_samples + metric_sample_load_by_rank = gathered_load_samples[:, :, 0, :] + self._metric_sample_load = metric_sample_load_by_rank.sum(dim=0) + self.state = EPLBManagerState.PLAN_PLACEMENT + + def _step_plan_placement(self) -> None: + """由 rank 0 使用已汇集的全局负载启动异步规划。""" + self.state = EPLBManagerState.WAIT_PLAN_PLACEMENT_FINISHED + if self.global_rank == 0: + assert self._pending_plan_load_samples is not None + # planner 消费保留 rank/sample 维的原始快照;指标使用其中 + # 一个 sample 行,确保 before/after 比较的是同一批负载。 + self._plan_task = EPLBPlanTask( + planner=self.planner, + logical_expert_load_samples=self._pending_plan_load_samples, + current_placement=self.current_placement, + ) + self._plan_task.start() + # 任务对象已持有 Tensor,manager 不再保留重复引用。 + self._pending_plan_load_samples = None + + def _step_wait_plan_placement_finished(self) -> None: + """等待 rank 0 完成规划并广播目标专家排布。""" + placement: Optional[ExpertPlacement] = None + if self.global_rank == 0: + assert self._plan_task is not None + if self._plan_task.is_finished(): + placement = self._plan_task.result + assert placement is not None + + values = [placement] + dist.broadcast_object_list(values, src=0, group=self.control_group) + placement = values[0] + if placement is None: + return + + if self.global_rank == 0: + assert self._metric_sample_load is not None + eplb_metrics.publish_rebalance_compute_metrics( + metric_client=self.metric_client, + sample_load=self._metric_sample_load, + current_placement=self.current_placement, + target_placement=placement, + expert_alignment=EPLB_EXPERT_ALIGNMENT, + ) + self._metric_sample_load = None + self._plan_task = None + + if placement == self.current_placement: + if self.global_rank == 0: + logger.info("eplb skip rearrangement because placement is unchanged") + self.state = EPLBManagerState.COLLECTING + return + + # 广播得到的 placement 已经是规划器新建的完整布局,没有外部持有者会 + # 再修改它,因此可直接保存,不需要逐层深拷贝。 + self.target_placement = placement + self.state = EPLBManagerState.PLAN_TRANSFER + + def _step_plan_transfer(self) -> None: + """启动异步传输规划,再进入完成状态轮询阶段。""" + # 所有 rank 使用相同的 current/target placement 独立生成确定性的传输 + # 批次,避免广播体积较大的任务列表;耗时的逐层依赖分析放到后台线程, + # 当前推理线程从下一次安全边界开始只需轮询完成状态。 + assert self.target_placement is not None + self._transfer_planner = EPLBTransferPlanner( + current_placement=self.current_placement, + target_placement=self.target_placement, + num_logical_experts=self.num_logical_experts, + world_size=self.world_size, + ) + self._transfer_planner.start() + self.state = EPLBManagerState.WAIT_PLAN_TRANSFER_FINISHED + + def _step_wait_plan_transfer_finished(self) -> None: + """等待所有 rank 异步生成相同的传输批次,再统一进入传输状态。""" + assert self._transfer_planner is not None + + # 即使本 rank 已经完成,也必须等待其他 rank 的镜像任务列表就绪;否则 + # 提前进入 TRANSFERRING 的 rank 可能发起尚无对端参与的点对点传输。 + if not self._all_ranks_finished(self._transfer_planner.is_finished()): + return + + pending_transfer_batches = self._transfer_planner.result + assert pending_transfer_batches is not None + self.pending_transfer_batches = pending_transfer_batches + self._transfer_planner = None + if not self.pending_transfer_batches: + raise RuntimeError("planned EPLB rearrangement must contain at least one transfer") + + self.state = EPLBManagerState.TRANSFERRING + if self.global_rank == 0: + assert self.target_placement is not None + changed_layer_count = sum( + current != target for current, target in zip(self.current_placement, self.target_placement) + ) + logger.info( + "eplb started steps=%s changed_layer_count=%s changed_slot_count=%s", + self.steps, + changed_layer_count, + sum(len(transfer_batch) for transfer_batch in self.pending_transfer_batches), + ) + + def _step_transferring(self) -> None: + """启动或轮询一个传输批次;整批完成后再统一提交。""" + if self.rebalance_started_at is None: + self.rebalance_started_at = time.monotonic() + + # 没有活动批次时,所有 rank 根据相同的 pending 列表构造下一批任务。 + if self.active_transfer_batch is None: + assert self.pending_transfer_batches is not None + transfer_batch = self.pending_transfer_batches.pop(0) if self.pending_transfer_batches else [] + + # 空批次表示公共任务列表已经耗尽,所有 rank 可以同时结束重排。 + if not transfer_batch: + assert self.target_placement is not None + self.current_placement = self.target_placement + elapsed = time.monotonic() - self.rebalance_started_at + if self.global_rank == 0: + self._persist_current_placement() + self._clear_prefill_route_samples() + self.completed_rebalance_count += 1 + self.pending_transfer_batches = None + self.target_placement = None + self.rebalance_started_at = None + self.state = EPLBManagerState.COLLECTING + if self.global_rank == 0: + logger.info( + "eplb completed wall_time=%.2fs completed_rebalance_count=%s max_rebalance_count=%s", + elapsed, + self.completed_rebalance_count, + self.max_rebalance_count, + ) + else: + self.active_transfer_batch = transfer_batch + + # 普通批次允许多个 rank 不冲突的任务并行,但每个 rank 最多参与 + # 一条;覆盖环批次可能要求同一 rank 同时保存多个源/目标的 pinned + # row,必须等整批传输完成后再统一覆盖 live 权重。 + self.active_transfers = [ + PinnedMemoryEPLBTransfer( + weights=self._weights, + transfer_group=self.transfer_group, + current_global_rank=self.global_rank, + transfer_info=transfer_info, + ) + for transfer_info in transfer_batch + if self.global_rank in (transfer_info.source_rank, transfer_info.dest_rank) + ] + for transfer in self.active_transfers: + transfer.start() + else: + # 已有活动批次时,本 step 只负责轮询;整批完成后才统一提交。 + self._poll_transfer_batch() + + def _poll_transfer_batch(self) -> None: + """等待当前批次全部完成,随后统一提交并释放本地任务。""" + # 每个 rank 只负责自己参与的任务;不参与当前批次的 rank,其本地任务 + # 列表为空,all([]) 自然为 True。所有 rank 汇总一个布尔值即可判断整批 + # 是否完成,无需重复传输并逐条匹配 EPLBTransferInfo。 + assert self.active_transfers is not None + if not self._all_ranks_finished(all(transfer.is_finished() for transfer in self.active_transfers)): + return + + # 所有 rank 使用相同的批次顺序提交,因此全局 placement 和 metadata + # 始终一致;只有 destination rank 会额外写入实际专家权重。 + from lightllm.server.router.model_infer.infer_batch import g_infer_context + + # 专家权重和路由 metadata 都由 overlap stream 上的 MoE forward + # 读取。将整批写操作排到同一条 stream,便可自然等待此前的 forward, + # 并保证后续 forward 只能看到完整提交后的权重与 metadata。 + assert self.active_transfer_batch is not None + with torch.cuda.stream(g_infer_context.get_overlap_stream()): + for transfer_info in self.active_transfer_batch: + self._commit_transfer(transfer_info) + for layer_index in {transfer_info.layer_index for transfer_info in self.active_transfer_batch}: + self._publish_layer_metadata(layer_index) + + self.active_transfers = None + self.active_transfer_batch = None + + def _commit_transfer(self, transfer_info: EPLBTransferInfo) -> None: + """把一条已完成传输提交到 live 权重和完整布局。""" + is_destination_rank = transfer_info.dest_rank == self.global_rank + if is_destination_rank: + assert self.active_transfers is not None + active_transfer = next( + (transfer for transfer in self.active_transfers if transfer.transfer_info == transfer_info), + None, + ) + assert active_transfer is not None, "EPLB destination rank has no matching completed transfer" + for tensor_buffer in active_transfer.tensor_buffers: + tensor_buffer.live_tensor[transfer_info.dest_local_expert_index].copy_( + tensor_buffer.pinned_row, + non_blocking=True, + ) + + layer_index = transfer_info.layer_index + layer_impl = self._eplb_impls[layer_index] + self.current_placement[layer_index][transfer_info.dest_rank][ + transfer_info.dest_local_expert_index + ] = transfer_info.source_logical_expert_id + if is_destination_rank: + layer_impl.local_logics_expert_ids_list[ + transfer_info.dest_local_expert_index + ] = transfer_info.source_logical_expert_id + + def _all_ranks_finished(self, local_finished: bool) -> bool: + """通过控制通信组判断所有 rank 的当前后台任务是否完成。""" + finished_by_rank = [False] * self.world_size + dist.all_gather_object( + finished_by_rank, + local_finished, + group=self.control_group, + ) + return all(finished_by_rank) + + def _publish_layer_metadata(self, layer_index: int) -> None: + """在整批槽位更新完成后发布该层路由 metadata。""" + layer_impl = self._eplb_impls[layer_index] + logical_to_physical_map = torch.tensor( + build_logical_to_physical_map( + self.current_placement[layer_index], + self.num_logical_experts, + current_rank=self.global_rank, + node_world_size=self.node_world_size, + ), + dtype=torch.int32, + pin_memory=True, + ) + layer_impl.logical_to_physical_map.copy_(logical_to_physical_map, non_blocking=True) + + def _clear_prefill_route_samples(self) -> None: + """在 overlap stream 上清空所有层的 prefill 路由样本和设备端索引。""" + from lightllm.server.router.model_infer.infer_batch import g_infer_context + + # prefill route sample 由 forward 中的 Triton kernel 在 overlap stream 上更新。 + # 将 zero_ 排到同一条 stream,可保证它位于此前 forward 之后、下一次 + # forward 之前,无需额外 synchronize,也不会与 atomic add 并发。 + with torch.cuda.stream(g_infer_context.get_overlap_stream()): + for impl in self._eplb_impls: + impl.prefill_route_counter.zero_() + impl.prefill_route_sample_index.zero_() + + def _persist_current_placement(self) -> None: + """由 rank 0 将当前完整布局写回启动时指定的输入/输出文件。""" + if self.config_path is None: + return + save_placement_config( + self.config_path, + layer_indexes=self.layer_indexes, + placement=self.current_placement, + num_logical_experts=self.num_logical_experts, + world_size=self.world_size, + num_redundant_experts_per_rank=self.num_redundant_experts_per_rank, + ) + + +def _find_fused_moe_weights(model: TpPartBaseModel) -> List[FusedMoeWeight]: + weights: List[FusedMoeWeight] = [] + for layer in model.trans_layers_weight: + weight = getattr(layer, "experts", None) + if isinstance(weight, FusedMoeWeight) and weight.enable_ep_moe: + weights.append(weight) + return weights diff --git a/lightllm/server/router/model_infer/mode_backend/redundancy_expert_manager.py b/lightllm/server/router/model_infer/mode_backend/redundancy_expert_manager.py deleted file mode 100644 index 596eca4f24..0000000000 --- a/lightllm/server/router/model_infer/mode_backend/redundancy_expert_manager.py +++ /dev/null @@ -1,158 +0,0 @@ -# 对于 deepseekv3 模型在 ep 运行模式下,自动分析统计各个专家的出现频率,然后 -# 自动更新当前的冗余专家为新的冗余专家。 -import torch -import time -import enum -import lightllm.utils.petrel_helper as utils -import threading -import json -from typing import List -from lightllm.common.basemodel.basemodel import TpPartBaseModel -from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.ep_redundancy import ( - FusedMoeWeightEPAutoRedundancy, -) -from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.fused_moe_weight import FusedMoeWeight -from lightllm.utils.envs_utils import get_env_start_args, get_redundancy_expert_update_interval -from lightllm.utils.envs_utils import get_redundancy_expert_update_max_load_count -from lightllm.utils.envs_utils import get_redundancy_expert_num -from lightllm.utils.dist_utils import get_global_rank -from lightllm.common.basemodel.layer_weights.hf_load_utils import load_func -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -class RedundancyExpertManager: - def __init__(self, model: TpPartBaseModel): - self.args = get_env_start_args() - self.model = model - self.ep_fused_moeweights: List[FusedMoeWeightEPAutoRedundancy] = [] - for layer in self.model.trans_layers_weight: - ep_weights = self._find_members_of_class(layer, FusedMoeWeight) - assert len(ep_weights) <= 1 - self.ep_fused_moeweights.extend([FusedMoeWeightEPAutoRedundancy(e) for e in ep_weights]) - - # save load params - self.use_safetensors = True - files = utils.PetrelHelper.list(self.args.model_dir, extension="all") - candidate_files = list(filter(lambda x: x.endswith(".safetensors"), files)) - if len(candidate_files) == 0: - self.use_safetensors = False - candidate_files = list(filter(lambda x: x.endswith(".bin"), files)) - assert len(candidate_files) != 0, "can only support pytorch tensor and safetensors format for weights." - self.candidate_files = candidate_files - - # state 1. check_to_update 2. prepare_update 3. start_load_hf_weights 4. wait_load_ready, 5. commit - self.state: _STATE = _STATE.CHECK_TO_UPDATE - self.update_time = time.time() - self.update_interval = get_redundancy_expert_update_interval() - self.load_thread: threading.Thread = None - self.global_rank = get_global_rank() - # 冗余专家的最大加载次数 - self.load_count = 0 - self.max_load_count = get_redundancy_expert_update_max_load_count() - - # 清理counter - self._clear_all_counter() - - self.rank0_redundancy_expert_config = { - "redundancy_expert_num": get_redundancy_expert_num(), - "default": list(range(get_redundancy_expert_num())), - } - - def step(self): - if self.load_count >= self.max_load_count: - return - - if self.state == _STATE.CHECK_TO_UPDATE: - cur_time = time.time() - if cur_time - self.update_time > self.update_interval: - self.update_time = cur_time - self.state = _STATE.PREPARE_UPDATE - logger.info(f"global_rank {self.global_rank} state to prepare update") - elif self.state == _STATE.PREPARE_UPDATE: - self._prepare_load_new_redundancy_expert() - self.state = _STATE.START_LOAD_HF_WEIGHTS - logger.info(f"global_rank {self.global_rank} state to start load hf weights") - - elif self.state == _STATE.START_LOAD_HF_WEIGHTS: - self.load_thread = threading.Thread(target=self._load_hf_weights, daemon=True) - self.load_thread.start() - self.state = _STATE.WAIT_LOAD_READY - logger.info(f"global_rank {self.global_rank} state to wait load ready") - - elif self.state == _STATE.WAIT_LOAD_READY: - if not self.load_thread.is_alive(): - self.load_thread = None - self.state = _STATE.COMMIT - logger.info(f"global_rank {self.global_rank} state to commit") - - elif self.state == _STATE.COMMIT: - self._commit() - self.state = _STATE.CHECK_TO_UPDATE - self.load_count += 1 - logger.info(f"global_rank {self.global_rank} state to check to update") - return - - def _prepare_load_new_redundancy_expert(self): - for w in self.ep_fused_moeweights: - topk_redundancy_expert_ids = w.prepare_redundancy_experts() - if self.global_rank == 0: - self.rank0_redundancy_expert_config[str(w._ep_w.layer_num)] = topk_redundancy_expert_ids - - if self.global_rank == 0: - try: - with open("./redundancy_expert_config.json", "w") as f: - json.dump(self.rank0_redundancy_expert_config, f, indent=4) - logger.info( - f"rank {self.global_rank} save redundancy_expert_config.json to ./redundancy_expert_config.json" - ) - except BaseException as e: - logger.exception(str(e)) - logger.error(f"global rank {self.global_rank} save redundancy_expert_config.json failed") - - return - - def _load_hf_weights(self): - start = time.time() - try: - for file in self.candidate_files: - load_func( - file, - use_safetensors=self.use_safetensors, - pre_post_layer=None, - transformer_layer_list=self.ep_fused_moeweights, - weight_dir=self.args.model_dir, - ) - except BaseException as e: - logger.exception(str(e)) - raise e - cost_time = time.time() - start - logger.info(f"global rank {self.global_rank} load redundancy_expert cost time: {cost_time} s") - return - - def _commit(self): - for w in self.ep_fused_moeweights: - w.commit() - return - - def _find_members_of_class(self, obj, cls): - members = [] - for attr in dir(obj): - value = getattr(obj, attr) - if isinstance(value, cls): - members.append(value) - return members - - def _clear_all_counter(self): - for w in self.ep_fused_moeweights: - w.clear_counter() - return - - -class _STATE(enum.Enum): - CHECK_TO_UPDATE = 0 - PREPARE_UPDATE = 1 - START_LOAD_HF_WEIGHTS = 2 - WAIT_LOAD_READY = 3 - COMMIT = 4 diff --git a/lightllm/server/router/model_infer/model_rpc.py b/lightllm/server/router/model_infer/model_rpc.py index 17dd96be85..29b8b857a3 100644 --- a/lightllm/server/router/model_infer/model_rpc.py +++ b/lightllm/server/router/model_infer/model_rpc.py @@ -25,7 +25,6 @@ PDDecodeNode, PDDPForDecodeNode, ) -from lightllm.server.router.model_infer.mode_backend.redundancy_expert_manager import RedundancyExpertManager from lightllm.server.router.model_infer.mode_backend.rl_backend_ops import RlBackendOps from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.log_utils import init_logger @@ -97,13 +96,6 @@ def exposed_init_model(self, kvargs): logger.info(f"use {self.backend.__class__.__name__}") self.backend.init_model(kvargs) self.rl_backend_ops = RlBackendOps(self.backend) if self.args.enable_rl else None - - # only deepseekv3 can support auto_update_redundancy_expert - if self.args.auto_update_redundancy_expert: - self.redundancy_expert_manager = RedundancyExpertManager(self.backend.model) - logger.info("init redundancy_expert_manager") - else: - self.redundancy_expert_manager = None return def exposed_get_max_total_token_num(self): diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index acd2ac6711..2521fa8b36 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -96,64 +96,13 @@ def get_lightllm_websocket_max_message_size(): return int(os.getenv("LIGHTLLM_WEBSOCKET_MAX_SIZE", 128 * 1024 * 1024)) -# get_redundancy_expert_ids and get_redundancy_expert_num are primarily -# used to obtain the IDs and number of redundant experts during inference. -# They depend on a configuration file specified by ep_redundancy_expert_config_path, -# which is a JSON formatted text file. -# The content format is as follows: -# { -# "redundancy_expert_num": 1, # Number of redundant experts per rank -# "0": [0], # Key: layer_index (string), -# # Value: list of original expert IDs that are redundant for this layer -# "1": [0], -# "default": [0] # Default list of redundant expert IDs if layer-specific entry is not found -# } - - -@lru_cache(maxsize=None) -def get_redundancy_expert_ids(layer_index: int): - """ - Get the redundancy expert ids from the environment variable. - :return: List of redundancy expert ids. - """ - args = get_env_start_args() - if args.ep_redundancy_expert_config_path is None: - return [] - - with open(args.ep_redundancy_expert_config_path, "r") as f: - config = json.load(f) - if str(layer_index) in config: - return config[str(layer_index)] - else: - return config.get("default", []) - - -@lru_cache(maxsize=None) -def get_redundancy_expert_num(): - """ - Get the number of redundancy experts from the environment variable. - :return: Number of redundancy experts. - """ - args = get_env_start_args() - if args.ep_redundancy_expert_config_path is None: - return 0 - - with open(args.ep_redundancy_expert_config_path, "r") as f: - config = json.load(f) - if "redundancy_expert_num" in config: - return config["redundancy_expert_num"] - else: - return 0 - - -@lru_cache(maxsize=None) -def get_redundancy_expert_update_interval(): - return int(os.getenv("LIGHTLLM_REDUNDANCY_EXPERT_UPDATE_INTERVAL", 30 * 60)) - - @lru_cache(maxsize=None) -def get_redundancy_expert_update_max_load_count(): - return int(os.getenv("LIGHTLLM_REDUNDANCY_EXPERT_UPDATE_MAX_LOAD_COUNT", 1)) +def get_eplb_step_interval(): + """返回两次 EPLB 评估之间的推理步数。""" + interval = int(os.getenv("LIGHTLLM_EPLB_STEP_INTERVAL", 20)) + if interval <= 0: + raise ValueError("LIGHTLLM_EPLB_STEP_INTERVAL must be greater than 0") + return interval @lru_cache(maxsize=None) diff --git a/test/advanced_config/redundancy_expert/test_redundancy_expert_config.json b/test/advanced_config/redundancy_expert/test_redundancy_expert_config.json deleted file mode 100644 index 241ab25ea3..0000000000 --- a/test/advanced_config/redundancy_expert/test_redundancy_expert_config.json +++ /dev/null @@ -1,180 +0,0 @@ -{ - "redundancy_expert_num": 1, - "default": [ - 0 - ], - "3": [ - 226 - ], - "4": [ - 123 - ], - "5": [ - 187 - ], - "6": [ - 138 - ], - "7": [ - 132 - ], - "8": [ - 240 - ], - "9": [ - 4 - ], - "10": [ - 88 - ], - "11": [ - 60 - ], - "12": [ - 161 - ], - "13": [ - 178 - ], - "14": [ - 80 - ], - "15": [ - 144 - ], - "16": [ - 195 - ], - "17": [ - 251 - ], - "18": [ - 226 - ], - "19": [ - 87 - ], - "20": [ - 149 - ], - "21": [ - 45 - ], - "22": [ - 214 - ], - "23": [ - 41 - ], - "24": [ - 46 - ], - "25": [ - 156 - ], - "26": [ - 112 - ], - "27": [ - 185 - ], - "28": [ - 58 - ], - "29": [ - 156 - ], - "30": [ - 147 - ], - "31": [ - 199 - ], - "32": [ - 16 - ], - "33": [ - 188 - ], - "34": [ - 227 - ], - "35": [ - 136 - ], - "36": [ - 84 - ], - "37": [ - 15 - ], - "38": [ - 204 - ], - "39": [ - 96 - ], - "40": [ - 226 - ], - "41": [ - 25 - ], - "42": [ - 69 - ], - "43": [ - 122 - ], - "44": [ - 152 - ], - "45": [ - 113 - ], - "46": [ - 98 - ], - "47": [ - 68 - ], - "48": [ - 13 - ], - "49": [ - 102 - ], - "50": [ - 214 - ], - "51": [ - 201 - ], - "52": [ - 182 - ], - "53": [ - 235 - ], - "54": [ - 162 - ], - "55": [ - 125 - ], - "56": [ - 62 - ], - "57": [ - 121 - ], - "58": [ - 105 - ], - "59": [ - 236 - ], - "60": [ - 117 - ] -} \ No newline at end of file diff --git a/unit_tests/common/basemodel/triton_kernel/test_redundancy_topk_ids_repair.py b/unit_tests/common/basemodel/triton_kernel/test_redundancy_topk_ids_repair.py deleted file mode 100644 index 16131ef935..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_redundancy_topk_ids_repair.py +++ /dev/null @@ -1,151 +0,0 @@ -import torch -import pytest -from lightllm.common.basemodel.triton_kernel.redundancy_topk_ids_repair import redundancy_topk_ids_repair -from lightllm.common.basemodel.triton_kernel.redundancy_topk_ids_repair import expert_id_counter -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -def test_redundancy_topk_ids_repair(): - ep_expert_num = 4 - global_rank = 0 - redundancy_expert_num = 1 - topk_ids = torch.tensor( - [ - [0, 1, 2, 3], - [7, 9, 10, 11], - [1, 3, 5, 7], - ], - dtype=torch.int64, - device="cuda", - ) - - redundancy_expert_ids = torch.tensor( - [ - 0, - ], - dtype=torch.int64, - device="cuda", - ) - - expert_id_counter = torch.zeros(12, dtype=torch.int64, device="cuda") - - redundancy_topk_ids_repair( - topk_ids=topk_ids, - redundancy_expert_ids=redundancy_expert_ids, - ep_expert_num=ep_expert_num, - global_rank=global_rank, - expert_counter=expert_id_counter, - enable_counter=True, - ) - - ans_topk_ids = torch.tensor( - [ - [0, 1, 2, 3], - [7, 9, 10, 11], - [1, 3, 5, 7], - ], - dtype=torch.int64, - device="cuda", - ) - ans_topk_ids = (ans_topk_ids // ep_expert_num) * redundancy_expert_num + ans_topk_ids - new_redundancy_expert_ids = (redundancy_expert_ids // ep_expert_num) * redundancy_expert_num + redundancy_expert_ids - ans_topk_ids[ans_topk_ids == new_redundancy_expert_ids[0]] = ( - (ep_expert_num + redundancy_expert_num) * global_rank + ep_expert_num + 0 - ) - - assert torch.equal(topk_ids, ans_topk_ids) - assert torch.equal( - expert_id_counter, torch.tensor([1, 2, 1, 2, 0, 1, 0, 2, 0, 1, 1, 1], dtype=torch.int64, device="cuda") - ) - - ep_expert_num = 4 - global_rank = 1 - redundancy_expert_num = 1 - topk_ids = torch.tensor( - [ - [0, 1, 2, 3], - [7, 9, 10, 11], - [1, 3, 5, 7], - ], - dtype=torch.int64, - device="cuda", - ) - - redundancy_expert_ids = torch.tensor( - [ - 5, - ], - dtype=torch.int64, - device="cuda", - ) - redundancy_topk_ids_repair( - topk_ids=topk_ids, - redundancy_expert_ids=redundancy_expert_ids, - ep_expert_num=ep_expert_num, - global_rank=global_rank, - ) - - ans_topk_ids = torch.tensor( - [ - [0, 1, 2, 3], - [7, 9, 10, 11], - [1, 3, 5, 7], - ], - dtype=torch.int64, - device="cuda", - ) - ans_topk_ids = (ans_topk_ids // ep_expert_num) * redundancy_expert_num + ans_topk_ids - new_redundancy_expert_ids = (redundancy_expert_ids // ep_expert_num) * redundancy_expert_num + redundancy_expert_ids - ans_topk_ids[ans_topk_ids == new_redundancy_expert_ids[0]] = ( - (ep_expert_num + redundancy_expert_num) * global_rank + ep_expert_num + 0 - ) - - assert torch.equal(topk_ids, ans_topk_ids) - - -def test_expert_id_counter(): - token_num = 256 - tok_ids = torch.randint( - low=0, - high=12, - size=(token_num, 8), - dtype=torch.int64, - device="cuda", - ) - expert_counter = torch.zeros(12, dtype=torch.int64, device="cuda") - expert_id_counter(topk_ids=tok_ids, expert_counter=expert_counter) - - ans_expert_counter = torch.zeros(12, dtype=torch.int64, device="cuda") - ids, counts = torch.unique(tok_ids.view(-1), return_counts=True) - ans_expert_counter[ids] = counts - - assert torch.equal(expert_counter, ans_expert_counter) - - # test speed - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - for _ in range(100): - tok_ids = torch.randint( - low=0, - high=12, - size=(token_num, 8), - dtype=torch.int64, - device="cuda", - ) - expert_counter = torch.zeros(12, dtype=torch.int64, device="cuda") - expert_id_counter(topk_ids=tok_ids, expert_counter=expert_counter) - graph.replay() - - start_event = torch.cuda.Event(enable_timing=True) - start_event.record() - graph.replay() - end_event = torch.cuda.Event(enable_timing=True) - end_event.record() - torch.cuda.synchronize() - logger.info(f"expert_id_counter time cost: {start_event.elapsed_time(end_event)} ms") - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/fused_moe/test_eplb.py b/unit_tests/common/fused_moe/test_eplb.py new file mode 100644 index 0000000000..8079353bfd --- /dev/null +++ b/unit_tests/common/fused_moe/test_eplb.py @@ -0,0 +1,3054 @@ +import threading +import time +from contextlib import nullcontext +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.server.router.model_infer.mode_backend.eplb.placement import ( + EPLBPlanner, + GreedyEPLBPlanner, + build_initial_local_expert_ids, + build_logical_to_physical_map, + create_eplb_planner, +) +from lightllm.server.api_cli import make_argument_parser +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.server.router.model_infer.infer_batch import g_infer_context +from lightllm.server.router.model_infer.mode_backend.eplb import ( + runtime_manager as manager_module, +) +from lightllm.server.router.model_infer.mode_backend.eplb import async_task as async_task_module +from lightllm.server.router.model_infer.mode_backend.eplb import ( + async_load_gather_task as load_gather_module, +) +from lightllm.server.router.model_infer.mode_backend.eplb import metrics as eplb_metrics +from lightllm.server.router.model_infer.mode_backend.eplb import ( + async_placement_plan_task as plan_module, +) +from lightllm.server.router.model_infer.mode_backend.eplb import ( + async_expert_transfer as transfer_module, +) +from lightllm.server.router.model_infer.mode_backend.eplb import ( + async_transfer_planner as transfer_planner_module, +) +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl import ( + deepgemm_impl as deepgemm_module, +) +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl import ( + create_fuse_moe_impl, + FuseMoeMarlin, + FuseMoeTriton, +) +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl.base_impl import ( + FuseMoeBaseImpl, +) +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe import ( + fused_moe_weight as fused_weight_module, +) +from lightllm.server.router.model_infer.mode_backend.eplb.eplb_utils import extract_eplb_expert_tensors +from lightllm.server.router.model_infer.mode_backend.eplb.async_expert_transfer import ( + EPLBTransferInfo, + ExpertTensorBuffer, + PinnedMemoryEPLBTransfer, + build_transfer_plan, +) + + +def _test_moe_impl( + *, + eplb=False, + num_logical_experts=128, + world_size=16, + num_redundant_experts_per_rank=1, + prefill_route_counter=None, + prefill_route_sample_index=None, + recording=False, +): + logical_to_physical_map = None + if eplb: + if prefill_route_counter is None: + prefill_route_counter = torch.zeros((24, num_logical_experts), dtype=torch.int64) + if prefill_route_sample_index is None: + prefill_route_sample_index = torch.zeros((2,), dtype=torch.int64) + logical_to_physical_map = torch.zeros((num_logical_experts, world_size + 3), dtype=torch.int32) + logical_to_physical_map[:, :3] = 1 + else: + num_redundant_experts_per_rank = 0 + return SimpleNamespace( + n_routed_experts=num_logical_experts, + num_total_physical_experts=(num_logical_experts + world_size * num_redundant_experts_per_rank), + num_redundant_experts_per_rank=num_redundant_experts_per_rank, + local_logics_expert_ids_list=list(range(num_logical_experts // world_size + num_redundant_experts_per_rank)), + logical_to_physical_map=logical_to_physical_map, + prefill_route_counter=prefill_route_counter, + prefill_route_sample_index=prefill_route_sample_index, + recording=recording, + ) + + +def _set_deepgemm_runtime(impl, runtime): + for name in ( + "num_total_physical_experts", + "num_redundant_experts_per_rank", + "logical_to_physical_map", + "prefill_route_counter", + "prefill_route_sample_index", + "recording", + ): + setattr(impl, name, getattr(runtime, name)) + + +def _initial_expert_placement(num_logical_experts, world_size, num_redundant_experts_per_rank): + return torch.tensor( + build_initial_local_expert_ids( + num_logical_experts, + world_size, + num_redundant_experts_per_rank, + ), + dtype=torch.int64, + ) + + +def _planner_load(load_by_layer, world_size): + """构造 [rank, layer, sample, logical expert] CPU planner 输入。""" + aggregated_load = torch.as_tensor(load_by_layer, dtype=torch.float64) + assert aggregated_load.ndim == 2 + load = torch.zeros( + (world_size, aggregated_load.shape[0], 1, aggregated_load.shape[1]), + dtype=torch.float64, + ) + load[0, :, 0] = aggregated_load + return load + + +def _rank_to_logic_expert_ids(redundant_placement, num_logical_experts): + num_ranks = len(redundant_placement) + num_primary_experts_per_rank = num_logical_experts // num_ranks + return [ + list( + range( + rank * num_primary_experts_per_rank, + (rank + 1) * num_primary_experts_per_rank, + ) + ) + + list(rank_redundant_expert_ids) + for rank, rank_redundant_expert_ids in enumerate(redundant_placement) + ] + + +def test_base_call_template_forwards_selection_and_capture_callback(): + class Impl(FuseMoeBaseImpl): + def _select_experts( + self, + input_tensor, + router_logits, + correction_bias, + top_k, + renormalize, + use_grouped_topk, + topk_group, + num_expert_group, + scoring_func, + per_expert_scale=None, + ): + return "weights", "logical_ids" + + def _prepare_expert_execution(self, topk_weights, topk_ids, is_prefill, shared_expert_gate=None): + seen["prepare"] = {"topk_ids": topk_ids, "is_prefill": is_prefill} + return topk_weights, "physical_ids" + + def _fused_experts( + self, + input_tensor, + w13, + w2, + topk_weights, + topk_ids, + is_prefill, + router_logits=None, + ): + seen["fused"] = {"topk_ids": topk_ids} + return "output" + + seen, captured = {}, [] + impl = Impl(4, 0, 1.0, SimpleNamespace()) + result = impl( + "input", + "logits", + "w13", + "w2", + None, + "softmax", + 2, + False, + False, + 0, + 0, + moe_capture_callback=captured.append, + is_prefill=True, + ) + assert result == "output" + assert captured == ["logical_ids"] + assert seen["prepare"]["topk_ids"] == "logical_ids" + assert seen["prepare"]["is_prefill"] is True + assert seen["fused"]["topk_ids"] == "physical_ids" + + with pytest.raises(AssertionError, match="is_prefill must be explicitly specified"): + impl( + "input", + "logits", + "w13", + "w2", + None, + "softmax", + 2, + False, + False, + 0, + 0, + ) + + +def test_factory_selects_all_paths_without_ep_constructor_state(monkeypatch): + plain_quant = SimpleNamespace(method_name="none") + marlin_quant = SimpleNamespace(method_name="awq_marlin") + monkeypatch.setattr(FuseMoeMarlin, "create_workspace", lambda self: None) + monkeypatch.setattr( + deepgemm_module, + "get_env_start_args", + lambda: SimpleNamespace(eplb_num_redundant_experts_per_rank=0), + ) + monkeypatch.setattr(deepgemm_module, "get_global_world_size", lambda: 2) + monkeypatch.setattr(deepgemm_module, "get_global_rank", lambda: 0) + ep_impl = create_fuse_moe_impl( + n_routed_experts=4, + num_fused_shared_experts=0, + routed_scaling_factor=1.0, + quant_method=plain_quant, + enable_ep_moe=True, + ) + assert isinstance(ep_impl, deepgemm_module.FuseMoeDeepGEMM) + assert ep_impl.num_total_physical_experts == 4 + assert not hasattr(ep_impl, "num_primary_experts_per_rank") + assert not hasattr(ep_impl, "prefill_route_counter") + assert not hasattr(ep_impl, "expert_parallel_state") + assert isinstance( + create_fuse_moe_impl( + n_routed_experts=4, + num_fused_shared_experts=0, + routed_scaling_factor=1.0, + quant_method=plain_quant, + ), + FuseMoeTriton, + ) + assert isinstance( + create_fuse_moe_impl( + n_routed_experts=4, + num_fused_shared_experts=0, + routed_scaling_factor=1.0, + quant_method=marlin_quant, + ), + FuseMoeMarlin, + ) + + +def test_find_fused_moe_weights_uses_layer_experts_in_model_order(monkeypatch): + class FakeFusedMoeWeight: + def __init__(self, layer_num, enable_ep_moe=True): + self.layer_num_ = layer_num + self.enable_ep_moe = enable_ep_moe + + monkeypatch.setattr(manager_module, "FusedMoeWeight", FakeFusedMoeWeight) + first = FakeFusedMoeWeight(1) + second = FakeFusedMoeWeight(3) + disabled = FakeFusedMoeWeight(2, enable_ep_moe=False) + model = SimpleNamespace( + trans_layers_weight=[ + SimpleNamespace(experts=first), + SimpleNamespace(), + SimpleNamespace(experts=disabled), + SimpleNamespace(experts=second), + ] + ) + + assert manager_module._find_fused_moe_weights(model) == [first, second] + + +def test_eplb_redundant_experts_default_to_disabled(): + parser = make_argument_parser() + + assert parser.parse_args([]).eplb_num_redundant_experts_per_rank == 0 + assert parser.parse_args(["--eplb_num_redundant_experts_per_rank", "3"]).eplb_num_redundant_experts_per_rank == 3 + assert StartArgs().eplb_num_redundant_experts_per_rank == 0 + assert parser.parse_args([]).eplb_plan_mode == "greedy" + assert parser.parse_args(["--eplb_plan_mode", "greedy"]).eplb_plan_mode == "greedy" + assert StartArgs().eplb_plan_mode == "greedy" + assert parser.parse_args([]).eplb_rebalance_count == 1 + assert parser.parse_args(["--eplb_rebalance_count", "-1"]).eplb_rebalance_count == -1 + assert parser.parse_args(["--eplb_rebalance_count", "0"]).eplb_rebalance_count == 0 + assert StartArgs().eplb_rebalance_count == 1 + assert parser.parse_args([]).eplb_config_path is None + assert parser.parse_args(["--eplb_config_path", "/tmp/eplb.json"]).eplb_config_path == "/tmp/eplb.json" + assert StartArgs().eplb_config_path is None + + +@pytest.mark.parametrize( + ("num_logical_experts", "num_ranks", "num_redundant_experts_per_rank", "expected"), + [ + (8, 4, 2, [[0, 1, 2, 3], [2, 3, 4, 5], [4, 5, 6, 7], [6, 7, 0, 1]]), + (6, 3, 4, [[0, 1, 2, 3, 4, 5], [2, 3, 4, 5, 0, 1], [4, 5, 0, 1, 2, 3]]), + ], +) +def test_build_initial_local_expert_ids( + num_logical_experts, + num_ranks, + num_redundant_experts_per_rank, + expected, +): + actual = build_initial_local_expert_ids( + num_logical_experts, + num_ranks, + num_redundant_experts_per_rank, + ) + + assert actual == expected + + +def test_build_initial_local_expert_ids_rejects_local_or_duplicate_replicas(): + with pytest.raises(AssertionError): + build_initial_local_expert_ids(8, 4, 7) + + +def test_eplb_planner_defines_an_abstract_planning_interface(): + with pytest.raises(TypeError): + EPLBPlanner() + + assert isinstance(GreedyEPLBPlanner(2, 1), EPLBPlanner) + + +def test_create_eplb_planner_selects_requested_algorithm(): + planner = create_eplb_planner( + "greedy", + 2, + 1, + expert_alignment=1, + ) + + assert isinstance(planner, GreedyEPLBPlanner) + + with pytest.raises(ValueError, match="unsupported EPLB plan mode"): + create_eplb_planner( + "unknown", + 2, + 1, + expert_alignment=1, + ) + + +def test_eplb_planner_builds_legal_concrete_slot_layout(): + planner = GreedyEPLBPlanner( + 4, + 1, + expert_alignment=1, + ) + current = _initial_expert_placement(8, 4, 1).unsqueeze(0).tolist() + load = torch.ones((4, 1, 1, 8), dtype=torch.int64) + load[:, :, :, 0] = 1000 + load[:, :, :, 4] = 500 + + result = planner.plan(logical_expert_load_samples=load, current_placement=current) + placement = result[0] + + for row in placement: + assert len(row) == 3 + assert len(row) == len(set(row)) + assert 0 in row + assert set(expert for row in placement for expert in row) == set(range(8)) + assert any(row[:2] != list(range(rank * 2, (rank + 1) * 2)) for rank, row in enumerate(placement)) + assert isinstance(result, list) + + +def test_eplb_planner_returns_deterministic_layout_for_zero_load_experts(): + planner = GreedyEPLBPlanner(2, 1) + current = [[[0, 1, 3], [2, 3, 1]]] + + result = planner.plan( + logical_expert_load_samples=_planner_load([[0, 0, 0, 0]], world_size=2), + current_placement=current, + ) + + assert result == [[[0, 1, 3], [2, 0, 1]]] + + +def test_eplb_planner_aggregates_rank_and_sample_dimensions(): + planner = GreedyEPLBPlanner(2, 1) + current = [[[0, 1, 3], [2, 3, 1]]] + raw_load = torch.tensor( + [ + [[[500, 1, 2, 3], [300, 4, 5, 6]]], + [[[100, 7, 8, 9], [200, 10, 11, 12]]], + ], + dtype=torch.int64, + ) + aggregated_load = raw_load.sum(dim=(0, 2)) + equivalent_raw_load = torch.zeros_like(raw_load) + equivalent_raw_load[0, :, 0] = aggregated_load + + result = planner.plan(logical_expert_load_samples=raw_load, current_placement=current) + + assert result == planner.plan( + logical_expert_load_samples=equivalent_raw_load, + current_placement=current, + ) + + +def test_eplb_planner_plans_each_layer_independently_then_combines_results(): + planner = GreedyEPLBPlanner(2, 1) + current_layer = [[0, 1, 3], [2, 3, 1]] + current = [[row[:] for row in current_layer], [row[:] for row in current_layer]] + + load = _planner_load( + [ + [1000, 1, 1, 1], + [0, 0, 0, 0], + ], + world_size=2, + ) + + result = planner.plan(logical_expert_load_samples=load, current_placement=current) + + assert result == [ + [[0, 1, 3], [2, 0, 1]], + [[0, 1, 3], [2, 0, 1]], + ] + + +def test_eplb_planner_iteratively_places_hot_expert_on_idle_rank(): + planner = GreedyEPLBPlanner(2, 1) + current = [[[0, 1, 3], [2, 3, 1]]] + + result = planner.plan( + logical_expert_load_samples=_planner_load([[1000, 1, 1, 1]], world_size=2), + current_placement=current, + ) + + assert result == [[[0, 1, 3], [2, 0, 1]]] + + +def test_eplb_planner_repeatedly_splits_the_hottest_remaining_expert(): + planner = GreedyEPLBPlanner(4, 3) + current = _initial_expert_placement(8, 4, 3).unsqueeze(0).tolist() + + result = planner.plan( + logical_expert_load_samples=_planner_load([[1000, 900, 800, 700, 1, 1, 1, 1]], world_size=4), + current_placement=current, + ) + + replica_counts = [sum(expert in row for row in result[0]) for expert in range(8)] + assert replica_counts == [4, 4, 4, 4, 1, 1, 1, 1] + + +def test_eplb_planner_balances_expert_groups_with_equal_replica_counts(): + planner = GreedyEPLBPlanner(4, 1) + + placement = planner._distribute_remaining_experts( + redundant_experts=[0], + expert_groups=[ + (1, 1, 8.0), + (2, 1, 7.0), + (3, 1, 6.0), + (4, 1, 5.0), + (5, 1, 4.0), + (6, 1, 3.0), + (7, 1, 2.0), + (8, 1, 1.0), + ], + ) + + assert placement == [ + [0, 1, 8], + [0, 2, 7], + [0, 3, 6], + [0, 4, 5], + ] + + +def test_eplb_planner_places_replicas_of_one_expert_on_distinct_ranks(): + planner = GreedyEPLBPlanner(2, 1) + + placement = planner._distribute_remaining_experts( + redundant_experts=[0], + expert_groups=[ + (1, 2, 5.0), + (2, 1, 8.0), + (3, 1, 1.0), + ], + ) + + assert placement == [ + [0, 1, 2], + [0, 1, 3], + ] + assert all(len(row) == len(set(row)) for row in placement) + + +def test_eplb_planner_places_single_replicas_by_rank_load_before_free_slots(): + planner = GreedyEPLBPlanner(4, 1) + + placement = planner._distribute_remaining_experts( + redundant_experts=[0], + expert_groups=[ + (1, 3, 10.0), + (2, 2, 1.0), + (3, 1, 8.0), + (4, 1, 7.0), + (5, 1, 6.0), + (6, 1, 5.0), + (7, 1, 4.0), + (8, 1, 3.0), + (9, 1, 2.0), + ], + ) + + # 多副本专家平铺后,rank 3 的剩余槽位比 rank 1、2 少,但负载最低; + # 因此它仍连续取得最热的两个单副本专家,并率先填满。 + assert placement == [ + [0, 1, 2, 7], + [0, 1, 5, 9], + [0, 1, 6, 8], + [0, 2, 3, 4], + ] + + +def test_eplb_planner_matches_documented_two_stage_distribution_example(): + planner = GreedyEPLBPlanner(4, 1) + expert_groups = [ + (1, 2, 6.0), + (2, 1, 9.0), + (3, 1, 8.0), + (4, 1, 7.0), + (5, 1, 5.0), + (6, 1, 4.0), + (7, 1, 3.0), + ] + + placement = planner._distribute_remaining_experts( + redundant_experts=[0], + expert_groups=expert_groups, + ) + + assert placement == [ + [0, 1, 4], + [0, 1, 5], + [0, 2, 7], + [0, 3, 6], + ] + load_per_replica = {expert: load for expert, _, load in expert_groups} + assert [sum(load_per_replica[expert] for expert in row[1:]) for row in placement] == [13.0, 11.0, 12.0, 12.0] + + +def test_eplb_planner_greedily_matches_candidate_ranks_before_reusing_slots(): + planner = GreedyEPLBPlanner(3, 1) + current = [ + [0, 1, 2], + [3, 4, 5], + [6, 7, 8], + ] + candidate = [ + [3, 4, 9], + [6, 7, 10], + [0, 1, 11], + ] + + placement = planner._reuse_current_slots(candidate, current) + + assert placement == [ + [0, 1, 11], + [3, 4, 9], + [6, 7, 10], + ] + + +def test_eplb_planner_keeps_selected_experts_in_their_current_slots(): + planner = GreedyEPLBPlanner(4, 2) + current = _initial_expert_placement(8, 4, 2).unsqueeze(0).tolist() + + result = planner.plan( + logical_expert_load_samples=_planner_load([[50, 98, 54, 6, 34, 66, 63, 52]], world_size=4), + current_placement=current, + ) + + # 只要专家仍分配在同一个 rank,就保留其原物理槽位。 + for current_row, target_row in zip(current[0], result[0]): + for slot, expert in enumerate(current_row): + if expert in target_row: + assert target_row[slot] == expert + + +def test_eplb_planner_fills_every_rank_with_distinct_nonlocal_experts(): + planner = GreedyEPLBPlanner( + 4, + 1, + ) + current = _initial_expert_placement(16, 4, 1).unsqueeze(0).tolist() + load_by_layer_and_rank = torch.randint( + 0, + 10000, + (1, 4, 16), + generator=torch.Generator().manual_seed(2), + ) + load = load_by_layer_and_rank.permute(1, 0, 2).unsqueeze(dim=2) + + result = planner.plan(logical_expert_load_samples=load, current_placement=current) + + assert len(result) == len(current) + assert all(len(actual) == len(expected) for actual, expected in zip(result[0], current[0])) + for row in result[0]: + assert len(row) == len(set(row)) == 5 + assert set(expert for row in result[0] for expert in row) == set(range(16)) + + +def test_eplb_planner_supports_multiple_redundant_experts_per_rank(): + planner = GreedyEPLBPlanner(4, 3) + current = _initial_expert_placement(16, 4, 3).unsqueeze(0).tolist() + load = _planner_load( + [ + [ + 22613, + 26852, + 21852, + 23480, + 13270, + 14695, + 28735, + 22303, + 15324, + 19604, + 21492, + 25458, + 14120, + 12130, + 18620, + 22888, + ] + ], + world_size=4, + ) + + result = planner.plan(logical_expert_load_samples=load, current_placement=current) + + for row in result[0]: + assert len(row) == len(set(row)) == 7 + replica_counts = [sum(expert in row for row in result[0]) for expert in range(16)] + assert replica_counts[1] == replica_counts[6] == replica_counts[11] == 4 + assert sum(replica_counts) == 28 + + +def test_fused_moe_loads_default_replicas_into_their_physical_rows(): + weight = object.__new__(fused_weight_module.FusedMoeWeight) + weight.lock = threading.Lock() + loaded = [] + + def load_weight(expert, local, _weights): + loaded.append(("weight", expert, local)) + + def load_scale(expert, local, _weights): + loaded.append(("scale", expert, local)) + + def load_zero_point(expert, local, _weights): + loaded.append(("zero", expert, local)) + + weight._load_expert = load_weight + weight._load_expert_scale = load_scale + weight._load_expert_zero_point = load_zero_point + local_logic_expert_ids_list = build_initial_local_expert_ids(8, 4, 2)[0] + + weight._load_weight(local_logic_expert_ids_list, {}) + + assert local_logic_expert_ids_list == [0, 1, 2, 3] + assert [entry for entry in loaded if entry[0] == "weight"] == [ + ("weight", 0, 0), + ("weight", 1, 1), + ("weight", 2, 2), + ("weight", 3, 3), + ] + + +def test_logical_to_physical_map_selects_one_physical_expert(): + rank_to_logic_expert_ids = [[0, 1, 2, 3], [2, 3, 0, 1]] + logical_to_physical = build_logical_to_physical_map( + rank_to_logic_expert_ids, + num_logical_experts=4, + current_rank=0, + node_world_size=2, + ) + + assert isinstance(logical_to_physical, list) + assert len(logical_to_physical) == 4 + assert all(len(row) == 11 for row in logical_to_physical) + assert [row[0] for row in logical_to_physical] == [2, 2, 2, 2] + assert [row[1] for row in logical_to_physical] == [2, 2, 2, 2] + assert [row[2] for row in logical_to_physical] == [1, 1, 1, 1] + assert [row[3] for row in logical_to_physical] == [0, 1, 2, 3] + assert all(physical_id >= 0 for row in logical_to_physical for physical_id in row[3 : 3 + row[0]]) + assert all(physical_id == -1 for row in logical_to_physical for physical_id in row[3 + row[0] :]) + + +def test_logical_to_physical_map_requires_expert_count_divisible_by_rank_count(): + rank_to_logic_expert_ids = [[0, 1, 0], [2, 3, 1]] + + with pytest.raises(AssertionError): + build_logical_to_physical_map( + rank_to_logic_expert_ids, + num_logical_experts=5, + current_rank=0, + node_world_size=2, + ) + + +def test_logical_to_physical_map_supports_all_redundant_slots_for_one_expert(): + logical_to_physical = build_logical_to_physical_map( + [[0, 1, 0, 0], [2, 3, 0, 0]], + num_logical_experts=4, + current_rank=0, + node_world_size=2, + ) + + # 1 个主副本加上 2 个 rank 的全部 4 个冗余槽。 + assert logical_to_physical[0][:3] == [5, 5, 3] + assert len(logical_to_physical[0][3:]) == 8 + assert len(set(logical_to_physical[0][3:8])) == 5 + assert logical_to_physical[0][8:] == [-1, -1, -1] + + +def test_logical_to_physical_map_prefers_current_rank_replica(): + redundant = [[4], [5], [0], [1]] + rank_to_logic_expert_ids = _rank_to_logic_expert_ids(redundant, 8) + rank0_map = build_logical_to_physical_map( + rank_to_logic_expert_ids, + num_logical_experts=8, + current_rank=0, + node_world_size=2, + ) + rank1_map = build_logical_to_physical_map( + rank_to_logic_expert_ids, + num_logical_experts=8, + current_rank=1, + node_world_size=2, + ) + fallback_redundant = [[4], [5], [0], [1], [2], [3]] + rank4_map = build_logical_to_physical_map( + _rank_to_logic_expert_ids(fallback_redundant, 12), + 12, + current_rank=4, + node_world_size=2, + ) + + assert rank0_map[0][:3] == [2, 1, 1] + assert rank1_map[0][:3] == [2, 1, 0] + assert rank0_map[0][3] == 0 + assert set(rank0_map[0][3:5]) == {0, 8} + assert set(rank1_map[0][3:5]) == {0, 8} + assert rank4_map[0][:3] == [2, 0, 0] + assert set(rank4_map[0][3:5]) == {0, 8} + + +def test_nonlocal_rank_without_same_node_replica_routes_across_all_replicas(): + redundant = [[4], [5], [0], [1]] + rank_to_logic_expert_ids = _rank_to_logic_expert_ids(redundant, 8) + maps = [ + build_logical_to_physical_map( + rank_to_logic_expert_ids, + 8, + current_rank=rank, + node_world_size=1, + ) + for rank in range(4) + ] + assert maps[0][0][:3] == [2, 1, 1] + assert maps[1][0][:3] == [2, 0, 0] + assert maps[2][0][:3] == [2, 1, 1] + assert maps[3][0][:3] == [2, 0, 0] + assert all(set(logical_map[0][3:5]) == {0, 8} for logical_map in maps) + + +def test_nonlocal_rank_prefers_same_node_replica(): + rank_to_logic_expert_ids = [ + [1, 2], + [0, 3], + [0, 4], + [5, 6], + [0, 7], + [1, 2], + [3, 4], + [5, 6], + ] + + rank0_map = build_logical_to_physical_map( + rank_to_logic_expert_ids, + num_logical_experts=8, + current_rank=0, + node_world_size=4, + ) + + # Expert 0 is on same-node ranks 1 and 2 (physical IDs 2 and 4), plus + # remote rank 4 (physical ID 8). Hash routing only uses the first two. + assert rank0_map[0][:6] == [3, 2, 0, 2, 4, 8] + + +def test_current_rank_moves_local_replica_to_front_without_changing_copies(): + # Expert 0 is primary on rank 0 and redundant on rank 1。两个 rank 都优先 + # 自己的本地副本,因此路由槽的起点不同,但候选集合和数量保持一致。 + redundant = [[1], [0], [3], [2]] + rank_to_logic_expert_ids = _rank_to_logic_expert_ids(redundant, 4) + rank0_map = build_logical_to_physical_map( + rank_to_logic_expert_ids, + 4, + current_rank=0, + node_world_size=1, + ) + rank1_map = build_logical_to_physical_map( + rank_to_logic_expert_ids, + 4, + current_rank=1, + node_world_size=1, + ) + + assert rank0_map[0] == [2, 1, 1, 0, 3, -1, -1, -1, -1, -1, -1] + assert rank1_map[0] == [2, 1, 1, 3, 0, -1, -1, -1, -1, -1, -1] + + +def test_current_rank_stably_moves_all_local_physical_ids_to_front(): + # Expert 0 在 rank 0/1/2 上依次对应 physical IDs [0, 1, 3, 5]。 + # 对 rank 1 构建路由表时,只把本地 ID 3 移到最前面;其余远端 ID + # 仍保持原来的 [0, 1, 5] 顺序。 + rank_to_logic_expert_ids = [[0, 0], [1, 0], [2, 0]] + + rank1_map = build_logical_to_physical_map( + rank_to_logic_expert_ids, + num_logical_experts=3, + current_rank=1, + node_world_size=1, + ) + + assert rank1_map[0] == [4, 1, 1, 3, 0, 1, 5, -1, -1] + + +def test_transfer_plan_respects_explicit_target_slots(): + current = [[0, 1, 4, 5], [2, 3, 6, 7], [4, 5, 0, 1], [6, 7, 2, 3]] + target = [[0, 1, 5, 4], [2, 3, 7, 6], [4, 5, 1, 0], [6, 7, 3, 2]] + + plan = build_transfer_plan(current, target, 3, num_logical_experts=8, world_size=4) + transfer_infos = [transfer_info for transfer_batch in plan for transfer_info in transfer_batch] + + assert all(info.layer_index == 3 for info in transfer_infos) + assert all( + current[info.source_rank][info.source_local_expert_index] == info.source_logical_expert_id + for info in transfer_infos + ) + assert {(info.dest_rank, info.source_logical_expert_id) for info in transfer_infos} == { + (rank, target[rank][slot]) for rank in range(4) for slot in range(2, 4) + } + + +def test_manager_evaluating_starts_raw_load_gather_without_modifying_counters(monkeypatch): + counters = [ + torch.tensor([[10, 11]], dtype=torch.int64), + torch.tensor([[40, 41]], dtype=torch.int64), + ] + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.state = manager_module.EPLBManagerState.EVALUATING + manager._eplb_impls = [ + _test_moe_impl( + eplb=True, + prefill_route_counter=counter, + num_logical_experts=2, + world_size=1, + ) + for counter in counters + ] + manager.num_logical_experts = 2 + manager.global_rank = 1 + manager.world_size = 1 + manager.load_gather_group = object() + manager.control_group = object() + manager.max_rebalance_count = -1 + manager.completed_rebalance_count = 0 + tasks = [] + + class LoadGatherTask: + def __init__(self, **kwargs): + self.kwargs = kwargs + self.started = False + tasks.append(self) + + def start(self): + self.started = True + + monkeypatch.setattr(manager_module, "EPLBLoadGatherTask", LoadGatherTask) + monkeypatch.setattr( + manager_module.dist, + "broadcast_object_list", + lambda values, **_kwargs: values.__setitem__(0, True), + ) + + manager._step_evaluating() + + assert manager.state is manager_module.EPLBManagerState.WAIT_LOAD_GATHER_FINISHED + assert len(tasks) == 1 + assert tasks[0].started + assert tasks[0].kwargs["load_gather_group"] is manager.load_gather_group + assert torch.equal( + tasks[0].kwargs["local_load_samples"], + torch.tensor([[[10, 11]], [[40, 41]]], dtype=torch.int64), + ) + assert torch.equal(counters[0], torch.tensor([[10, 11]], dtype=torch.int64)) + assert torch.equal(counters[1], torch.tensor([[40, 41]], dtype=torch.int64)) + + +def test_load_gather_task_preserves_rank_layer_and_sample_dimensions(monkeypatch): + local_load = torch.tensor( + [ + [[1, 2], [3, 4]], + [[5, 6], [7, 8]], + ], + dtype=torch.int64, + ) + load_gather_group = object() + calls = [] + + def all_gather(output, local, *, group): + calls.append((local, group)) + output[0].copy_(local) + output[1].copy_(local + 100) + + monkeypatch.setattr(load_gather_module.dist, "all_gather", all_gather) + monkeypatch.setattr(load_gather_module.dist, "get_world_size", lambda *, group: 2) + task = load_gather_module.EPLBLoadGatherTask( + local_load_samples=local_load, + load_gather_group=load_gather_group, + ) + + task._run() + + assert task.is_finished() + assert len(calls) == 1 + assert torch.equal(calls[0][0], local_load) + assert calls[0][1] is load_gather_group + assert task.result is not None + assert task.result.shape == (2, 2, 2, 2) + assert torch.equal(task.result[0], local_load) + assert torch.equal(task.result[1], local_load + 100) + + +def test_manager_delegates_distribution_planning_to_planner_class(): + current_placement = [[[0, 1]]] + logical_load = torch.tensor([[[[10, 20]]]]) + calls = [] + planned_placement = [[[0, 1]]] + planner = SimpleNamespace( + plan=lambda **kwargs: ( + calls.append((kwargs["logical_expert_load_samples"], kwargs["current_placement"])) or planned_placement + ) + ) + task = plan_module.EPLBPlanTask( + planner=planner, + logical_expert_load_samples=logical_load, + current_placement=current_placement, + ) + + task._run() + + assert task.status == "succeeded" + assert task.result == planned_placement + assert len(calls) == 1 + assert calls[0][0] is logical_load + assert calls[0][1] == current_placement + + +def test_plan_task_exits_process_on_failure(monkeypatch): + def fail(**_kwargs): + raise RuntimeError("planning boom") + + task = plan_module.EPLBPlanTask( + planner=SimpleNamespace(plan=fail), + logical_expert_load_samples=torch.tensor([[[[10, 20]]]]), + current_placement=[[[1]]], + ) + exits = [] + logs = [] + monkeypatch.setattr(async_task_module.os, "_exit", exits.append) + monkeypatch.setattr(async_task_module.logger, "exception", logs.append) + + task._run() + + assert exits == [1] + assert logs == ["EPLBPlanTask failed"] + + +def test_transfer_planner_combines_all_layer_batches(monkeypatch): + current_placement = [[[0, 1], [2, 3]], [[0, 2], [1, 3]]] + target_placement = [[[2, 1], [0, 3]], [[0, 3], [1, 2]]] + transfer_infos = [ + EPLBTransferInfo(0, 2, 1, 0, 0, 0), + EPLBTransferInfo(1, 3, 1, 1, 0, 1), + ] + calls = [] + + def build_plan(**kwargs): + calls.append(kwargs) + return [[transfer_infos[kwargs["layer_index"]]]] + + monkeypatch.setattr(transfer_planner_module, "build_transfer_plan", build_plan) + planner = transfer_planner_module.EPLBTransferPlanner( + current_placement=current_placement, + target_placement=target_placement, + num_logical_experts=4, + world_size=2, + ) + + planner._run() + + assert planner.status == "succeeded" + assert planner.result == [[transfer_infos[0]], [transfer_infos[1]]] + assert calls == [ + { + "current_placement": current_placement[0], + "target_placement": target_placement[0], + "layer_index": 0, + "num_logical_experts": 4, + "world_size": 2, + }, + { + "current_placement": current_placement[1], + "target_placement": target_placement[1], + "layer_index": 1, + "num_logical_experts": 4, + "world_size": 2, + }, + ] + + +def test_transfer_planner_exits_process_on_failure(monkeypatch): + def fail(**_kwargs): + raise RuntimeError("transfer planning boom") + + planner = transfer_planner_module.EPLBTransferPlanner( + current_placement=[[[0], [1]]], + target_placement=[[[1], [0]]], + num_logical_experts=2, + world_size=2, + ) + exits = [] + logs = [] + monkeypatch.setattr(transfer_planner_module, "build_transfer_plan", fail) + monkeypatch.setattr(async_task_module.os, "_exit", exits.append) + monkeypatch.setattr(async_task_module.logger, "exception", logs.append) + + planner._run() + + assert exits == [1] + assert logs == ["EPLBTransferPlanner failed"] + + +def test_compute_critical_overhead_ratio_estimates_rank_pressure(): + load = torch.tensor([[384, 128, 128, 128]], dtype=torch.int64) + placement = [[[0, 1], [2, 3]]] + + ratio = eplb_metrics.compute_critical_overhead_ratio( + logical_expert_load=load, + placement=placement, + expert_alignment=128, + ) + + # rank loads are [512, 256], so excess critical / balanced is 128 / 384. + assert ratio == pytest.approx(1 / 3) + + +def test_compute_critical_overhead_ratio_preserves_layer_boundaries(): + load = torch.tensor( + [ + [1, 0], + [0, 1], + ], + dtype=torch.int64, + ) + + ratio = eplb_metrics.compute_critical_overhead_ratio( + logical_expert_load=load, + placement=[ + [[0], [1]], + [[0], [1]], + ], + expert_alignment=128, + ) + + # 两层的热点 rank 相反;逐层取关键路径时仍各有 100% 开销,不能相互抵消。 + assert ratio == pytest.approx(1.0) + + +def test_compute_critical_overhead_ratio_is_zero_without_load(): + assert ( + eplb_metrics.compute_critical_overhead_ratio( + logical_expert_load=torch.zeros((1, 2), dtype=torch.int64), + placement=[[[0], [1]]], + expert_alignment=128, + ) + == 0.0 + ) + + +def test_logical_expert_imbalance_percentiles_report_layer_distribution(): + load = torch.tensor( + [ + [3, 3, 3, 3], + [6, 2, 2, 2], + [9, 1, 1, 1], + [12, 0, 0, 0], + [0, 0, 0, 0], + ], + dtype=torch.int64, + ) + + assert eplb_metrics.logical_expert_imbalance_percentiles(expert_load=load) == { + 25: pytest.approx(1.0), + 50: pytest.approx(2.0), + 100: pytest.approx(4.0), + } + + +def test_logical_expert_imbalance_percentiles_are_zero_without_load(): + assert eplb_metrics.logical_expert_imbalance_percentiles(expert_load=torch.zeros((3, 4), dtype=torch.int64)) == { + 25: 0.0, + 50: 0.0, + 100: 0.0, + } + + +def test_publish_expert_load_metrics(): + calls = [] + metric_client = SimpleNamespace(gauge_set=lambda name, value: calls.append((name, value))) + + eplb_metrics.publish_expert_load_metrics( + metric_client=metric_client, + expert_load=torch.tensor([[384, 128, 128, 128]]), + ) + + assert calls == [ + (eplb_metrics.EXPERT_IMBALANCE_RATIO_METRICS[25], pytest.approx(2.0)), + (eplb_metrics.EXPERT_IMBALANCE_RATIO_METRICS[50], pytest.approx(2.0)), + (eplb_metrics.EXPERT_IMBALANCE_RATIO_METRICS[100], pytest.approx(2.0)), + ] + + +def test_publish_rebalance_compute_metrics_from_sample_load(): + calls = [] + metric_client = SimpleNamespace(gauge_set=lambda name, value: calls.append((name, value))) + + eplb_metrics.publish_rebalance_compute_metrics( + metric_client=metric_client, + sample_load=torch.tensor([[384, 128, 128, 128]]), + current_placement=[[[0, 1, 2], [1, 2, 3]]], + target_placement=[[[0, 1, 2], [0, 2, 3]]], + expert_alignment=128, + ) + + assert calls == [ + ( + eplb_metrics.COMPUTE_CRITICAL_OVERHEAD_RATIO_BEFORE_REBALANCE_METRIC, + pytest.approx(0.25), + ), + ( + eplb_metrics.COMPUTE_CRITICAL_OVERHEAD_RATIO_AFTER_REBALANCE_METRIC, + pytest.approx(0.0), + ), + ] + + +def test_eplb_prefill_route_counter_has_24_samples_per_logical_expert(monkeypatch): + args = type( + "Args", + (), + { + "eplb_num_redundant_experts_per_rank": 2, + "eplb_config_path": None, + }, + )() + monkeypatch.setattr(deepgemm_module, "get_env_start_args", lambda: args) + monkeypatch.setattr(deepgemm_module, "get_global_world_size", lambda: 2) + monkeypatch.setattr(deepgemm_module, "get_global_rank", lambda: 0) + monkeypatch.setattr(deepgemm_module, "get_node_world_size", lambda: 2) + monkeypatch.setattr(torch.Tensor, "cuda", lambda tensor: tensor) + original_zeros = torch.zeros + + def cpu_zeros(*shape, **kwargs): + kwargs.pop("device", None) + return original_zeros(*shape, **kwargs) + + monkeypatch.setattr(deepgemm_module.torch, "zeros", cpu_zeros) + + impl = deepgemm_module.FuseMoeDeepGEMM(4, 0, 1.0, SimpleNamespace()) + + assert impl.prefill_route_counter.shape == (24, 4) + assert impl.prefill_route_sample_index.shape == (2,) + assert torch.equal(impl.prefill_route_sample_index, torch.zeros(2, dtype=torch.int64)) + + +def test_ep_without_eplb_creates_layout_without_eplb_runtime_state(monkeypatch): + args = type( + "Args", + (), + {"eplb_num_redundant_experts_per_rank": 0}, + )() + monkeypatch.setattr(deepgemm_module, "get_env_start_args", lambda: args) + monkeypatch.setattr(deepgemm_module, "get_global_world_size", lambda: 2) + monkeypatch.setattr(deepgemm_module, "get_global_rank", lambda: 0) + + impl = deepgemm_module.FuseMoeDeepGEMM(4, 0, 1.0, SimpleNamespace()) + + assert impl.num_redundant_experts_per_rank == 0 + assert impl.num_total_physical_experts == impl.n_routed_experts + assert impl.local_logics_expert_ids_list == [0, 1] + assert not hasattr(impl, "num_primary_experts_per_rank") + assert not hasattr(impl, "initial_local_expert_ids_by_rank") + assert not hasattr(impl, "logical_to_physical_map") + assert not hasattr(impl, "prefill_route_counter") + assert not hasattr(impl, "prefill_route_sample_index") + assert not hasattr(impl, "recording") + + +def test_manager_wait_load_gather_aggregates_rank_and_sample_dimensions(monkeypatch): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager._eplb_impls = [object()] + manager.global_rank = 0 + manager.world_size = 4 + manager.num_logical_experts = 4 + manager.control_group = object() + manager.state = manager_module.EPLBManagerState.WAIT_LOAD_GATHER_FINISHED + gathered_load = torch.zeros((4, 1, 24, 4), dtype=torch.int64) + gathered_load[0, 0, 0].fill_(100) + gathered_load[1, 0, 1].fill_(100) + gathered_load[2, 0, 0].fill_(50) + manager._load_gather_task = SimpleNamespace( + is_finished=lambda: True, + result=gathered_load, + ) + seen = {} + + def all_gather_object(output, local_finished, **kwargs): + seen["local_finished"] = local_finished + seen["group"] = kwargs["group"] + output[:] = [True] * manager.world_size + + monkeypatch.setattr(manager_module.dist, "all_gather_object", all_gather_object) + + manager._step_wait_load_gather_finished() + + assert seen["group"] is manager.control_group + assert seen["local_finished"] is True + assert manager.state is manager_module.EPLBManagerState.PLAN_PLACEMENT + assert manager._pending_plan_load_samples is gathered_load + # metrics 只汇总各 rank 的 sample 0;rank 1 在 sample 1 中的负载 + # 不应累加进来。 + assert torch.equal(manager._metric_sample_load, torch.full((1, 4), 150, dtype=torch.int64)) + assert manager._load_gather_task is None + + +def test_manager_wait_load_gather_does_not_advance_until_every_rank_finishes(monkeypatch): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.state = manager_module.EPLBManagerState.WAIT_LOAD_GATHER_FINISHED + manager.world_size = 2 + manager.control_group = object() + load_gather_task = SimpleNamespace( + is_finished=lambda: True, + result=torch.ones((2, 1, 1, 2), dtype=torch.int64), + ) + manager._load_gather_task = load_gather_task + monkeypatch.setattr( + manager_module.dist, + "all_gather_object", + lambda output, _local_finished, **_kwargs: output.__setitem__(slice(None), [True, False]), + ) + + manager._step_wait_load_gather_finished() + + assert manager.state is manager_module.EPLBManagerState.WAIT_LOAD_GATHER_FINISHED + assert manager._load_gather_task is load_gather_task + + +def test_decode_dispatch_uses_physical_ids_and_total_expert_count(monkeypatch): + class Buffer: + def low_latency_dispatch(self, **kwargs): + calls.append(kwargs) + return "recv", "masked", "handle", "event", "hook" + + impl = object.__new__(deepgemm_module.FuseMoeDeepGEMM) + impl.quant_method = type("Quant", (), {"method_name": "fp8"})() + impl.n_routed_experts = 128 + _set_deepgemm_runtime(impl, _test_moe_impl(eplb=True)) + logical_ids = torch.tensor([[0, 127]], dtype=torch.int32) + physical_ids = torch.tensor([[128, 143]], dtype=torch.int32) + impl._select_experts = lambda **_kwargs: ( + torch.ones((1, 2)), + logical_ids, + ) + calls, repairs = [], [] + + def repair(**kwargs): + repairs.append(kwargs) + return physical_ids + + monkeypatch.setattr(deepgemm_module, "eplb_repair_topk_ids", repair) + monkeypatch.setattr( + deepgemm_module, + "get_deepep_num_max_dispatch_tokens_per_rank_decode", + lambda: 16, + ) + monkeypatch.setattr(deepgemm_module.dist_group_manager, "ep_low_latency_buffer", Buffer()) + + result = impl.low_latency_dispatch( + torch.empty((1, 4)), + torch.empty((1, 128)), + None, + False, + 2, + False, + 0, + 0, + "softmax", + ) + + assert result[2].tolist() == [[128, 143]] + assert repairs[0]["logical_topk_ids"] is logical_ids + assert repairs[0]["mode"] == "current_gpu_first" + assert calls[0]["num_experts"] == 144 + + +def test_select_returns_logical_ids_and_applies_expert_scale(monkeypatch): + from lightllm.common.basemodel.triton_kernel.fused_moe import topk_select + + impl = object.__new__(deepgemm_module.FuseMoeDeepGEMM) + impl.routed_scaling_factor = 2.0 + _set_deepgemm_runtime(impl, _test_moe_impl(eplb=True)) + logical_ids = torch.tensor([[3, 4]], dtype=torch.int32) + calls = [] + + def select(**kwargs): + calls.append(kwargs) + return torch.tensor([[0.5, 0.25]]), logical_ids + + monkeypatch.setattr(topk_select, "select_experts", select) + weights, selected = impl._select_experts( + torch.empty((1, 4)), + torch.empty((1, 128)), + None, + 2, + False, + False, + 0, + 0, + "softmax", + per_expert_scale=torch.tensor([1.0, 1.0, 1.0, 3.0, 5.0]), + ) + + assert len(calls) == 1 + assert weights.tolist() == [[3.0, 2.5]] + assert selected is logical_ids + + +def test_eplb_prefill_repairs_ids_after_selection(monkeypatch): + impl = object.__new__(deepgemm_module.FuseMoeDeepGEMM) + impl.routed_scaling_factor = 1.0 + impl.quant_method = object() + _set_deepgemm_runtime(impl, _test_moe_impl(eplb=True)) + logical_ids = torch.tensor([[3, 4]], dtype=torch.int32) + physical_ids = torch.tensor([[130, 131]], dtype=torch.long) + calls = [] + impl._select_experts = lambda **_kwargs: (torch.ones((1, 2)), logical_ids) + + def repair(**kwargs): + calls.append(kwargs) + return physical_ids + + monkeypatch.setattr(deepgemm_module, "eplb_repair_topk_ids", repair) + monkeypatch.setattr(deepgemm_module, "quantize_fused_experts_input", lambda *_args: "qinput") + + weights, topk_idx, qinput = impl.select_experts_and_quant_input( + torch.empty((1, 4)), + torch.empty((1, 128)), + None, + object(), + False, + 2, + False, + 0, + 0, + "softmax", + ) + + assert weights.tolist() == [[1.0, 1.0]] + assert topk_idx is physical_ids + assert topk_idx.dtype is torch.long + assert qinput == "qinput" + assert calls[0]["logical_topk_ids"] is logical_ids + assert not calls[0]["update_prefill_route_counter"] + assert calls[0]["mode"] == "global_first" + + +def test_eplb_prefill_dispatch_consumes_physical_ids_and_event(monkeypatch): + class Buffer: + def dispatch(self, _qinput, **kwargs): + calls.append(kwargs) + return ( + (torch.empty((4, 2)),), + "recv_idx", + "recv_weight", + SimpleNamespace(num_recv_tokens_per_expert_list=[4]), + SimpleNamespace(current_stream_wait=lambda: None), + ) + + impl = object.__new__(deepgemm_module.FuseMoeDeepGEMM) + impl.routed_scaling_factor = 1.0 + impl.quant_method = object() + runtime = _test_moe_impl( + eplb=True, + prefill_route_counter=torch.zeros((24, 128), dtype=torch.int64), + recording=True, + ) + _set_deepgemm_runtime(impl, runtime) + calls, repair_calls = [], [] + logical_ids = torch.tensor([[3, 4]], dtype=torch.int32) + physical_ids = torch.tensor([[130, 131]], dtype=torch.long) + impl._select_experts = lambda **_kwargs: (torch.ones((1, 2)), logical_ids) + + def repair(**kwargs): + repair_calls.append(kwargs) + return physical_ids + + monkeypatch.setattr(deepgemm_module, "eplb_repair_topk_ids", repair) + monkeypatch.setattr(deepgemm_module, "quantize_fused_experts_input", lambda *_args: "qinput") + monkeypatch.setattr(deepgemm_module.dist_group_manager, "ep_buffer", Buffer()) + monkeypatch.setattr( + deepgemm_module, + "get_deepep_num_max_dispatch_tokens_per_rank_prefill", + lambda: 16, + ) + monkeypatch.setattr(deepgemm_module, "get_ep_num_sms", lambda: 8) + + weights, topk_idx, qinput = impl.select_experts_and_quant_input( + torch.empty((1, 4)), + torch.empty((1, 128)), + None, + object(), + True, + 2, + False, + 1, + 8, + "sigmoid", + ) + caller_event = object() + impl.dispatch( + qinput, + topk_idx, + weights, + overlap_event=caller_event, + ) + + assert topk_idx is physical_ids + assert len(repair_calls) == 1 + assert repair_calls[0]["logical_topk_ids"] is logical_ids + assert repair_calls[0]["update_prefill_route_counter"] + assert repair_calls[0]["mode"] == "global_first" + assert calls[0]["topk_idx"] is physical_ids + assert calls[0]["topk_idx"].dtype is torch.long + assert calls[0]["previous_event"] is caller_event + + +def test_prefill_dispatch_preserves_event(monkeypatch): + class Buffer: + def dispatch(self, _qinput, **kwargs): + calls.append(kwargs) + return ( + (torch.empty((4, 2)),), + "recv_idx", + "recv_weight", + SimpleNamespace(num_recv_tokens_per_expert_list=[4]), + SimpleNamespace(current_stream_wait=lambda: None), + ) + + impl = object.__new__(deepgemm_module.FuseMoeDeepGEMM) + _set_deepgemm_runtime(impl, _test_moe_impl(eplb=True)) + calls = [] + caller_event = object() + monkeypatch.setattr(deepgemm_module.dist_group_manager, "ep_buffer", Buffer()) + monkeypatch.setattr( + deepgemm_module, + "get_deepep_num_max_dispatch_tokens_per_rank_prefill", + lambda: 16, + ) + monkeypatch.setattr(deepgemm_module, "get_ep_num_sms", lambda: 8) + + impl.dispatch( + "qinput", + torch.tensor([[1, 2]], dtype=torch.long), + torch.ones((1, 2)), + caller_event, + ) + + assert calls[0]["previous_event"] is caller_event + assert calls[0]["topk_idx"].dtype is torch.long + + +def test_deepgemm_constructor_owns_eplb_runtime(monkeypatch): + monkeypatch.setattr( + deepgemm_module, + "get_env_start_args", + lambda: SimpleNamespace( + eplb_num_redundant_experts_per_rank=1, + eplb_config_path=None, + ), + ) + monkeypatch.setattr(deepgemm_module, "get_global_world_size", lambda: 2) + monkeypatch.setattr(deepgemm_module, "get_global_rank", lambda: 0) + monkeypatch.setattr(deepgemm_module, "get_node_world_size", lambda: 2) + monkeypatch.setattr(torch.Tensor, "cuda", lambda tensor: tensor) + original_zeros = torch.zeros + + def cpu_zeros(*shape, **kwargs): + kwargs.pop("device", None) + return original_zeros(*shape, **kwargs) + + monkeypatch.setattr(deepgemm_module.torch, "zeros", cpu_zeros) + impl = deepgemm_module.FuseMoeDeepGEMM(4, 0, 1.0, SimpleNamespace()) + + assert impl.num_redundant_experts_per_rank == 1 + assert impl.num_total_physical_experts == 6 + assert impl.prefill_route_counter.shape == (24, 4) + assert impl.prefill_route_sample_index.shape == (2,) + assert impl.recording + assert impl.local_logics_expert_ids_list == [0, 1, 2] + assert not hasattr(impl, "initial_local_expert_ids_by_rank") + assert not hasattr(impl, "expert_parallel_state") + + +def test_deepgemm_constructor_loads_saved_layout_before_weight_initialization(monkeypatch): + saved_placement = [[1, 0, 3], [2, 3, 1]] + monkeypatch.setattr( + deepgemm_module, + "get_env_start_args", + lambda: SimpleNamespace( + eplb_num_redundant_experts_per_rank=1, + eplb_config_path="/tmp/eplb.json", + ), + ) + monkeypatch.setattr(deepgemm_module, "get_global_world_size", lambda: 2) + monkeypatch.setattr(deepgemm_module, "get_global_rank", lambda: 0) + monkeypatch.setattr(deepgemm_module, "get_node_world_size", lambda: 2) + monkeypatch.setattr(torch.Tensor, "cuda", lambda tensor: tensor) + monkeypatch.setattr( + deepgemm_module, + "load_layer_placement", + lambda path, **kwargs: ( + saved_placement + if path == "/tmp/eplb.json" + and kwargs + == { + "layer_index": 7, + "num_logical_experts": 4, + "world_size": 2, + "num_redundant_experts_per_rank": 1, + } + else None + ), + ) + original_zeros = torch.zeros + monkeypatch.setattr( + deepgemm_module.torch, + "zeros", + lambda *shape, **kwargs: original_zeros(*shape, dtype=kwargs.get("dtype")), + ) + + impl = deepgemm_module.FuseMoeDeepGEMM(4, 0, 1.0, SimpleNamespace(), layer_index=7) + + assert impl.local_logics_expert_ids_list == saved_placement[0] + expected_map = build_logical_to_physical_map(saved_placement, 4, current_rank=0, node_world_size=2) + assert impl.logical_to_physical_map.tolist() == expected_map + + +def test_deepgemm_keeps_route_recording_when_rebalance_count_is_zero(monkeypatch): + monkeypatch.setattr( + deepgemm_module, + "get_env_start_args", + lambda: SimpleNamespace( + eplb_num_redundant_experts_per_rank=1, + eplb_rebalance_count=0, + eplb_config_path=None, + ), + ) + monkeypatch.setattr(deepgemm_module, "get_global_world_size", lambda: 2) + monkeypatch.setattr(deepgemm_module, "get_global_rank", lambda: 0) + monkeypatch.setattr(deepgemm_module, "get_node_world_size", lambda: 2) + monkeypatch.setattr(torch.Tensor, "cuda", lambda tensor: tensor) + monkeypatch.setattr( + deepgemm_module.torch, + "zeros", + lambda *shape, **kwargs: torch.full(shape, 0, dtype=kwargs.get("dtype")), + ) + + impl = deepgemm_module.FuseMoeDeepGEMM(4, 0, 1.0, SimpleNamespace()) + + assert impl.recording + + +def test_eplb_prepare_repairs_logical_ids(monkeypatch): + impl = object.__new__(deepgemm_module.FuseMoeDeepGEMM) + runtime = _test_moe_impl(eplb=True, recording=True) + _set_deepgemm_runtime(impl, runtime) + logical_ids = torch.tensor([[3, 4]], dtype=torch.int32) + physical_ids = torch.tensor([[13, 14]], dtype=torch.int32) + calls = [] + + def repair(**kwargs): + calls.append(kwargs) + return physical_ids + + monkeypatch.setattr(deepgemm_module, "eplb_repair_topk_ids", repair) + weights, selected = impl._prepare_expert_execution(torch.ones((1, 2)), logical_ids, is_prefill=True) + + assert weights.tolist() == [[1.0, 1.0]] + assert selected is physical_ids + assert calls[0]["logical_topk_ids"] is logical_ids + assert calls[0]["update_prefill_route_counter"] + assert calls[0]["mode"] == "global_first" + + +def test_decode_masked_group_gemm_uses_all_physical_rows_when_eplb_is_enabled( + monkeypatch, +): + impl = object.__new__(deepgemm_module.FuseMoeDeepGEMM) + _set_deepgemm_runtime(impl, _test_moe_impl(eplb=True, num_logical_experts=8, world_size=1)) + captured = {} + + def masked(*args, **kwargs): + captured["w13"] = args[3] + captured["w13_scale"] = args[4] + captured["w2"] = args[5] + captured["w2_scale"] = args[6] + return "out" + + monkeypatch.setattr(deepgemm_module, "masked_group_gemm", masked) + pack = lambda: type( + "Pack", + (), + {"weight": torch.empty((10, 4)), "weight_scale": torch.empty((10, 1))}, + )() + + assert impl.masked_group_gemm((torch.empty((1, 4)),), pack(), pack(), torch.empty(8), torch.float16, 1) == "out" + assert captured["w13"].shape[0] == captured["w2"].shape[0] == 10 + assert captured["w13_scale"].shape[0] == captured["w2_scale"].shape[0] == 10 + + +def test_decode_fused_experts_uses_full_weight_packs_and_physical_experts( + monkeypatch, +): + impl = object.__new__(deepgemm_module.FuseMoeDeepGEMM) + impl.n_routed_experts = 128 + _set_deepgemm_runtime( + impl, + _test_moe_impl( + eplb=True, + num_logical_experts=128, + world_size=16, + num_redundant_experts_per_rank=2, + ), + ) + impl.quant_method = object() + captured = [] + + def fused(**kwargs): + captured.append(kwargs) + return "out" + + monkeypatch.setattr(deepgemm_module, "fused_experts", fused) + pack = lambda: type( + "Pack", + (), + { + "weight": torch.empty((10, 4)), + "weight_scale": torch.empty((10, 1)), + "weight_zero_point": None, + }, + )() + w13, w2 = pack(), pack() + + for _ in range(2): + assert ( + impl._fused_experts( + torch.empty((1, 4)), + w13, + w2, + torch.ones((1, 2)), + torch.zeros((1, 2), dtype=torch.int64), + is_prefill=False, + ) + == "out" + ) + + assert [call["num_experts"] for call in captured] == [160, 160] + assert all(call["w13"] is w13 and call["w2"] is w2 for call in captured) + + +def test_transfer_plan_uses_stable_current_expert_source(): + current = [[0, 1, 4, 5], [2, 3, 6, 7], [4, 5, 0, 1], [6, 7, 2, 3]] + target = [[0, 1, 6, 5], [2, 3, 6, 7], [4, 5, 0, 4], [6, 7, 2, 3]] + plan = build_transfer_plan(current, target, 5, num_logical_experts=8, world_size=4) + assert plan == [ + [ + EPLBTransferInfo(5, 6, 1, 2, 0, 2), + EPLBTransferInfo(5, 4, 2, 0, 2, 3), + ], + ] + + +def test_transfer_plan_reuses_stable_source_for_repeated_expert(): + current = [[0, 1, 0, 1], [2, 3, 2, 3], [4, 5, 4, 5], [6, 7, 4, 7]] + target = [[0, 1, 4, 4], [2, 3, 2, 3], [4, 5, 4, 5], [6, 7, 4, 7]] + first = build_transfer_plan(current, target, 5, 8, 4) + second = build_transfer_plan(current, target, 5, 8, 4) + assert first == second + assert first == [ + [EPLBTransferInfo(5, 4, 2, 0, 0, 2)], + [EPLBTransferInfo(5, 4, 2, 2, 0, 3)], + ] + + +def test_transfer_plan_keeps_primary_slot_swap_in_one_atomic_batch(): + current = [[0, 1], [2, 3]] + target = [[2, 1], [0, 3]] + + plan = build_transfer_plan(current, target, 0, num_logical_experts=4, world_size=2) + + assert plan == [ + [ + EPLBTransferInfo(0, 2, 1, 0, 0, 0), + EPLBTransferInfo(0, 0, 0, 0, 1, 0), + ] + ] + + +def test_transfer_plan_keeps_three_way_cycle_in_one_atomic_batch(): + current = [[0], [1], [2]] + target = [[1], [2], [0]] + + plan = build_transfer_plan(current, target, 0, num_logical_experts=3, world_size=3) + + assert plan == [ + [ + EPLBTransferInfo(0, 1, 1, 0, 0, 0), + EPLBTransferInfo(0, 0, 0, 0, 2, 0), + EPLBTransferInfo(0, 2, 2, 0, 1, 0), + ] + ] + + +def test_p2p_message_tag_is_stable_and_identifies_transfer_tensor(): + transfer = object.__new__(PinnedMemoryEPLBTransfer) + transfer.transfer_info = EPLBTransferInfo(5, 4, 1, 0, 0, 2) + weight_tag = transfer._build_p2p_message_tag("w13.weight") + + assert weight_tag == transfer._build_p2p_message_tag("w13.weight") + assert 0 <= weight_tag <= 0x7FFFFFFF + assert weight_tag != transfer._build_p2p_message_tag("w13.weight_scale") + transfer.transfer_info = EPLBTransferInfo(5, 6, 1, 0, 0, 2) + assert weight_tag != transfer._build_p2p_message_tag("w13.weight") + transfer.transfer_info = EPLBTransferInfo(5, 4, 1, 1, 0, 2) + assert weight_tag != transfer._build_p2p_message_tag("w13.weight") + transfer.transfer_info = EPLBTransferInfo(5, 4, 1, 0, 0, 3) + assert weight_tag != transfer._build_p2p_message_tag("w13.weight") + + +def test_extract_expert_tensors_includes_quantization_metadata_in_order(): + class Pack: + def __init__(self, offset, scale=True, zero_point=True): + self.weight = torch.full((3, 2), offset) + self.weight_scale = torch.full((3, 1), offset + 1) if scale else None + self.weight_zero_point = torch.full((3, 1), offset + 2) if zero_point else None + + weight = type("Weight", (), {"w13": Pack(1), "w2": Pack(10, scale=False, zero_point=False)})() + tensors = extract_eplb_expert_tensors(weight) + assert [name for name, _ in tensors] == [ + "w13.weight", + "w13.weight_scale", + "w13.weight_zero_point", + "w2.weight", + ] + + +def test_manager_commits_transfer_rows_and_metadata(monkeypatch): + original_copy = torch.Tensor.copy_ + non_blocking_values = [] + copy_sources = [] + + def record_copy(tensor, source, non_blocking=False): + non_blocking_values.append(non_blocking) + copy_sources.append(source) + return original_copy(tensor, source, non_blocking=non_blocking) + + monkeypatch.setattr(torch.Tensor, "copy_", record_copy) + live = torch.arange(20).reshape(5, 4) + original_primary = live[:3].clone() + local_expert_ids = [0, 1, 2, 3, 2] + target_placement = [[[0, 1, 2, 4, 5], [3, 4, 5, 0, 2]]] + expected_metadata = torch.tensor( + build_logical_to_physical_map( + target_placement[0], + 6, + current_rank=0, + node_world_size=2, + ), + dtype=torch.int32, + ) + logical_to_physical_map = torch.zeros_like(expected_metadata) + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.global_rank = 0 + manager.world_size = 2 + manager.node_world_size = 2 + manager.num_logical_experts = 6 + manager.target_placement = target_placement + manager.current_placement = [[[0, 1, 2, 3, 2], [3, 4, 5, 0, 2]]] + manager._eplb_impls = [ + SimpleNamespace( + local_logics_expert_ids_list=local_expert_ids, + logical_to_physical_map=logical_to_physical_map, + ) + ] + transfers = [ + SimpleNamespace( + transfer_info=EPLBTransferInfo(0, 4, 1, 1, 0, 3), + tensor_buffers=[ExpertTensorBuffer("weight", live, torch.full((4,), -4))], + ), + SimpleNamespace( + transfer_info=EPLBTransferInfo(0, 5, 1, 2, 0, 4), + tensor_buffers=[ExpertTensorBuffer("weight", live, torch.full((4,), -5))], + ), + ] + manager.active_transfers = [transfers[0]] + manager._commit_transfer(transfers[0].transfer_info) + assert manager.current_placement[0][0] == [0, 1, 2, 4, 2] + manager.active_transfers = [transfers[1]] + manager._commit_transfer(transfers[1].transfer_info) + manager._publish_layer_metadata(0) + + assert torch.equal(live[:3], original_primary) + assert torch.equal(live[3], torch.full((4,), -4)) + assert torch.equal(live[4], torch.full((4,), -5)) + assert local_expert_ids == [0, 1, 2, 4, 5] + assert torch.equal(logical_to_physical_map, expected_metadata) + assert non_blocking_values == [True, True, True] + assert copy_sources[-1].is_pinned() + + +def test_manager_transfers_only_local_tasks_and_gathers_global_status(monkeypatch): + class Transfer: + def __init__(self, transfer_info): + self.transfer_info = transfer_info + self.finished = finished_by_info[transfer_info] + + def start(self): + starts.append(self.transfer_info) + + def is_finished(self): + return self.finished + + remote_info = EPLBTransferInfo(0, 2, 0, 0, 2, 2) + local_info0 = EPLBTransferInfo(0, 3, 1, 1, 3, 2) + local_info1 = EPLBTransferInfo(1, 4, 0, 0, 1, 2) + finished_by_info = {local_info0: False, local_info1: True} + starts = [] + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.control_group = object() + manager.transfer_group = object() + manager._weights = [object(), object()] + manager.world_size = 4 + manager.pending_transfer_batches = [[remote_info, local_info0], [local_info1]] + manager.target_placement = [ + [[0, 1, 2], [2, 3, 3], [4, 5, 2], [6, 7, 0]], + [[0, 1, 4], [2, 3, 5], [4, 5, 6], [6, 7, 1]], + ] + manager.current_placement = [ + [[0, 1, 4], [2, 3, 5], [4, 5, 6], [6, 7, 0]], + [[0, 1, 2], [2, 3, 4], [4, 5, 6], [6, 7, 0]], + ] + manager.global_rank = 1 + manager.state = manager_module.EPLBManagerState.TRANSFERRING + manager.max_rebalance_count = -1 + manager.completed_rebalance_count = 0 + manager.rebalance_started_at = None + manager.active_transfer_batch = None + manager.active_transfers = None + committed = [] + active_streams = [] + cleared_route_counters = [] + + def commit_transfer(transfer_info): + assert active_streams == [overlap_stream] + committed.append(transfer_info) + + def publish_layer_metadata(_layer_index): + assert active_streams == [overlap_stream] + + manager._commit_transfer = commit_transfer + manager._publish_layer_metadata = publish_layer_metadata + manager._clear_prefill_route_samples = lambda: cleared_route_counters.append(True) + used_streams = [] + overlap_stream = object() + + class StreamContext: + def __init__(self, stream): + self.stream = stream + + def __enter__(self): + used_streams.append(self.stream) + active_streams.append(self.stream) + + def __exit__(self, *_args): + active_streams.pop() + + monkeypatch.setattr( + manager_module.torch.cuda, + "stream", + StreamContext, + ) + monkeypatch.setattr(g_infer_context, "get_overlap_stream", lambda: overlap_stream) + monkeypatch.setattr( + manager_module, + "PinnedMemoryEPLBTransfer", + lambda **kwargs: Transfer(kwargs["transfer_info"]), + ) + + gathered_states = [ + [True, False, True, False], + [True, True, True, True], + [True, True, True, True], + ] + local_states = [] + + def all_gather_object(output, local_state, **_kwargs): + local_states.append(local_state) + output[:] = gathered_states.pop(0) + + monkeypatch.setattr(manager_module.dist, "all_gather_object", all_gather_object) + + manager._step_transferring() + assert starts == [local_info0] + assert committed == [] + assert manager.active_transfer_batch == [remote_info, local_info0] + assert manager.pending_transfer_batches == [[local_info1]] + + manager._step_transferring() + assert committed == [] + + manager.active_transfers[0].finished = True + manager._step_transferring() + assert committed == [remote_info, local_info0] + assert manager.active_transfers is None + + manager._step_transferring() + assert starts == [local_info0, local_info1] + + manager._step_transferring() + assert committed == [remote_info, local_info0, local_info1] + + manager._step_transferring() + assert manager.state is manager_module.EPLBManagerState.COLLECTING + assert manager.completed_rebalance_count == 1 + assert manager.current_placement == [ + [[0, 1, 2], [2, 3, 3], [4, 5, 2], [6, 7, 0]], + [[0, 1, 4], [2, 3, 5], [4, 5, 6], [6, 7, 1]], + ] + assert manager.pending_transfer_batches is None + assert manager.target_placement is None + assert manager.rebalance_started_at is None + assert cleared_route_counters == [True] + assert local_states == [False, True, True] + assert used_streams == [overlap_stream, overlap_stream] + + +def test_manager_returns_to_collecting_after_reaching_rebalance_limit(): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + impls = [SimpleNamespace(recording=True), SimpleNamespace(recording=True)] + target_placement = [[[0, 1], [1, 0]]] + manager.state = manager_module.EPLBManagerState.TRANSFERRING + manager.global_rank = 0 + manager._eplb_impls = impls + manager.current_placement = [[[0, 1], [0, 1]]] + manager.target_placement = target_placement + manager.pending_transfer_batches = [] + manager.max_rebalance_count = 1 + manager.completed_rebalance_count = 0 + manager.rebalance_started_at = None + manager.active_transfer_batch = None + manager.active_transfers = None + manager._clear_prefill_route_samples = lambda: None + persisted_placements = [] + manager._persist_current_placement = lambda: persisted_placements.append(manager.current_placement) + + manager._step_transferring() + + assert manager.current_placement is target_placement + assert manager.completed_rebalance_count == 1 + assert manager.state is manager_module.EPLBManagerState.COLLECTING + assert persisted_placements == [target_placement] + assert all(impl.recording for impl in impls) + + +def test_wait_plan_finish_broadcasts_pending_status(monkeypatch): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.state = manager_module.EPLBManagerState.WAIT_PLAN_PLACEMENT_FINISHED + manager.global_rank = 0 + manager.control_group = object() + manager._plan_task = SimpleNamespace( + is_finished=lambda: False, + result=None, + ) + broadcasts = [] + + def broadcast(values, **_kwargs): + broadcasts.append(values[0]) + + monkeypatch.setattr(manager_module.dist, "broadcast_object_list", broadcast) + manager._step_wait_plan_placement_finished() + assert broadcasts == [None] + assert manager.state is manager_module.EPLBManagerState.WAIT_PLAN_PLACEMENT_FINISHED + + +def test_wait_plan_finish_publishes_before_and_after_metrics(monkeypatch): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.state = manager_module.EPLBManagerState.WAIT_PLAN_PLACEMENT_FINISHED + manager.global_rank = 0 + manager.control_group = object() + manager.current_placement = [[[0, 1], [2, 3]]] + target_placement = [[[0, 2], [1, 3]]] + manager._pending_plan_load_samples = None + manager._metric_sample_load = torch.tensor([[4, 3, 2, 1]], dtype=torch.int64) + manager._plan_task = SimpleNamespace( + is_finished=lambda: True, + result=target_placement, + ) + manager.metric_client = object() + published = [] + monkeypatch.setattr( + eplb_metrics, + "publish_rebalance_compute_metrics", + lambda **kwargs: published.append((kwargs["sample_load"], kwargs["target_placement"])), + ) + monkeypatch.setattr(manager_module.dist, "broadcast_object_list", lambda _values, **_kwargs: None) + + manager._step_wait_plan_placement_finished() + + assert manager.state is manager_module.EPLBManagerState.PLAN_TRANSFER + assert manager.target_placement is target_placement + assert len(published) == 1 + assert torch.equal(published[0][0], torch.tensor([[4, 3, 2, 1]], dtype=torch.int64)) + assert published[0][1] is target_placement + assert manager._pending_plan_load_samples is None + assert manager._metric_sample_load is None + assert manager._plan_task is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_manager_transfer_task_commit_orders_live_weights_between_overlap_forwards( + monkeypatch, +): + class Transfer: + def __init__(self, live, received, transfer_info): + self.tensor_buffers = [ExpertTensorBuffer("weight", live, received)] + self.transfer_info = transfer_info + self.status = "succeeded" + + def is_finished(self): + return True + + live = torch.tensor([1.0], device="cuda") + received = torch.tensor(2.0, pin_memory=True) + previous_read = torch.empty_like(live) + next_read = torch.empty_like(live) + source_stream = torch.cuda.Stream(device=live.device) + destination_stream = torch.cuda.Stream(device=live.device) + initial_stream = torch.cuda.current_stream(device=live.device) + original_overlap_stream = g_infer_context.overlap_stream + + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + transfer_info = EPLBTransferInfo(0, 0, 0, 0, 0, 0) + transfer = Transfer(live, received, transfer_info) + manager.active_transfers = [transfer] + manager.active_transfer_batch = [transfer_info] + manager.control_group = object() + manager.world_size = 1 + manager.node_world_size = 1 + manager.pending_transfer_batches = [] + manager.num_logical_experts = 1 + manager.global_rank = 0 + manager.target_placement = [[[0]]] + manager._eplb_impls = [ + SimpleNamespace( + local_logics_expert_ids_list=[0], + logical_to_physical_map=torch.zeros((1, 1), dtype=torch.int32, device="cuda"), + ) + ] + manager.current_placement = [[[0]]] + manager.rebalance_started_at = time.time() + manager.state = manager_module.EPLBManagerState.TRANSFERRING + monkeypatch.setattr( + manager_module.dist, + "all_gather_object", + lambda output, local_ready, **_kwargs: output.__setitem__(slice(None), [local_ready]), + ) + monkeypatch.setattr(manager_module, "build_logical_to_physical_map", lambda *_args, **_kwargs: [[0]]) + + try: + g_infer_context.overlap_stream = source_stream + with torch.cuda.stream(source_stream): + source_stream.wait_stream(initial_stream) + torch.cuda._sleep(20_000_000) + previous_read.copy_(live, non_blocking=True) + with torch.cuda.stream(destination_stream): + manager._step_transferring() + with torch.cuda.stream(source_stream): + source_stream.wait_stream(destination_stream) + next_read.copy_(live, non_blocking=True) + source_stream.synchronize() + + assert previous_read.item() == 1.0 + assert next_read.item() == 2.0 + finally: + g_infer_context.overlap_stream = original_overlap_stream + + +def test_manager_step_advances_inflight_transfer(): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.state = manager_module.EPLBManagerState.TRANSFERRING + calls = [] + manager._step_transferring = lambda: calls.append("transfer") + + manager.step() + + assert calls == ["transfer"] + + +def test_manager_evaluates_only_after_entering_evaluating_state(monkeypatch): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + prefill_route_counter = torch.tensor([[1, 2]], dtype=torch.int64) + load_gather_tasks = [] + manager.state = manager_module.EPLBManagerState.COLLECTING + manager.global_rank = 1 + manager.steps = 0 + manager.step_interval = 3 + manager.next_evaluation_step = 3 + manager.num_logical_experts = 2 + manager._eplb_impls = [SimpleNamespace(prefill_route_counter=prefill_route_counter)] + manager.world_size = 1 + manager.control_group = object() + manager.max_rebalance_count = -1 + manager.completed_rebalance_count = 0 + monkeypatch.setattr( + manager_module, + "EPLBLoadGatherTask", + lambda **_kwargs: load_gather_tasks.append(True), + ) + monkeypatch.setattr( + manager_module.dist, + "broadcast_object_list", + lambda values, **_kwargs: values.__setitem__(0, False), + ) + + manager.step() + manager.step() + assert manager.state is manager_module.EPLBManagerState.COLLECTING + manager.step() + + assert manager.state is manager_module.EPLBManagerState.EVALUATING + assert load_gather_tasks == [] + assert manager.next_evaluation_step == 6 + + manager.step() + + assert manager.state is manager_module.EPLBManagerState.COLLECTING + assert load_gather_tasks == [] + assert manager.next_evaluation_step == 6 + assert torch.equal(prefill_route_counter, torch.tensor([[1, 2]], dtype=torch.int64)) + + +def test_manager_step_uses_explicit_state_instead_of_pending_work(): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.state = manager_module.EPLBManagerState.TRANSFERRING + manager.pending_transfer_batches = [[object()]] + manager._plan_task = object() + calls = [] + manager._step_transferring = lambda: calls.append("transfer") + manager._step_evaluating = lambda: calls.append("evaluation") + + manager.step() + + assert calls == ["transfer"] + + +def test_manager_plans_transfers_asynchronously_before_entering_transferring(monkeypatch): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.state = manager_module.EPLBManagerState.WAIT_PLAN_PLACEMENT_FINISHED + manager.global_rank = 1 + manager.control_group = object() + manager.transfer_group = object() + manager.world_size = 2 + manager.num_logical_experts = 4 + manager.pending_transfer_batches = None + placement = [ + [[0, 1, 2], [2, 3, 3]], + [[0, 1, 3], [2, 3, 0]], + ] + manager.current_placement = [ + [[0, 1, 3], [2, 3, 2]], + [[0, 1, 2], [2, 3, 1]], + ] + monkeypatch.setattr( + manager_module.dist, + "broadcast_object_list", + lambda values, **_kwargs: values.__setitem__(0, placement), + ) + + transfer_infos = [ + EPLBTransferInfo(0, 2, 0, 0, 1, 2), + EPLBTransferInfo(1, 0, 0, 0, 1, 2), + ] + transfer_planners = [] + + class TransferPlanner: + def __init__(self, **kwargs): + self.current = kwargs["current_placement"] + self.target = kwargs["target_placement"] + self.num_logical_experts = kwargs["num_logical_experts"] + self.world_size = kwargs["world_size"] + self.result = [[transfer_infos[0]], [transfer_infos[1]]] + self.started = False + self.finished = False + transfer_planners.append(self) + + def start(self): + self.started = True + + def is_finished(self): + return self.finished + + monkeypatch.setattr(manager_module, "EPLBTransferPlanner", TransferPlanner) + monkeypatch.setattr( + manager_module, + "PinnedMemoryEPLBTransfer", + lambda **_kwargs: pytest.fail("transfer object must not be built while planning transfers"), + ) + + manager._step_wait_plan_placement_finished() + + assert manager.state is manager_module.EPLBManagerState.PLAN_TRANSFER + assert manager.target_placement is placement + assert transfer_planners == [] + assert manager.pending_transfer_batches is None + + manager.step() + + assert manager.state is manager_module.EPLBManagerState.WAIT_PLAN_TRANSFER_FINISHED + assert len(transfer_planners) == 1 + transfer_planner = transfer_planners[0] + assert transfer_planner.current is manager.current_placement + assert transfer_planner.target is placement + assert transfer_planner.num_logical_experts == manager.num_logical_experts + assert transfer_planner.world_size == manager.world_size + assert transfer_planner.started + + remote_finished = False + + def gather_finished(output, local_finished, **_kwargs): + output[:] = [local_finished, remote_finished] + + monkeypatch.setattr(manager_module.dist, "all_gather_object", gather_finished) + transfer_planner.finished = True + manager.step() + + assert manager.state is manager_module.EPLBManagerState.WAIT_PLAN_TRANSFER_FINISHED + assert manager.pending_transfer_batches is None + + remote_finished = True + manager.step() + + assert manager.state is manager_module.EPLBManagerState.TRANSFERRING + assert manager.pending_transfer_batches == [[transfer_infos[0]], [transfer_infos[1]]] + assert manager._transfer_planner is None + + +def test_manager_evaluation_with_insufficient_tokens_returns_to_collecting(monkeypatch): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.state = manager_module.EPLBManagerState.EVALUATING + manager._eplb_impls = [SimpleNamespace(prefill_route_counter=torch.full((1, 4), 100, dtype=torch.int64))] + manager.num_logical_experts = 4 + manager.steps = 11 + manager.step_interval = 20 + manager.next_evaluation_step = 31 + manager.global_rank = 0 + manager.world_size = 1 + manager.control_group = object() + manager.load_gather_group = object() + manager.max_rebalance_count = -1 + manager.completed_rebalance_count = 0 + manager.metric_client = object() + broadcast_decisions = [] + monkeypatch.setattr(eplb_metrics, "publish_expert_load_metrics", lambda **_kwargs: None) + monkeypatch.setattr( + manager_module, + "EPLBLoadGatherTask", + lambda **_kwargs: pytest.fail("insufficient rank-0 load must skip all-gather"), + ) + monkeypatch.setattr( + manager_module.dist, + "broadcast_object_list", + lambda values, **kwargs: broadcast_decisions.append((values[0], kwargs["src"], kwargs["group"])), + ) + + manager.step() + + assert manager.state is manager_module.EPLBManagerState.COLLECTING + assert manager.next_evaluation_step == 31 + assert broadcast_decisions == [(False, 0, manager.control_group)] + + +def test_manager_evaluation_with_enough_tokens_enters_planning(monkeypatch): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + local_load = torch.full((1, 4), 256, dtype=torch.int64) + plan_tasks = [] + published_loads = [] + manager.state = manager_module.EPLBManagerState.EVALUATING + manager.global_rank = 0 + manager.num_logical_experts = 4 + manager._eplb_impls = [SimpleNamespace(prefill_route_counter=local_load)] + manager.world_size = 1 + manager.control_group = object() + manager.load_gather_group = object() + manager.max_rebalance_count = -1 + manager.completed_rebalance_count = 0 + + class LoadGatherTask: + def __init__(self, **kwargs): + self.local_load_samples = kwargs["local_load_samples"] + self.result = self.local_load_samples.unsqueeze(0) + self.started = False + + def start(self): + self.started = True + + def is_finished(self): + return True + + monkeypatch.setattr(manager_module, "EPLBLoadGatherTask", LoadGatherTask) + monkeypatch.setattr(manager_module.dist, "broadcast_object_list", lambda _values, **_kwargs: None) + monkeypatch.setattr( + manager_module.dist, + "all_gather_object", + lambda output, local_finished, **_kwargs: output.__setitem__(slice(None), [local_finished]), + ) + + class PlanTask: + def __init__(self, **kwargs): + self.planner = kwargs["planner"] + self.logical_expert_load_samples = kwargs["logical_expert_load_samples"] + self.current_placement = kwargs["current_placement"] + self.started = False + plan_tasks.append(self) + + def start(self): + self.started = True + + manager.planner = object() + manager.current_placement = [[[0, 1, 2, 3]]] + manager.metric_client = object() + monkeypatch.setattr( + eplb_metrics, + "publish_expert_load_metrics", + lambda **kwargs: published_loads.append(kwargs["expert_load"]), + ) + monkeypatch.setattr(manager_module, "EPLBPlanTask", PlanTask) + + manager.step() + + assert manager.state is manager_module.EPLBManagerState.WAIT_LOAD_GATHER_FINISHED + assert manager._load_gather_task.started + assert torch.equal(manager._load_gather_task.local_load_samples, local_load.unsqueeze(0)) + assert len(published_loads) == 1 + assert torch.equal(published_loads[0], local_load) + assert plan_tasks == [] + assert getattr(manager, "_plan_task", None) is None + + manager.step() + + assert manager.state is manager_module.EPLBManagerState.PLAN_PLACEMENT + planning_load_samples = manager._pending_plan_load_samples + assert torch.equal(planning_load_samples, local_load.unsqueeze(0).unsqueeze(0)) + assert torch.equal(manager._metric_sample_load, local_load) + assert manager._load_gather_task is None + assert plan_tasks == [] + + manager.step() + + assert manager.state is manager_module.EPLBManagerState.WAIT_PLAN_PLACEMENT_FINISHED + assert plan_tasks[0].logical_expert_load_samples is planning_load_samples + assert torch.equal(plan_tasks[0].logical_expert_load_samples, local_load.unsqueeze(0).unsqueeze(0)) + assert plan_tasks[0].planner is manager.planner + assert plan_tasks[0].current_placement is manager.current_placement + assert plan_tasks[0].started + assert manager._plan_task is plan_tasks[0] + assert manager._pending_plan_load_samples is None + assert torch.equal(manager._metric_sample_load, local_load) + assert len(published_loads) == 1 + + +def test_manager_keeps_reporting_after_reaching_rebalance_limit(monkeypatch): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + local_load = torch.tensor([[10, 20, 30, 40]], dtype=torch.int64) + published_loads = [] + cleared_counters = [] + manager.state = manager_module.EPLBManagerState.EVALUATING + manager.global_rank = 0 + manager.num_logical_experts = 4 + manager._eplb_impls = [SimpleNamespace(prefill_route_counter=local_load)] + manager.max_rebalance_count = 1 + manager.completed_rebalance_count = 1 + manager.metric_client = object() + monkeypatch.setattr( + eplb_metrics, + "publish_expert_load_metrics", + lambda **kwargs: published_loads.append(kwargs["expert_load"]), + ) + manager._clear_prefill_route_samples = lambda: cleared_counters.append(True) + monkeypatch.setattr( + manager_module, + "EPLBLoadGatherTask", + lambda **_kwargs: pytest.fail("load gathering must stop after reaching the limit"), + ) + + manager.step() + + assert manager.state is manager_module.EPLBManagerState.COLLECTING + assert len(published_loads) == 1 + assert torch.equal(published_loads[0], local_load) + assert cleared_counters == [True] + + +def test_nonzero_rank_enters_plan_wait_without_starting_planner(): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.state = manager_module.EPLBManagerState.PLAN_PLACEMENT + manager.global_rank = 1 + + manager.step() + + assert manager.state is manager_module.EPLBManagerState.WAIT_PLAN_PLACEMENT_FINISHED + assert getattr(manager, "_plan_task", None) is None + assert getattr(manager, "_pending_plan_load_samples", None) is None + assert getattr(manager, "_metric_sample_load", None) is None + + +def test_manager_planning_without_changes_returns_to_collecting(monkeypatch): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.state = manager_module.EPLBManagerState.WAIT_PLAN_PLACEMENT_FINISHED + manager.global_rank = 1 + manager.control_group = object() + manager.steps = 11 + manager.step_interval = 20 + manager.next_evaluation_step = 31 + manager.current_placement = [[[0, 1], [1, 0]]] + result = None + + def broadcast(values, **_kwargs): + values[0] = result + + monkeypatch.setattr(manager_module.dist, "broadcast_object_list", broadcast) + + manager.step() + assert manager.state is manager_module.EPLBManagerState.WAIT_PLAN_PLACEMENT_FINISHED + + result = [[[0, 1], [1, 0]]] + manager.step() + + assert manager.state is manager_module.EPLBManagerState.COLLECTING + assert manager.next_evaluation_step == 31 + assert getattr(manager, "_plan_task", None) is None + + +def test_manager_exposes_one_lifecycle_step_entrypoint(): + assert hasattr(manager_module.EPLBManager, "step") + assert not hasattr(manager_module.EPLBManager, "poll") + + +def test_pinned_transfer_copies_source_row_and_sends_to_destination(monkeypatch): + class Stream: + def __init__(self): + self.synchronize_count = 0 + + def synchronize(self): + self.synchronize_count += 1 + + transfer = object.__new__(PinnedMemoryEPLBTransfer) + transfer._device = "cuda:0" + transfer._is_source_rank = True + transfer._is_destination_rank = False + transfer._p2p_group = object() + transfer.transfer_info = EPLBTransferInfo(0, 5, 0, 1, 1, 2) + transfer._device_to_host_stream = Stream() + transfer.tensor_buffers = [ + ExpertTensorBuffer( + "weight", + torch.tensor([[1.0, 2.0], [3.0, 4.0]]), + torch.empty(2), + ) + ] + transfer.status = "running" + sends = [] + monkeypatch.setattr(transfer_module.torch.cuda, "set_device", lambda _device: None) + monkeypatch.setattr(transfer_module.torch.cuda, "stream", lambda _stream: nullcontext()) + monkeypatch.setattr( + transfer_module.dist, + "send", + lambda tensor, dst, group, tag: sends.append((tensor.clone(), dst, group, tag)), + ) + + transfer._run() + + assert transfer.status == "succeeded" + assert len(sends) == 1 + assert torch.equal(sends[0][0], torch.tensor([3.0, 4.0])) + expected_tag = transfer._build_p2p_message_tag("weight") + assert sends[0][1:] == (1, transfer._p2p_group, expected_tag) + assert torch.equal(transfer.tensor_buffers[0].pinned_row, torch.tensor([3.0, 4.0])) + assert transfer._device_to_host_stream.synchronize_count == 1 + + +def test_pinned_transfer_skips_p2p_for_local_destination(monkeypatch): + class Stream: + def synchronize(self): + pass + + transfer = object.__new__(PinnedMemoryEPLBTransfer) + transfer._device = "cuda:0" + transfer._is_source_rank = True + transfer._is_destination_rank = True + transfer._p2p_group = object() + transfer.transfer_info = EPLBTransferInfo(0, 5, 0, 0, 0, 1) + transfer._device_to_host_stream = Stream() + transfer.tensor_buffers = [ + ExpertTensorBuffer( + "weight", + torch.tensor([[3.0, 4.0]]), + torch.empty(2), + ) + ] + transfer.status = "running" + p2p_calls = [] + monkeypatch.setattr(transfer_module.torch.cuda, "set_device", lambda _device: None) + monkeypatch.setattr(transfer_module.torch.cuda, "stream", lambda _stream: nullcontext()) + monkeypatch.setattr(transfer_module.dist, "send", lambda *_args, **_kwargs: p2p_calls.append("send")) + monkeypatch.setattr(transfer_module.dist, "recv", lambda *_args, **_kwargs: p2p_calls.append("recv")) + + transfer._run() + + assert transfer.is_finished() + assert p2p_calls == [] + assert torch.equal(transfer.tensor_buffers[0].pinned_row, torch.tensor([3.0, 4.0])) + + +def test_pinned_transfer_exits_process_on_failure(monkeypatch): + transfer = object.__new__(PinnedMemoryEPLBTransfer) + transfer._device = "cuda:0" + transfer._is_source_rank = False + transfer._is_destination_rank = True + transfer.transfer_info = EPLBTransferInfo(0, 3, 1, 0, 0, 2) + transfer._p2p_group = object() + transfer.tensor_buffers = [ + ExpertTensorBuffer( + "weight", + torch.empty((1, 1)), + torch.empty(1), + ) + ] + transfer.status = "running" + logged_messages = [] + exit_codes = [] + + def fail_recv(*_args, **_kwargs): + raise RuntimeError("recv failed") + + monkeypatch.setattr(transfer_module.torch.cuda, "set_device", lambda _device: None) + monkeypatch.setattr(transfer_module.dist, "recv", fail_recv) + monkeypatch.setattr(async_task_module.logger, "exception", logged_messages.append) + monkeypatch.setattr(async_task_module.os, "_exit", exit_codes.append) + + transfer._run() + + assert logged_messages == ["PinnedMemoryEPLBTransfer failed"] + assert exit_codes == [1] + assert not transfer.is_finished() + + +def test_pinned_transfer_is_single_use_and_exposes_pinned_rows(monkeypatch): + class Stream: + def synchronize(self): + pass + + transfer = object.__new__(PinnedMemoryEPLBTransfer) + transfer._device = "cuda:0" + transfer._is_source_rank = False + transfer._is_destination_rank = True + transfer.transfer_info = EPLBTransferInfo(0, 3, 1, 0, 0, 2) + transfer._p2p_group = object() + transfer._device_to_host_stream = Stream() + transfer.tensor_buffers = [ + ExpertTensorBuffer( + "weight", + torch.empty((1, 1)), + torch.tensor([3.0]), + ) + ] + transfer.status = "idle" + transfer._thread = threading.Thread(target=transfer._run, daemon=True) + receives = [] + monkeypatch.setattr(transfer_module.torch.cuda, "set_device", lambda _device: None) + monkeypatch.setattr( + transfer_module.dist, + "recv", + lambda tensor, src, group, tag: receives.append((tensor, src, group, tag)), + ) + + assert not transfer.is_finished() + transfer.start() + deadline = time.monotonic() + 2 + while transfer.status == "running" and time.monotonic() < deadline: + time.sleep(0.001) + assert transfer.status == "succeeded" + assert transfer.is_finished() + assert torch.equal(transfer.tensor_buffers[0].pinned_row, torch.tensor([3.0])) + assert len(receives) == 1 + assert receives[0][0] is transfer.tensor_buffers[0].pinned_row + expected_tag = transfer._build_p2p_message_tag("weight") + assert receives[0][1:] == (1, transfer._p2p_group, expected_tag) + with pytest.raises(AssertionError, match="already been started"): + transfer.start() + + +def test_manager_requires_more_than_one_rank(monkeypatch): + monkeypatch.setattr(manager_module, "is_sm100_gpu", lambda: False) + monkeypatch.setattr(manager_module, "_find_fused_moe_weights", lambda model: [object()]) + monkeypatch.setattr(manager_module, "get_global_rank", lambda: 0) + monkeypatch.setattr(manager_module, "get_global_world_size", lambda: 1) + + with pytest.raises(AssertionError, match="more than one rank"): + manager_module.EPLBManager(type("Model", (), {})()) + + +def test_manager_rejects_sm100_before_initialization(monkeypatch): + monkeypatch.setattr(manager_module, "is_sm100_gpu", lambda: True) + + with pytest.raises(AssertionError, match="EPLB does not support SM100"): + manager_module.EPLBManager(type("Model", (), {})()) + + +def test_manager_clears_all_prefill_route_samples_on_overlap_stream(monkeypatch): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + counters = [torch.tensor([[1, 2]]), torch.tensor([[3, 4]])] + sample_indices = [torch.tensor([7, 3]), torch.tensor([9, 4])] + manager._eplb_impls = [ + SimpleNamespace(prefill_route_counter=counter, prefill_route_sample_index=sample_index) + for counter, sample_index in zip(counters, sample_indices) + ] + overlap_stream = object() + used_streams = [] + monkeypatch.setattr(g_infer_context, "get_overlap_stream", lambda: overlap_stream) + monkeypatch.setattr( + manager_module.torch.cuda, + "stream", + lambda stream: (used_streams.append(stream) or nullcontext()), + ) + + manager._clear_prefill_route_samples() + + assert used_streams == [overlap_stream] + assert all(torch.count_nonzero(counter) == 0 for counter in counters) + assert all(torch.count_nonzero(sample_index) == 0 for sample_index in sample_indices) + + +def test_manager_initializes_without_transfer_task(monkeypatch): + weight = type( + "Weight", + (), + { + "n_routed_experts": 4, + "layer_num_": 0, + "fuse_moe_impl": _test_moe_impl( + eplb=True, + recording=True, + num_logical_experts=4, + world_size=2, + num_redundant_experts_per_rank=2, + prefill_route_counter=torch.zeros((24, 4), dtype=torch.int64), + ), + }, + )() + groups = [object(), object(), object()] + new_group_calls = [] + monkeypatch.setattr(manager_module, "is_sm100_gpu", lambda: False) + monkeypatch.setattr(manager_module, "_find_fused_moe_weights", lambda model: [weight]) + monkeypatch.setattr(manager_module, "get_global_rank", lambda: 0) + monkeypatch.setattr(manager_module, "get_global_world_size", lambda: 2) + monkeypatch.setattr(manager_module, "get_node_world_size", lambda: 2) + monkeypatch.setattr(manager_module, "get_eplb_step_interval", lambda: 20) + clear_calls = [] + monkeypatch.setattr( + manager_module.EPLBManager, + "_clear_prefill_route_samples", + lambda manager: clear_calls.append(manager), + ) + monkeypatch.setattr(manager_module, "get_shm_port_args", lambda: SimpleNamespace(metric_port=1234)) + metric_client = SimpleNamespace() + metric_client_ports = [] + monkeypatch.setattr( + manager_module, + "MetricClient", + lambda port: (metric_client_ports.append(port) or metric_client), + ) + + def new_group(*args, **kwargs): + new_group_calls.append((args, kwargs)) + return groups[len(new_group_calls) - 1] + + monkeypatch.setattr(manager_module.dist, "new_group", new_group) + all_gather_calls = [] + + def all_gather_object(output, local_expert_ids_by_layer, group): + all_gather_calls.append((local_expert_ids_by_layer, group)) + output[:] = [local_expert_ids_by_layer, [[2, 3, 0, 1]]] + + monkeypatch.setattr(manager_module.dist, "all_gather_object", all_gather_object) + logs = [] + monkeypatch.setattr(manager_module.logger, "info", lambda message: logs.append(message)) + monkeypatch.setattr( + manager_module, + "save_placement_config", + lambda *_args, **_kwargs: pytest.fail("manager initialization must not save the placement"), + ) + manager = manager_module.EPLBManager(type("Model", (), {})(), config_path="/tmp/eplb.json") + assert manager._plan_task is None + assert manager.pending_transfer_batches is None + assert manager.state is manager_module.EPLBManagerState.COLLECTING + assert (manager.load_gather_group, manager.control_group, manager.transfer_group) == tuple(groups) + assert new_group_calls == [(([0, 1],), {"backend": "gloo"})] * 3 + assert all_gather_calls == [([[0, 1, 2, 3]], groups[1])] + assert manager.current_placement == [[[0, 1, 2, 3], [2, 3, 0, 1]]] + assert manager.metric_client is metric_client + assert metric_client_ports == [1234] + assert manager.next_evaluation_step == manager.step_interval + assert manager.max_rebalance_count == 1 + assert manager.completed_rebalance_count == 0 + assert manager.plan_mode == "greedy" + assert clear_calls == [manager] + assert isinstance(manager.planner, GreedyEPLBPlanner) + assert "plan_mode=greedy" in logs[0] + assert "planner=GreedyEPLBPlanner" in logs[0] + assert weight.fuse_moe_impl.recording + assert manager._eplb_impls[0] is weight.fuse_moe_impl + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for the Triton EPLB kernel") +@pytest.mark.parametrize("update_prefill_route_counter", [False, True]) +@pytest.mark.parametrize("tokens", [1, 32]) +@pytest.mark.parametrize("mode", ["current_gpu_first", "current_node_first", "global_first"]) +def test_eplb_repair_topk_ids_maps_and_counts(update_prefill_route_counter, tokens, mode): + from lightllm.common.basemodel.triton_kernel.fused_moe.eplb_topk_ids import ( + eplb_repair_topk_ids, + ) + + topk = 4 + experts = 64 + logical_ids = (torch.arange(tokens * topk, dtype=torch.int32, device="cuda") % experts).view(tokens, topk) + original_logical_ids = logical_ids.clone() + logical_experts = torch.arange(experts, dtype=torch.int32, device="cuda") + replica_counts = torch.where( + logical_experts % 3 == 0, + torch.full_like(logical_experts, 3), + torch.ones_like(logical_experts), + ) + num_current_gpu_replicas = torch.where( + logical_experts % 5 == 0, + torch.ones_like(logical_experts), + torch.zeros_like(logical_experts), + ) + num_node_replicas = torch.where( + num_current_gpu_replicas > 0, + torch.where( + replica_counts >= 2, + torch.full_like(logical_experts, 2), + num_current_gpu_replicas, + ), + torch.where( + (logical_experts % 7 == 0) & (replica_counts == 3), + torch.full_like(logical_experts, 2), + torch.zeros_like(logical_experts), + ), + ) + logical_to_physical = torch.stack( + ( + replica_counts, + num_node_replicas, + num_current_gpu_replicas, + logical_experts, + torch.where(replica_counts >= 2, logical_experts + experts, logical_experts), + torch.where(replica_counts == 3, logical_experts + 2 * experts, logical_experts), + ), + dim=1, + ) + counter = torch.zeros((24, experts), dtype=torch.int64, device="cuda") + sample_index = torch.zeros((2,), dtype=torch.int64, device="cuda") + expected_counter = torch.zeros_like(counter) + + logical_ids_long = logical_ids.to(torch.long) + token_indices = torch.arange(tokens, device="cuda", dtype=torch.int64).unsqueeze(1) + if mode == "current_gpu_first": + num_preferred_replicas = torch.where( + num_current_gpu_replicas > 0, + num_current_gpu_replicas, + replica_counts, + ) + elif mode == "current_node_first": + num_preferred_replicas = torch.where(num_node_replicas > 0, num_node_replicas, replica_counts) + else: + num_preferred_replicas = replica_counts + hash_values = token_indices ^ ((logical_ids.to(torch.int64) + 1) * 0x9E3779B9) + hash_values &= 0xFFFFFFFF + hash_values ^= hash_values >> 16 + hash_values = (hash_values * 0x7FEB352D) & 0xFFFFFFFF + hash_values ^= hash_values >> 15 + hash_values = (hash_values * 0x846CA68B) & 0xFFFFFFFF + hash_values ^= hash_values >> 16 + replica_indices = hash_values % num_preferred_replicas[logical_ids_long].to(torch.int64) + expected_ids = logical_to_physical[logical_ids_long, replica_indices + 3] + if update_prefill_route_counter: + expected_counter[0].scatter_add_( + 0, + logical_ids.reshape(-1).to(torch.long), + torch.ones(logical_ids.numel(), dtype=torch.int64, device="cuda"), + ) + + physical_ids = eplb_repair_topk_ids( + logical_topk_ids=logical_ids, + logical_to_physical_map=logical_to_physical, + prefill_route_counter=counter, + prefill_route_sample_index=sample_index, + update_prefill_route_counter=update_prefill_route_counter, + mode=mode, + ) + torch.cuda.synchronize() + + assert torch.equal(logical_ids, original_logical_ids) + assert torch.equal(physical_ids, expected_ids) + assert torch.equal(counter, expected_counter) + assert sample_index.tolist() == ([1, 0] if update_prefill_route_counter else [0, 0]) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for the Triton EPLB kernel") +def test_eplb_prefill_route_counter_multigrid_ring_wrap_is_exact(): + from lightllm.common.basemodel.triton_kernel.fused_moe.eplb_topk_ids import ( + eplb_repair_topk_ids, + ) + + capacity = 24 + experts = 256 + tokens = 1025 + topk = 4 + num_samples = capacity + 2 + base_ids = torch.arange(tokens * topk, dtype=torch.int32, device="cuda").view(tokens, topk) + logical_experts = torch.arange(experts, dtype=torch.int32, device="cuda") + logical_to_physical = torch.stack( + ( + torch.ones_like(logical_experts), + torch.ones_like(logical_experts), + torch.ones_like(logical_experts), + logical_experts, + ), + dim=1, + ) + counter = torch.zeros((capacity, experts), dtype=torch.int64, device="cuda") + sample_index = torch.zeros((2,), dtype=torch.int64, device="cuda") + expected = torch.zeros_like(counter) + + for sample in range(num_samples): + logical_ids = (base_ids + sample) % experts + physical_ids = eplb_repair_topk_ids( + logical_topk_ids=logical_ids, + logical_to_physical_map=logical_to_physical, + prefill_route_counter=counter, + prefill_route_sample_index=sample_index, + update_prefill_route_counter=True, + mode="global_first", + ) + expected[sample % capacity] = torch.bincount( + logical_ids.reshape(-1).to(torch.long), + minlength=experts, + ) + assert torch.equal(physical_ids, logical_ids) + + torch.cuda.synchronize() + + assert sample_index.tolist() == [num_samples, 0] + assert torch.equal(counter, expected) + assert int(counter.sum().item()) == capacity * tokens * topk + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for the Triton EPLB kernel") +def test_eplb_repair_topk_ids_spreads_strided_expert_tokens(): + from lightllm.common.basemodel.triton_kernel.fused_moe.eplb_topk_ids import ( + eplb_repair_topk_ids, + ) + + num_tokens = 4096 + token_indices = torch.arange(num_tokens, dtype=torch.int32, device="cuda") + logical_ids = torch.where(token_indices % 4 == 0, 0, 1).view(-1, 1) + logical_to_physical = torch.tensor( + [ + [4, 4, 4, 0, 2, 3, 4], + [1, 1, 1, 1, -1, -1, -1], + ], + dtype=torch.int32, + device="cuda", + ) + counter = torch.zeros((24, 2), dtype=torch.int64, device="cuda") + sample_index = torch.zeros((2,), dtype=torch.int64, device="cuda") + + physical_ids = eplb_repair_topk_ids( + logical_topk_ids=logical_ids, + logical_to_physical_map=logical_to_physical, + prefill_route_counter=counter, + prefill_route_sample_index=sample_index, + update_prefill_route_counter=False, + mode="global_first", + ) + torch.cuda.synchronize() + + strided_token_outputs = physical_ids[token_indices % 4 == 0, 0] + replica_counts = torch.stack([(strided_token_outputs == physical_id).sum() for physical_id in (0, 2, 3, 4)]) + assert torch.all(replica_counts > strided_token_outputs.numel() * 0.2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for the Triton EPLB kernel") +@pytest.mark.parametrize("num_replicas", [2, 3, 4, 5, 101, 127, 128, 251]) +def test_eplb_replica_hash_is_uniform_across_tokens_and_experts(num_replicas): + from lightllm.common.basemodel.triton_kernel.fused_moe.eplb_topk_ids import ( + eplb_repair_topk_ids, + ) + + num_tokens = 4097 + num_experts = 256 + expert_ids = torch.arange(num_experts, dtype=torch.int32, device="cuda") + logical_ids = expert_ids.expand(num_tokens, -1).contiguous() + replica_counts = torch.full((num_experts, 3), num_replicas, dtype=torch.int32, device="cuda") + physical_ids = expert_ids.unsqueeze(1) + num_experts * torch.arange( + num_replicas, + dtype=torch.int32, + device="cuda", + ) + logical_to_physical = torch.cat((replica_counts, physical_ids), dim=1) + counter = torch.zeros((24, num_experts), dtype=torch.int64, device="cuda") + sample_index = torch.zeros((2,), dtype=torch.int64, device="cuda") + + routed_physical_ids = eplb_repair_topk_ids( + logical_topk_ids=logical_ids, + logical_to_physical_map=logical_to_physical, + prefill_route_counter=counter, + prefill_route_sample_index=sample_index, + update_prefill_route_counter=False, + mode="global_first", + ) + torch.cuda.synchronize() + + routed_replica_indices = routed_physical_ids // num_experts + observed_counts = torch.stack( + [(routed_replica_indices == replica_index).sum(dim=0) for replica_index in range(num_replicas)], + dim=1, + ) + expected_count = num_tokens / num_replicas + deviations = observed_counts - expected_count + if num_replicas <= 5: + max_relative_deviation = (deviations.abs() / expected_count).max().item() + assert max_relative_deviation < 0.12 + else: + # 副本很多时单槽期望样本较少,使用每个 expert 的归一化卡方值 + # 检查整体形状,并额外检查跨 expert 汇总后的单槽最大偏差。 + normalized_chi_square = (deviations.square() / expected_count).sum(dim=1) / (num_replicas - 1) + assert normalized_chi_square.max().item() < 1.75 + + aggregate_expected_count = num_tokens * num_experts / num_replicas + aggregate_max_relative_deviation = ( + (observed_counts.sum(dim=0) - aggregate_expected_count).abs() / aggregate_expected_count + ).max() + assert aggregate_max_relative_deviation.item() < 0.06 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for the Triton EPLB kernel") +def test_eplb_repair_topk_ids_empty_input_skips_kernel(): + from lightllm.common.basemodel.triton_kernel.fused_moe.eplb_topk_ids import ( + eplb_repair_topk_ids, + ) + + experts = 64 + logical_ids = torch.empty((0, 4), dtype=torch.int32, device="cuda") + counter = torch.zeros((24, experts), dtype=torch.int64, device="cuda") + sample_index = torch.zeros((2,), dtype=torch.int64, device="cuda") + logical_to_physical = torch.stack( + ( + torch.ones((experts,), dtype=torch.int32, device="cuda"), + torch.ones((experts,), dtype=torch.int32, device="cuda"), + torch.ones((experts,), dtype=torch.int32, device="cuda"), + torch.arange(experts, dtype=torch.int32, device="cuda"), + ), + dim=1, + ) + physical_ids = eplb_repair_topk_ids( + logical_topk_ids=logical_ids, + logical_to_physical_map=logical_to_physical, + prefill_route_counter=counter, + prefill_route_sample_index=sample_index, + update_prefill_route_counter=True, + mode="current_gpu_first", + ) + + assert physical_ids.shape == (0, 4) + assert physical_ids.dtype is torch.int32 + assert torch.equal(counter, torch.zeros_like(counter)) + assert torch.equal(sample_index, torch.zeros_like(sample_index)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for the Triton EPLB kernel") +def test_eplb_repair_topk_ids_rejects_unknown_dispatch_mode(): + from lightllm.common.basemodel.triton_kernel.fused_moe.eplb_topk_ids import ( + eplb_repair_topk_ids, + ) + + logical_ids = torch.empty((0, 1), dtype=torch.int32, device="cuda") + logical_to_physical = torch.tensor([[1, 1, 1, 0]], dtype=torch.int32, device="cuda") + counter = torch.zeros((24, 1), dtype=torch.int64, device="cuda") + sample_index = torch.zeros((2,), dtype=torch.int64, device="cuda") + + with pytest.raises(AssertionError, match="unsupported EPLB dispatch mode"): + eplb_repair_topk_ids( + logical_topk_ids=logical_ids, + logical_to_physical_map=logical_to_physical, + prefill_route_counter=counter, + prefill_route_sample_index=sample_index, + update_prefill_route_counter=False, + mode="unknown", + ) diff --git a/unit_tests/common/fused_moe/test_eplb_placement_config.py b/unit_tests/common/fused_moe/test_eplb_placement_config.py new file mode 100644 index 0000000000..509431abec --- /dev/null +++ b/unit_tests/common/fused_moe/test_eplb_placement_config.py @@ -0,0 +1,202 @@ +import json + +from lightllm.server.router.model_infer.mode_backend.eplb.placement import ( + load_layer_placement, + save_placement_config, +) +from lightllm.server.router.model_infer.mode_backend.eplb.placement import config as config_module + + +def test_placement_config_round_trip(tmp_path): + config_path = tmp_path / "nested" / "eplb-placement.json" + placement = [ + [[0, 1, 2], [2, 3, 0]], + [[1, 0, 3], [2, 3, 1]], + ] + + assert save_placement_config( + str(config_path), + layer_indexes=[3, 7], + placement=placement, + num_logical_experts=4, + world_size=2, + num_redundant_experts_per_rank=1, + ) + assert ( + load_layer_placement( + str(config_path), + layer_index=7, + num_logical_experts=4, + world_size=2, + num_redundant_experts_per_rank=1, + ) + == placement[1] + ) + + saved_config = json.loads(config_path.read_text(encoding="utf-8")) + assert saved_config == { + "version": 1, + "num_logical_experts": 4, + "world_size": 2, + "num_redundant_experts_per_rank": 1, + "layers": {"3": placement[0], "7": placement[1]}, + } + assert not (config_path.parent / f"{config_path.name}.lock").exists() + + +def test_placement_config_lock_prevents_concurrent_write(tmp_path, monkeypatch): + config_path = tmp_path / "eplb-placement.json" + lock_path = tmp_path / "eplb-placement.json.lock" + config_path.write_text("original", encoding="utf-8") + lock_path.write_text("another-process", encoding="utf-8") + warnings = [] + monkeypatch.setattr(config_module.logger, "warning", lambda *args: warnings.append(args)) + + assert not save_placement_config( + str(config_path), + layer_indexes=[3], + placement=[[[0, 1, 2], [2, 3, 0]]], + num_logical_experts=4, + world_size=2, + num_redundant_experts_per_rank=1, + ) + assert config_path.read_text(encoding="utf-8") == "original" + assert lock_path.read_text(encoding="utf-8") == "another-process" + assert len(warnings) == 1 + + +def test_missing_placement_config_warns_and_falls_back(tmp_path, monkeypatch): + config_path = tmp_path / "missing.json" + warnings = [] + config_module._read_config.cache_clear() + monkeypatch.setattr(config_module.logger, "warning", lambda *args: warnings.append(args)) + + assert ( + load_layer_placement( + str(config_path), + layer_index=3, + num_logical_experts=4, + world_size=2, + num_redundant_experts_per_rank=1, + ) + is None + ) + assert len(warnings) == 1 + assert "using the default initial placement" in warnings[0][0] + + +def test_malformed_json_warns_and_falls_back(tmp_path, monkeypatch): + config_path = tmp_path / "malformed.json" + config_path.write_text("{not-json", encoding="utf-8") + warnings = [] + config_module._read_config.cache_clear() + monkeypatch.setattr(config_module.logger, "warning", lambda *args: warnings.append(args)) + + assert ( + load_layer_placement( + str(config_path), + layer_index=3, + num_logical_experts=4, + world_size=2, + num_redundant_experts_per_rank=1, + ) + is None + ) + assert len(warnings) == 1 + assert "using the default initial placement" in warnings[0][0] + + +def test_invalid_placement_config_warns_and_falls_back(tmp_path, monkeypatch): + config_path = tmp_path / "invalid.json" + config_path.write_text( + json.dumps( + { + "version": 1, + "num_logical_experts": 4, + "world_size": 2, + "num_redundant_experts_per_rank": 1, + "layers": {"3": [[0, 1, 1], [2, 3, 0]]}, + } + ), + encoding="utf-8", + ) + warnings = [] + config_module._read_config.cache_clear() + monkeypatch.setattr(config_module.logger, "warning", lambda *args: warnings.append(args)) + + assert ( + load_layer_placement( + str(config_path), + layer_index=3, + num_logical_experts=4, + world_size=2, + num_redundant_experts_per_rank=1, + ) + is None + ) + assert len(warnings) == 1 + assert "contains duplicate expert IDs" in str(warnings[0][3]) + + +def test_missing_layer_warns_and_falls_back(tmp_path, monkeypatch): + config_path = tmp_path / "missing-layer.json" + config_path.write_text( + json.dumps( + { + "version": 1, + "num_logical_experts": 4, + "world_size": 2, + "num_redundant_experts_per_rank": 1, + "layers": {}, + } + ), + encoding="utf-8", + ) + warnings = [] + config_module._read_config.cache_clear() + monkeypatch.setattr(config_module.logger, "warning", lambda *args: warnings.append(args)) + + assert ( + load_layer_placement( + str(config_path), + layer_index=3, + num_logical_experts=4, + world_size=2, + num_redundant_experts_per_rank=1, + ) + is None + ) + assert len(warnings) == 1 + assert "is missing" in str(warnings[0][3]) + + +def test_topology_mismatch_warns_and_falls_back(tmp_path, monkeypatch): + config_path = tmp_path / "mismatch.json" + config_path.write_text( + json.dumps( + { + "version": 1, + "num_logical_experts": 8, + "world_size": 2, + "num_redundant_experts_per_rank": 1, + "layers": {}, + } + ), + encoding="utf-8", + ) + warnings = [] + config_module._read_config.cache_clear() + monkeypatch.setattr(config_module.logger, "warning", lambda *args: warnings.append(args)) + + assert ( + load_layer_placement( + str(config_path), + layer_index=3, + num_logical_experts=4, + world_size=2, + num_redundant_experts_per_rank=1, + ) + is None + ) + assert len(warnings) == 1 + assert "does not match the current deployment" in warnings[0][0] diff --git a/unit_tests/common/fused_moe/test_eplb_transfer_gpu.py b/unit_tests/common/fused_moe/test_eplb_transfer_gpu.py new file mode 100644 index 0000000000..61a5448ea3 --- /dev/null +++ b/unit_tests/common/fused_moe/test_eplb_transfer_gpu.py @@ -0,0 +1,345 @@ +"""Multi-GPU correctness test for pinned-memory EPLB transfers.""" + +import gc +import os +import socket +import time +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from lightllm.server.router.model_infer.mode_backend.eplb.async_expert_transfer import ( + EPLBTransferInfo, + PinnedMemoryEPLBTransfer, + build_transfer_plan, +) + + +class _Pack: + def __init__(self, weight, weight_scale): + self.weight = weight + self.weight_scale = weight_scale + self.weight_zero_point = None + + +class _FakeWeight: + def __init__(self, rank, layer_index): + self.layer_num_ = layer_index + logical_ids = ([0, 1, 2], [2, 3, 0])[rank] + self.fuse_moe_impl = SimpleNamespace( + num_redundant_experts_per_rank=1, + local_logics_expert_ids_list=list(logical_ids), + ) + self.w13 = self._pack(logical_ids, layer_index, 0) + self.w2 = self._pack(logical_ids, layer_index, 10) + + @staticmethod + def _pack(logical_ids, layer_index, offset): + values = torch.tensor( + [[layer_index * 100 + expert + offset] for expert in logical_ids], + dtype=torch.float16, + device="cuda", + ) + scales = values.to(torch.float32) + 0.5 + return _Pack(values, scales) + + +class _StressFakeWeight: + """为并发传输压力测试生成可识别 source rank 的专家权重。""" + + def __init__(self, rank, layer_index, num_logical_experts): + self.layer_num_ = layer_index + logical_ids = list(range(num_logical_experts)) + self.fuse_moe_impl = SimpleNamespace(local_logics_expert_ids_list=logical_ids) + self.w13 = self._pack(rank, logical_ids, layer_index, 0) + self.w2 = self._pack(rank, logical_ids, layer_index, 100) + + @staticmethod + def _pack(rank, logical_ids, layer_index, offset): + # rank、layer、expert 和 tensor 类型都编码进数值;任何 recv 串包都会 + # 在目标 rank 的逐任务校验中表现为数值不一致。 + values = torch.tensor( + [[rank * 100_000 + layer_index * 1_000 + expert + offset] for expert in logical_ids], + dtype=torch.float32, + device="cuda", + ) + scales = values + 0.5 + return _Pack(values, scales) + + +def _free_port(): + sock = socket.socket() + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + sock.close() + return port + + +def _wait_for_transfer(transfer, control_group): + deadline = time.monotonic() + 30 + while time.monotonic() < deadline: + ready_count = torch.tensor([int(transfer.is_finished())], dtype=torch.int32) + dist.all_reduce(ready_count, op=dist.ReduceOp.MIN, group=control_group) + if int(ready_count.item()) == 1: + return + time.sleep(0.001) + raise TimeoutError("EPLB transfer worker did not finish globally") + + +def _wait_for_all_transfers(transfers, control_group): + """等待每个 rank 参与的全部并发传输完成。""" + deadline = time.monotonic() + 120 + while time.monotonic() < deadline: + local_finished = all(transfer.is_finished() for transfer in transfers) + globally_finished = torch.tensor([int(local_finished)], dtype=torch.int32) + dist.all_reduce(globally_finished, op=dist.ReduceOp.MIN, group=control_group) + if int(globally_finished.item()) == 1: + return + time.sleep(0.001) + raise TimeoutError("concurrent EPLB transfer workers did not finish globally") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_pytorch_keeps_unreferenced_pinned_source_alive_until_async_copy_finishes(): + """验证 pinned allocator 不会提前复用仍被异步 H2D 读取的内存。 + + copy stream 中先排入一个长任务,确保 H2D copy 在 Python 引用释放时仍未 + 完成。随后删除 pinned tensor 的唯一引用,并立刻申请同尺寸 pinned 内存: + + * 如果旧地址尚未复用,新 buffer 可以立即覆盖而不影响 H2D; + * 如果 allocator 返回了旧地址,对应 copy event 必须已经完成; + * 最终 GPU 数据必须保持为源 buffer 的原始内容。 + + 该测试验证当前 PyTorch/CUDA 组合的运行时行为;allocator 的正确性契约 + 仍由 PyTorch ``copy_`` 中的 host ``record_event`` 实现提供。 + """ + torch.cuda.synchronize() + num_elements = 8 * 1024 * 1024 + source = torch.full((num_elements,), 7, dtype=torch.int32, pin_memory=True) + source_ptr = source.data_ptr() + destination = torch.empty_like(source, device="cuda") + + copy_stream = torch.cuda.Stream() + copy_finished = torch.cuda.Event() + with torch.cuda.stream(copy_stream): + # copy 与 sleep 位于同一 stream,必须等 sleep 完成后才能开始。 + torch.cuda._sleep(1_000_000_000) + destination.copy_(source, non_blocking=True) + copy_finished.record() + + assert not copy_finished.query(), "test setup failed to leave the H2D copy pending" + + del source + gc.collect() + + replacement = torch.empty((num_elements,), dtype=torch.int32, pin_memory=True) + if replacement.data_ptr() == source_ptr: + # 相同地址只有在原 copy 已经结束、allocator 确认可以复用后才合法。 + assert copy_finished.query() + replacement.fill_(-3) + + copy_finished.synchronize() + assert torch.all(destination == 7) + + +def _worker(rank, port): + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + torch.cuda.set_device(rank) + dist.init_process_group("gloo", rank=rank, world_size=2) + control_group = dist.new_group([0, 1], backend="gloo") + transfer_group = dist.new_group([0, 1], backend="gloo") + + weights = [_FakeWeight(rank, layer_index) for layer_index in range(2)] + current = [[0, 1, 2], [2, 3, 0]] + target = [[0, 1, 3], [2, 3, 1]] + for expected_layer in range(2): + transfer_batches = build_transfer_plan( + current_placement=current, + target_placement=target, + layer_index=expected_layer, + num_logical_experts=4, + world_size=2, + ) + for transfer_batch in transfer_batches: + transfers = [ + PinnedMemoryEPLBTransfer( + weights=weights, + transfer_group=transfer_group, + current_global_rank=rank, + transfer_info=info, + ) + for info in transfer_batch + ] + for transfer in transfers: + assert all(buffer.pinned_row.is_pinned() for buffer in transfer.tensor_buffers) + transfer.start() + for transfer in transfers: + _wait_for_transfer(transfer, control_group) + if transfer.transfer_info.dest_rank == rank: + expected_expert = 3 if rank == 0 else 1 + expected_w13 = expected_layer * 100 + expected_expert + expected_w2 = expected_w13 + 10 + assert [buffer.name for buffer in transfer.tensor_buffers] == [ + "w13.weight", + "w13.weight_scale", + "w2.weight", + "w2.weight_scale", + ] + pinned_rows = [buffer.pinned_row for buffer in transfer.tensor_buffers] + assert torch.all(pinned_rows[0][0] == expected_w13) + assert torch.all(pinned_rows[1][0] == expected_w13 + 0.5) + assert torch.all(pinned_rows[2][0] == expected_w2) + assert torch.all(pinned_rows[3][0] == expected_w2 + 0.5) + + # 主槽位互换会形成覆盖环。两个方向必须同时完成 GPU -> pinned memory + # 传输后才能 commit,验证同一 rank 上并发的 send/recv 任务可以正常结束。 + swap_target = [[0, 3, 2], [2, 1, 0]] + swap_plan = build_transfer_plan( + current_placement=current, + target_placement=swap_target, + layer_index=0, + num_logical_experts=4, + world_size=2, + ) + assert len(swap_plan) == 1 + swap_infos = swap_plan[0] + assert len(swap_infos) == 2 + swap_transfers = [ + PinnedMemoryEPLBTransfer( + weights=weights, + transfer_group=transfer_group, + current_global_rank=rank, + transfer_info=info, + ) + for info in swap_infos + ] + for transfer in swap_transfers: + transfer.start() + for transfer in swap_transfers: + _wait_for_transfer(transfer, control_group) + + for transfer in swap_transfers: + if transfer.transfer_info.dest_rank == rank: + expected_expert = transfer.transfer_info.source_logical_expert_id + assert torch.all(transfer.tensor_buffers[0].pinned_row[0] == expected_expert) + dist.destroy_process_group() + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.device_count() < 2, + reason="requires two CUDA GPUs", +) +def test_eplb_pinned_memory_transfer_two_gpu_correctness(): + mp.spawn(_worker, args=(_free_port(),), nprocs=2, join=True) + + +def _many_concurrent_p2p_worker(rank, port): + """同时运行大量、重复 rank 对的 PinnedMemoryEPLBTransfer。""" + world_size = 4 + num_layers = 8 + num_logical_experts = 32 + transfers_per_rank_pair_per_layer = 8 + + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + torch.cuda.set_device(rank) + dist.init_process_group("gloo", rank=rank, world_size=world_size) + control_group = dist.new_group(list(range(world_size)), backend="gloo") + transfer_group = dist.new_group(list(range(world_size)), backend="gloo") + + try: + weights = [_StressFakeWeight(rank, layer_index, num_logical_experts) for layer_index in range(num_layers)] + + # 每层覆盖全部 12 个有向 rank 对,每个 rank 对重复 8 次。任务 identity + # 中的 layer、expert 和目标槽位不同,因此应该获得独立的 Gloo tag; + # source/destination rank 对则会被大量重复使用。 + transfer_infos = [] + for layer_index in range(num_layers): + for source_rank in range(world_size): + for dest_rank in range(world_size): + if source_rank == dest_rank: + continue + source_peers = [peer_rank for peer_rank in range(world_size) if peer_rank != dest_rank] + source_peer_index = source_peers.index(source_rank) + for repeat_index in range(transfers_per_rank_pair_per_layer): + expert_id = source_rank * transfers_per_rank_pair_per_layer + repeat_index + dest_local_expert_index = source_peer_index * transfers_per_rank_pair_per_layer + repeat_index + transfer_infos.append( + EPLBTransferInfo( + source_rank=source_rank, + layer_index=layer_index, + source_logical_expert_id=expert_id, + dest_rank=dest_rank, + dest_local_expert_index=dest_local_expert_index, + ) + ) + + local_transfer_infos = [ + transfer_info + for transfer_info in transfer_infos + if rank in (transfer_info.source_rank, transfer_info.dest_rank) + ] + transfers = [ + PinnedMemoryEPLBTransfer( + weights=weights, + transfer_group=transfer_group, + current_global_rank=rank, + transfer_info=transfer_info, + ) + for transfer_info in local_transfer_infos + ] + assert len(transfer_infos) == 768 + assert len(transfers) == 384 + + # 同一个 source/destination 对上的并发消息必须具有不同 tag,否则不同 + # 专家或张量可能被错误匹配。不同 rank 对可以安全复用相同整数 tag。 + message_keys = [] + for transfer in transfers: + transfer_info = transfer.transfer_info + for tensor_buffer in transfer.tensor_buffers: + message_keys.append( + ( + transfer_info.source_rank, + transfer_info.dest_rank, + transfer._build_p2p_message_tag(tensor_buffer.name), + ) + ) + assert len(message_keys) == len(set(message_keys)) + + dist.barrier(group=control_group) + for transfer in transfers: + transfer.start() + _wait_for_all_transfers(transfers, control_group) + + destination_transfers = [transfer for transfer in transfers if transfer.transfer_info.dest_rank == rank] + assert len(destination_transfers) == 192 + for transfer in destination_transfers: + transfer_info = transfer.transfer_info + expected_w13 = ( + transfer_info.source_rank * 100_000 + + transfer_info.layer_index * 1_000 + + transfer_info.source_logical_expert_id + ) + expected_values = [expected_w13, expected_w13 + 0.5, expected_w13 + 100, expected_w13 + 100.5] + assert [buffer.name for buffer in transfer.tensor_buffers] == [ + "w13.weight", + "w13.weight_scale", + "w2.weight", + "w2.weight_scale", + ] + for tensor_buffer, expected_value in zip(transfer.tensor_buffers, expected_values): + assert torch.all(tensor_buffer.pinned_row == expected_value) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.device_count() < 4, + reason="requires four CUDA GPUs", +) +def test_eplb_pinned_memory_transfer_four_gpu_many_concurrent_p2p(): + mp.spawn(_many_concurrent_p2p_worker, args=(_free_port(),), nprocs=4, join=True) diff --git a/unit_tests/server/test_api_start_eplb.py b/unit_tests/server/test_api_start_eplb.py new file mode 100644 index 0000000000..74984b5b62 --- /dev/null +++ b/unit_tests/server/test_api_start_eplb.py @@ -0,0 +1,145 @@ +import pytest + +from lightllm.server import api_start +from lightllm.server.core.objs.start_args_type import StartArgs + + +def test_eplb_redundant_expert_count_must_not_be_negative(monkeypatch): + args = StartArgs( + eplb_num_redundant_experts_per_rank=-1, + disable_vision=True, + disable_audio=True, + disable_shm_warning=True, + ) + + monkeypatch.setattr(api_start, "_set_envs_and_config", lambda args: None) + monkeypatch.setattr(api_start, "auto_set_max_req_total_len", lambda args: None) + monkeypatch.setattr(api_start, "auto_set_fused_shared_experts", lambda args: None) + monkeypatch.setattr(api_start, "set_unique_server_name", lambda args: None) + + with pytest.raises( + AssertionError, + match="--eplb_num_redundant_experts_per_rank must be greater than or equal to 0", + ): + api_start._launch_subprocesses(args) + + +def test_eplb_rebalance_count_must_not_be_less_than_negative_one(monkeypatch): + args = StartArgs( + eplb_rebalance_count=-2, + disable_vision=True, + disable_audio=True, + disable_shm_warning=True, + ) + + monkeypatch.setattr(api_start, "_set_envs_and_config", lambda args: None) + monkeypatch.setattr(api_start, "auto_set_max_req_total_len", lambda args: None) + monkeypatch.setattr(api_start, "auto_set_fused_shared_experts", lambda args: None) + monkeypatch.setattr(api_start, "set_unique_server_name", lambda args: None) + + with pytest.raises( + AssertionError, + match="--eplb_rebalance_count must be greater than or equal to -1", + ): + api_start._launch_subprocesses(args) + + +def test_eplb_redundant_experts_require_ep_moe(monkeypatch): + args = StartArgs( + enable_ep_moe=False, + eplb_num_redundant_experts_per_rank=1, + disable_vision=True, + disable_audio=True, + disable_shm_warning=True, + ) + + monkeypatch.setattr(api_start, "_set_envs_and_config", lambda args: None) + monkeypatch.setattr(api_start, "auto_set_max_req_total_len", lambda args: None) + monkeypatch.setattr(api_start, "auto_set_fused_shared_experts", lambda args: None) + monkeypatch.setattr(api_start, "set_unique_server_name", lambda args: None) + + with pytest.raises(AssertionError, match="EPLB requires --enable_ep_moe"): + api_start._launch_subprocesses(args) + + +def test_eplb_prefill_cudagraph_is_rejected_before_starting_subprocesses(monkeypatch): + args = StartArgs( + enable_ep_moe=True, + eplb_num_redundant_experts_per_rank=2, + enable_prefill_cudagraph=True, + disable_vision=True, + disable_audio=True, + disable_shm_warning=True, + ) + + monkeypatch.setattr(api_start, "_set_envs_and_config", lambda args: None) + monkeypatch.setattr(api_start, "auto_set_max_req_total_len", lambda args: None) + monkeypatch.setattr(api_start, "auto_set_fused_shared_experts", lambda args: None) + monkeypatch.setattr(api_start, "set_unique_server_name", lambda args: None) + monkeypatch.setattr( + api_start.process_manager, + "start_submodule_processes", + lambda *args, **kwargs: pytest.fail("subprocess startup must not be reached"), + ) + + with pytest.raises(AssertionError, match="EPLB does not support --enable_prefill_cudagraph"): + api_start._launch_subprocesses(args) + + +def test_eplb_redundant_experts_cannot_be_combined_with_rl(monkeypatch): + args = StartArgs( + enable_ep_moe=True, + enable_rl=True, + eplb_num_redundant_experts_per_rank=2, + disable_vision=True, + disable_audio=True, + disable_shm_warning=True, + ) + + monkeypatch.setattr(api_start, "_set_envs_and_config", lambda args: None) + monkeypatch.setattr(api_start, "auto_set_max_req_total_len", lambda args: None) + monkeypatch.setattr(api_start, "auto_set_fused_shared_experts", lambda args: None) + monkeypatch.setattr(api_start, "set_unique_server_name", lambda args: None) + monkeypatch.setattr( + api_start.process_manager, + "start_submodule_processes", + lambda *args, **kwargs: pytest.fail("subprocess startup must not be reached"), + ) + + with pytest.raises(AssertionError, match="EPLB redundant experts do not support --enable_rl"): + api_start._launch_subprocesses(args) + + +def test_eplb_mtp_combination_is_not_rejected_before_starting_subprocesses(monkeypatch): + args = StartArgs( + model_dir="test-model", + enable_ep_moe=True, + eplb_num_redundant_experts_per_rank=2, + mtp_mode="vanilla_no_att", + mtp_step=1, + eos_id=0, + data_type="float16", + disable_vision=True, + disable_audio=True, + disable_shm_warning=True, + ) + + monkeypatch.setattr(api_start, "_set_envs_and_config", lambda args: None) + monkeypatch.setattr(api_start, "auto_set_max_req_total_len", lambda args: None) + monkeypatch.setattr(api_start, "auto_set_fused_shared_experts", lambda args: None) + monkeypatch.setattr(api_start, "set_unique_server_name", lambda args: None) + monkeypatch.setattr(api_start, "auto_set_response_parsers", lambda args: None) + monkeypatch.setattr(api_start, "auto_configure_allreduce_flags_from_args", lambda args: None) + monkeypatch.setattr(api_start, "validate_ports", lambda ports: None) + monkeypatch.setattr(api_start, "set_env_start_args", lambda args: None) + monkeypatch.setattr(api_start, "get_shm_port_args", lambda create=False: None) + monkeypatch.setattr(api_start, "send_and_receive_node_ip", lambda args: None) + monkeypatch.setattr( + api_start.process_manager, + "start_submodule_processes", + lambda *args, **kwargs: (object(), None), + ) + monkeypatch.setattr(api_start.process_manager, "setup_exit_controller", lambda: None) + monkeypatch.setattr(api_start.process_manager, "register_process_tree", lambda process: None) + + api_start._launch_subprocesses(args)