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