mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-29 07:17:08 +08:00
141 lines
7.7 KiB
Python
141 lines
7.7 KiB
Python
"""
|
||
MoE 负载均衡策略：根据 token 统计重算 physical expert 放置。
|
||
|
||
术语：
|
||
- logical expert：模型定义的 expert，简写 log
|
||
- physical expert：GPU 上实际放置的 expert 副本，简写 phy，数量 = num_logical + num_redundant
|
||
"""
|
||
|
||
import torch
|
||
|
||
|
||
def _batch_greedy_indices(
|
||
tokens_per_expert: torch.Tensor,
|
||
physical_to_logical_map: torch.Tensor,
|
||
logical_count: torch.Tensor,
|
||
capacity: int,
|
||
) -> torch.Tensor:
|
||
"""
|
||
确定 physical expert 如何复制后，根据历史 batch 的信息贪心排布到各 rank 上
|
||
|
||
Parameters:
|
||
tokens_per_expert: [num_steps, num_moe_layers, num_logical_experts]
|
||
physical_to_logical_map: [num_moe_layers, num_physical_experts]
|
||
logical_count : [num_moe_layers, num_logical_experts]
|
||
capacity: 每个 rank 放置的 physical expert 数，即 num_local_physical_experts
|
||
|
||
return:
|
||
indices: [num_moe_layers, num_physical_experts]，每层为原 physical expert 编号的一个排列（0-based）
|
||
每连续 capacity 个元素表示一个 rank 放置的 physical experts
|
||
"""
|
||
weights = tokens_per_expert.sum(dim=0, dtype=torch.float32) / logical_count
|
||
# 按每副本历史总负载降序排列，同负载时保留原物理槽位顺序
|
||
physical_experts = physical_to_logical_map.to(torch.int64)
|
||
physical_order = weights.gather(-1, physical_experts).argsort(dim=-1, descending=True, stable=True)
|
||
expert_order = physical_experts.gather(-1, physical_order)
|
||
vectors = tokens_per_expert.permute(1, 2, 0).to(dtype=torch.float32, memory_format=torch.contiguous_format, copy=True)
|
||
vectors.div_(logical_count.unsqueeze(-1))
|
||
|
||
num_layers, _, num_steps = vectors.shape
|
||
num_physical = physical_order.size(1)
|
||
assert num_physical % capacity == 0, f"{num_physical=} {capacity=}"
|
||
num_ranks = num_physical // capacity
|
||
# 优化目标是逐 layer 独立的，对固定的 layer 而言，目标是将这层的 physical expert 分到各个 rank 上
|
||
# 并且使得 sum_{step i} max_{rank j} (sum_{physical expert e} [e 被分到 rank j] e 在 step i 的负载) 尽量小
|
||
# 具体贪心算法: 每次尝试将要考虑的 expert 放到各个 rank 上，选取当前最小化上述目标的放置方案
|
||
loads = vectors.new_zeros((num_layers, num_ranks, num_steps))
|
||
peaks = vectors.new_zeros((num_layers, num_steps))
|
||
occupancy = torch.zeros((num_layers, num_ranks), dtype=torch.int64, device=vectors.device)
|
||
candidates = torch.empty_like(loads)
|
||
scores = vectors.new_empty((num_layers, num_ranks))
|
||
indices = torch.empty_like(physical_order)
|
||
layers = torch.arange(num_layers, device=vectors.device)
|
||
for job_index in range(num_physical):
|
||
experts = expert_order[:, job_index]
|
||
vector = vectors[layers, experts]
|
||
# candidates[layer][j][step] 描述将当前 expert 放到 rank j 上后，各历史 step 在 rank j 上的负载
|
||
# 因此需和 peaks 取 max 得到当前 expert 放到 rank j 上后，各历史 step 的 rank 最大负载
|
||
torch.add(loads, vector.unsqueeze(1), out=candidates)
|
||
torch.maximum(candidates, peaks.unsqueeze(1), out=candidates)
|
||
torch.sum(candidates, dim=-1, out=scores)
|
||
# 排除已满的 rank 后选择最低分，argmin 在同分时返回最小 rank ID
|
||
scores.masked_fill_(occupancy == capacity, float("inf"))
|
||
rank = scores.argmin(dim=-1)
|
||
occupied = occupancy[layers, rank]
|
||
indices[layers, rank * capacity + occupied] = physical_order[:, job_index]
|
||
occupancy[layers, rank] = occupied + 1
|
||
updated = loads[layers, rank] + vector
|
||
loads[layers, rank] = updated
|
||
torch.maximum(peaks, updated, out=peaks)
|
||
return indices
|
||
|
||
|
||
def rebalance_experts_global(
|
||
tokens_per_expert: torch.Tensor,
|
||
num_physical_experts: int,
|
||
num_local_physical_experts: int,
|
||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||
"""
|
||
全局负载均衡：给热点 logical expert 分配冗余副本，再把所有副本均衡到各 GPU。
|
||
|
||
Parameters:
|
||
tokens_per_expert: [num_steps, num_moe_layers, num_logical_experts]
|
||
|
||
return:
|
||
physical_to_logical_map: [num_moe_layers, num_physical_experts]
|
||
logical_to_physical_map: [num_moe_layers, num_logical_experts, num_redundant_experts + 1]，无效位 -1
|
||
logical_count: [num_moe_layers, num_logical_experts]，各 logical expert 的副本数
|
||
"""
|
||
_, num_moe_layers, num_logical_experts = tokens_per_expert.shape
|
||
num_redundant_experts = num_physical_experts - num_logical_experts
|
||
device = tokens_per_expert.device
|
||
|
||
physical_to_logical_map = torch.empty(num_moe_layers, num_physical_experts, dtype=torch.int, device=device)
|
||
logical_to_physical_map = torch.full(
|
||
(num_moe_layers, num_logical_experts, num_redundant_experts + 1), -1, dtype=torch.int, device=device
|
||
)
|
||
logical_count = torch.ones(num_moe_layers, num_logical_experts, dtype=torch.int, device=device)
|
||
|
||
# 初始布局：前 num_logical_experts 个槽位放 logical experts，其余槽位放冗余副本
|
||
arange_layers = torch.arange(num_moe_layers, dtype=torch.int, device=device)
|
||
arange_logical = torch.arange(num_logical_experts, dtype=torch.int, device=device)
|
||
physical_to_logical_map[:, :num_logical_experts] = arange_logical
|
||
logical_to_physical_map[:, :, 0] = arange_logical
|
||
|
||
# 贪心分配冗余副本：每轮选一个 expert 复制。
|
||
# score[..., j] 是 expert j 当前的每副本负载，score_added 是给 j 加一个副本后的每副本负载；
|
||
# 把 argmax 位置替换成其 score_added 后，values[..., j] 即“副本给 j”时的最大每副本负载
|
||
# （j 非 argmax 时最大值不变，j 为 argmax 时是替换后重算的 max），按 step 求和取 argmin 得最优复制对象。
|
||
tokens_per_expert_all_diff = tokens_per_expert + arange_logical * 1e-4
|
||
for i in range(num_redundant_experts):
|
||
score = tokens_per_expert_all_diff / logical_count # NOTE: Values in score must be different from each other
|
||
score_added = tokens_per_expert / (logical_count + 1)
|
||
values, indices = score.max(-1, keepdim=True)
|
||
values = values.expand_as(score).contiguous()
|
||
score.scatter_(-1, indices, score_added.gather(-1, indices))
|
||
values.scatter_(-1, indices, score.max(-1, keepdim=True).values)
|
||
redundancy_indices = values.sum(0).argmin(-1) # [num_moe_layers]，本轮被复制的 logical 编号
|
||
|
||
physical_index = num_logical_experts + i # 本轮冗余副本占用的 physical 槽位
|
||
physical_to_logical_map[:, physical_index] = redundancy_indices
|
||
redundancy_count = logical_count.gather(-1, redundancy_indices.unsqueeze(-1)).squeeze(-1)
|
||
logical_to_physical_map[arange_layers, redundancy_indices, redundancy_count] = physical_index
|
||
logical_count[arange_layers, redundancy_indices] += 1
|
||
|
||
if num_local_physical_experts > 1:
|
||
assert num_physical_experts % num_local_physical_experts == 0
|
||
indices = _batch_greedy_indices(tokens_per_expert, physical_to_logical_map, logical_count, num_local_physical_experts)
|
||
physical_to_logical_map = physical_to_logical_map.gather(-1, indices)
|
||
# indices.argsort 给出旧槽位 -> 新槽位的映射，用它重映射 logical_to_physical_map 里的槽位编号（-1 先临时置 0 再还原）
|
||
mask = logical_to_physical_map == -1
|
||
logical_to_physical_map[mask] = 0
|
||
logical_to_physical_map = (
|
||
indices.argsort(-1)
|
||
.gather(-1, logical_to_physical_map.view(num_moe_layers, -1).to(torch.int64))
|
||
.view_as(logical_to_physical_map)
|
||
.to(torch.int)
|
||
)
|
||
logical_to_physical_map[mask] = -1
|
||
|
||
return physical_to_logical_map, logical_to_physical_map, logical_count
|