Megatron-LM MoE Layer 源码精读

def forward(self, hidden_states: torch.Tensor):
    def custom_forward(hidden_states):
        # ① Shared expert 先于路由执行（使用原始 hidden_states）
        shared_expert_output = self.shared_experts_compute(hidden_states)

        # ② Route：计算路由概率与选择掩码
        probs, routing_map = self.route(hidden_states)

        # ③ Preprocess：token dropping / padding 等前处理
        hidden_states, probs, residual = self.preprocess(
            hidden_states, probs, routing_map
        )

        # ④ Dispatch：All-to-All 分发 token 到对应 EP rank
        dispatched_input, probs = self.dispatch(hidden_states, probs)

        # ⑤ Compute：本 rank 上的 local experts 计算
        output, mlp_bias = self.routed_experts_compute(
            dispatched_input, probs, residual
        )

        # ⑥ Combine：All-to-All 收集 + 加权合并 + 叠加 shared expert
        output = self.combine(output, shared_expert_output)
        return output, mlp_bias

    if self.moe_layer_recompute:
        outputs = tensor_parallel.checkpoint(custom_forward, False, hidden_states)
    else:
        outputs = custom_forward(hidden_states)
    return outputs

class MoELayer(BaseMoELayer):
    def __init__(self, config, submodules, layer_number, pg_collection):
        super().__init__(config=config, ...)

        # ---- 1. Router ----
        self.router = TopKRouter(config=self.config, pg_collection=pg_collection)

        # ---- 2. Token Dispatcher ----
        if config.moe_token_dispatcher_type == "alltoall":
            self.token_dispatcher = MoEAlltoAllTokenDispatcher(...)
        elif config.moe_token_dispatcher_type == "allgather":
            self.token_dispatcher = MoEAllGatherTokenDispatcher(...)
        elif config.moe_token_dispatcher_type == "flex":
            self.token_dispatcher = MoEFlexTokenDispatcher(...)

        # ---- 3. Routed Experts ----
        self.experts = build_module(submodules.experts, self.num_local_experts, ...)

        # ---- 4. Shared Experts (可选) ----
        if self.use_shared_expert:
            self.shared_experts = build_module(submodules.shared_experts, ...)
            if self.shared_expert_overlap:
                self.token_dispatcher.set_shared_experts(self.shared_experts)

class BaseMoELayer(MegatronModule, ABC):
    def __init__(self, config, layer_number, pg_collection):
        super().__init__(config)
        self.ep_group = pg_collection.ep
        ep_size = utils.get_pg_size(self.ep_group)
        ep_rank = utils.get_pg_rank(self.ep_group)

        # 每个 EP rank 持有的 local expert 数量
        assert self.config.num_moe_experts % ep_size == 0
        self.num_local_experts = self.config.num_moe_experts // ep_size

        # 本 rank 负责的 expert 索引范围
        local_expert_indices_offset = ep_rank * self.num_local_experts
        self.local_expert_indices = [
            local_expert_indices_offset + i
            for i in range(self.num_local_experts)
        ]

def topk_routing_with_score_function(
    logits, topk, use_pre_softmax=False, num_groups=None,
    group_topk=None, scaling_factor=None, score_function="softmax",
    expert_bias=None, fused=False,
):
    if score_function == "sigmoid":
        scores = torch.sigmoid(logits.float()).type_as(logits)
        if expert_bias is not None:
            # expert_bias 在 top-k 选择时加到 scores 上
            scores_for_routing = scores + expert_bias
            _, top_indices = compute_topk(scores_for_routing, topk, ...)
            # 但 probs 仍然使用原始 scores（不含 bias）
            scores = torch.gather(scores, dim=1, index=top_indices)
        else:
            scores, top_indices = compute_topk(scores, topk, ...)

        if topk > 1:
            probs = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20)
        else:
            probs = scores

    routing_probs = torch.zeros_like(logits).scatter(1, top_indices, probs)
    routing_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool()
    return routing_probs, routing_map

def _apply_aux_loss(self, probs, scores_for_aux_loss, routing_map):
    tokens_per_expert = routing_map.sum(dim=0)  # [num_experts]

    # EP > 1: 跨 tp_ep 组 reduce，获取完整 EP 组的 token 分布
    # EP = 1: 跨 tp_cp 组 reduce
    if self.ep_group.size() > 1:
        reduce_group = self.tp_ep_group     # expert_tp x EP
    else:
        reduce_group = self.tp_cp_group     # TP x CP

    tokens_per_expert = reduce_from_tensor_model_parallel_region(
        tokens_per_expert, reduce_group
    )

    aux_loss = switch_load_balancing_loss_func(
        probs=scores_for_aux_loss,
        tokens_per_expert=tokens_per_expert,
        total_num_tokens=num_tokens * reduce_group.size(),
        topk=self.topk,
        num_experts=self.config.num_moe_experts,
        moe_aux_loss_coeff=aux_loss_coeff,
    )
    # 通过 MoEAuxLossAutoScaler 将 aux loss 附加到 probs 的梯度图上
    probs = MoEAuxLossAutoScaler.apply(probs, aux_loss)
    return probs

def token_dispatch(self, permutated_local_input_tokens, permuted_probs):
    """All-to-All 通信：在 EP group 内交换 token"""
    # input_splits: [ep_size] 本 rank 发给各 EP rank 的 token 数
    # output_splits: [ep_size] 本 rank 从各 EP rank 接收的 token 数
    global_input_tokens = all_to_all(
        self.ep_group,
        permutated_local_input_tokens,
        self.output_splits,    # 接收 split
        self.input_splits      # 发送 split
    )
    global_probs = all_to_all(
        self.ep_group,
        permuted_probs,
        self.output_splits,
        self.input_splits
    )
    return global_input_tokens, global_probs

def forward(self, permuted_local_hidden_states, tokens_per_expert, permuted_probs):
    # 将权重 reshape 为 [num_local_experts, H, ffn_hidden]
    w1 = self.weight1.view(self.num_local_experts, self.config.hidden_size, -1)
    w2 = self.weight2.view(self.num_local_experts, -1, self.config.hidden_size)

    # GroupedGEMM: fc1
    # input: [total_tokens, H], w1: [num_experts, H, ffn_hidden*2]
    fc1_output = gg.ops.gmm(
        permuted_local_hidden_states, w1, tokens_per_expert, trans_b=False
    )

    # SwiGLU activation + 路由权重
    intermediate = self.activation_func_with_probs(
        fc1_output, permuted_probs.unsqueeze(-1)
    )

    # GroupedGEMM: fc2
    # input: [total_tokens, ffn_hidden], w2: [num_experts, ffn_hidden, H]
    fc2_output = gg.ops.gmm(
        intermediate, w2, tokens_per_expert, trans_b=False
    )
    return fc2_output, None

class SharedExpertMLP(MLP):
    def forward(self, hidden_states):
        # 标准 MLP 前向：fc1 -> activation -> fc2
        if self.activation_recompute:
            output, _ = self._forward_with_activation_recompute(hidden_states)
        else:
            output, _ = super().forward(hidden_states)

        # 可选：sigmoid gate 调节输出
        if self.use_shared_expert_gate:
            logits = F.linear(hidden_states, self.gate_weight)
            gate_score = torch.sigmoid(logits)
            output = output * gate_score

        return output