deepinfer-eplb-policy
deepinfer/deepinfer/model_executor/eplb/policy.py · a679bb8 · L12–140 · 固定 commit 上游
这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。
12def _batch_greedy_indices(
13 tokens_per_expert: torch.Tensor,
14 physical_to_logical_map: torch.Tensor,
15 logical_count: torch.Tensor,
16 capacity: int,
17) -> torch.Tensor:
18 """
19 确定 physical expert 如何复制后，根据历史 batch 的信息贪心排布到各 rank 上
20
21 Parameters:
22 tokens_per_expert: [num_steps, num_moe_layers, num_logical_experts]
23 physical_to_logical_map: [num_moe_layers, num_physical_experts]
24 logical_count : [num_moe_layers, num_logical_experts]
25 capacity: 每个 rank 放置的 physical expert 数，即 num_local_physical_experts
26
27 return:
28 indices: [num_moe_layers, num_physical_experts]，每层为原 physical expert 编号的一个排列（0-based）
29 每连续 capacity 个元素表示一个 rank 放置的 physical experts
30 """
31 weights = tokens_per_expert.sum(dim=0, dtype=torch.float32) / logical_count
32 # 按每副本历史总负载降序排列，同负载时保留原物理槽位顺序
33 physical_experts = physical_to_logical_map.to(torch.int64)
34 physical_order = weights.gather(-1, physical_experts).argsort(dim=-1, descending=True, stable=True)
35 expert_order = physical_experts.gather(-1, physical_order)
36 vectors = tokens_per_expert.permute(1, 2, 0).to(dtype=torch.float32, memory_format=torch.contiguous_format, copy=True)
37 vectors.div_(logical_count.unsqueeze(-1))
38
39 num_layers, _, num_steps = vectors.shape
40 num_physical = physical_order.size(1)
41 assert num_physical % capacity == 0, f"{num_physical=} {capacity=}"
42 num_ranks = num_physical // capacity
43 # 优化目标是逐 layer 独立的，对固定的 layer 而言，目标是将这层的 physical expert 分到各个 rank 上
44 # 并且使得 sum_{step i} max_{rank j} (sum_{physical expert e} [e 被分到 rank j] e 在 step i 的负载) 尽量小
45 # 具体贪心算法: 每次尝试将要考虑的 expert 放到各个 rank 上，选取当前最小化上述目标的放置方案
46 loads = vectors.new_zeros((num_layers, num_ranks, num_steps))
47 peaks = vectors.new_zeros((num_layers, num_steps))
48 occupancy = torch.zeros((num_layers, num_ranks), dtype=torch.int64, device=vectors.device)
49 candidates = torch.empty_like(loads)
50 scores = vectors.new_empty((num_layers, num_ranks))
51 indices = torch.empty_like(physical_order)
52 layers = torch.arange(num_layers, device=vectors.device)
53 for job_index in range(num_physical):
54 experts = expert_order[:, job_index]
55 vector = vectors[layers, experts]
56 # candidates[layer][j][step] 描述将当前 expert 放到 rank j 上后，各历史 step 在 rank j 上的负载
57 # 因此需和 peaks 取 max 得到当前 expert 放到 rank j 上后，各历史 step 的 rank 最大负载
58 torch.add(loads, vector.unsqueeze(1), out=candidates)
59 torch.maximum(candidates, peaks.unsqueeze(1), out=candidates)
60 torch.sum(candidates, dim=-1, out=scores)
61 # 排除已满的 rank 后选择最低分，argmin 在同分时返回最小 rank ID
62 scores.masked_fill_(occupancy == capacity, float("inf"))
63 rank = scores.argmin(dim=-1)
64 occupied = occupancy[layers, rank]
65 indices[layers, rank * capacity + occupied] = physical_order[:, job_index]
66 occupancy[layers, rank] = occupied + 1
67 updated = loads[layers, rank] + vector
68 loads[layers, rank] = updated
69 torch.maximum(peaks, updated, out=peaks)
70 return indices
71
72
73def rebalance_experts_global(
74 tokens_per_expert: torch.Tensor,
75 num_physical_experts: int,
76 num_local_physical_experts: int,
77) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
78 """
79 全局负载均衡：给热点 logical expert 分配冗余副本，再把所有副本均衡到各 GPU。
80
81 Parameters:
82 tokens_per_expert: [num_steps, num_moe_layers, num_logical_experts]
83
84 return:
85 physical_to_logical_map: [num_moe_layers, num_physical_experts]
86 logical_to_physical_map: [num_moe_layers, num_logical_experts, num_redundant_experts + 1]，无效位 -1
87 logical_count: [num_moe_layers, num_logical_experts]，各 logical expert 的副本数
88 """
89 _, num_moe_layers, num_logical_experts = tokens_per_expert.shape
90 num_redundant_experts = num_physical_experts - num_logical_experts
91 device = tokens_per_expert.device
92
93 physical_to_logical_map = torch.empty(num_moe_layers, num_physical_experts, dtype=torch.int, device=device)
94 logical_to_physical_map = torch.full(
95 (num_moe_layers, num_logical_experts, num_redundant_experts + 1), -1, dtype=torch.int, device=device
96 )
97 logical_count = torch.ones(num_moe_layers, num_logical_experts, dtype=torch.int, device=device)
98
99 # 初始布局：前 num_logical_experts 个槽位放 logical experts，其余槽位放冗余副本
100 arange_layers = torch.arange(num_moe_layers, dtype=torch.int, device=device)
101 arange_logical = torch.arange(num_logical_experts, dtype=torch.int, device=device)
102 physical_to_logical_map[:, :num_logical_experts] = arange_logical
103 logical_to_physical_map[:, :, 0] = arange_logical
104
105 # 贪心分配冗余副本：每轮选一个 expert 复制。
106 # score[..., j] 是 expert j 当前的每副本负载，score_added 是给 j 加一个副本后的每副本负载；
107 # 把 argmax 位置替换成其 score_added 后，values[..., j] 即“副本给 j”时的最大每副本负载
108 # （j 非 argmax 时最大值不变，j 为 argmax 时是替换后重算的 max），按 step 求和取 argmin 得最优复制对象。
109 tokens_per_expert_all_diff = tokens_per_expert + arange_logical * 1e-4
110 for i in range(num_redundant_experts):
111 score = tokens_per_expert_all_diff / logical_count # NOTE: Values in score must be different from each other
112 score_added = tokens_per_expert / (logical_count + 1)
113 values, indices = score.max(-1, keepdim=True)
114 values = values.expand_as(score).contiguous()
115 score.scatter_(-1, indices, score_added.gather(-1, indices))
116 values.scatter_(-1, indices, score.max(-1, keepdim=True).values)
117 redundancy_indices = values.sum(0).argmin(-1) # [num_moe_layers]，本轮被复制的 logical 编号
118
119 physical_index = num_logical_experts + i # 本轮冗余副本占用的 physical 槽位
120 physical_to_logical_map[:, physical_index] = redundancy_indices
121 redundancy_count = logical_count.gather(-1, redundancy_indices.unsqueeze(-1)).squeeze(-1)
122 logical_to_physical_map[arange_layers, redundancy_indices, redundancy_count] = physical_index
123 logical_count[arange_layers, redundancy_indices] += 1
124
125 if num_local_physical_experts > 1:
126 assert num_physical_experts % num_local_physical_experts == 0
127 indices = _batch_greedy_indices(tokens_per_expert, physical_to_logical_map, logical_count, num_local_physical_experts)
128 physical_to_logical_map = physical_to_logical_map.gather(-1, indices)
129 # indices.argsort 给出旧槽位 -> 新槽位的映射，用它重映射 logical_to_physical_map 里的槽位编号（-1 先临时置 0 再还原）
130 mask = logical_to_physical_map == -1
131 logical_to_physical_map[mask] = 0
132 logical_to_physical_map = (
133 indices.argsort(-1)
134 .gather(-1, logical_to_physical_map.view(num_moe_layers, -1).to(torch.int64))
135 .view_as(logical_to_physical_map)
136 .to(torch.int)
137 )
138 logical_to_physical_map[mask] = -1
139
140 return physical_to_logical_map, logical_to_physical_map, logical_count