MoE 原理介绍

import torch
import torch.nn as nn
import torch.nn.functional as F

class Expert(nn.Module):
    """一个专家 = 一个参数独立的小 FFN。结构相同、权重互不共享"""
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff, bias=False)
        self.w2 = nn.Linear(d_ff, d_model, bias=False)
        self.act = nn.SiLU()

    def forward(self, x):                     # x: [n, d_model]
        return self.w2(self.act(self.w1(x)))


class Router(nn.Module):
    """门控网络：打分 → Top-K 选路 → 顺带计算辅助负载均衡损失"""
    def __init__(self, d_model, num_experts, top_k=2, aux_alpha=0.01):
        super().__init__()
        self.top_k = top_k
        self.aux_alpha = aux_alpha            # α 一般取 0.01
        self.gate = nn.Linear(d_model, num_experts, bias=False)

    def forward(self, x):                     # x: [T, d_model]
        logits = self.gate(x)                               # [T, N]
        gate_probs = F.softmax(logits, dim=-1)              # 完整概率，用于 aux loss
        topk_val, topk_idx = torch.topk(logits, self.top_k, dim=-1)
        topk_val = F.softmax(topk_val, dim=-1)              # 只在被选中的 K 个内归一化

        # Switch Transformer 式辅助损失: L_aux = N · Σ f_i·P_i
        # f_i: 专家 i 实际被选中的比例（不可导，衡量真实负载）
        # P_i: 平均门控概率（可导，提供梯度）
        N, T = self.gate.out_features, x.size(0)
        f = torch.zeros(N, device=x.device)
        f.scatter_add_(0, topk_idx.reshape(-1),
                       torch.ones(topk_idx.numel(), device=x.device))
        f = f / (T * self.top_k)
        P = gate_probs.mean(dim=0)
        aux = self.aux_alpha * N * (f * P).sum()

        return topk_val, topk_idx, aux

class MoELayer(nn.Module):
    def __init__(self, d_model, d_ff, num_experts=8, top_k=2):
        super().__init__()
        self.top_k = top_k
        self.router = Router(d_model, num_experts, top_k)
        self.experts = nn.ModuleList(
            [Expert(d_model, d_ff) for _ in range(num_experts)])

    def forward(self, x):                     # x: [b, s, d]
        b, s, d = x.shape
        T = b * s
        x_flat = x.reshape(T, d)

        topk_val, topk_idx, aux = self.router(x_flat)
        flat_eid = topk_idx.reshape(-1)                    # [T*K] 目标专家编号
        flat_tok = x_flat.repeat_interleave(self.top_k, dim=0)
        flat_w = topk_val.reshape(-1, 1)                   # [T*K, 1] 门控权重

        # 按专家编号 permute，把同一专家的 token 拼在一起批量算
        # （grouped GEMM 的朴素版，真实实现用融合 CUDA kernel）
        perm = flat_eid.argsort()
        sorted_tok, sorted_eid = flat_tok[perm], flat_eid[perm]
        sorted_out = torch.empty_like(sorted_tok)
        for e, expert in enumerate(self.experts):
            mask = sorted_eid == e
            if mask.any():
                sorted_out[mask] = expert(sorted_tok[mask])

        # unpermute 并按门控权重加权合并
        weighted = torch.empty_like(sorted_out)
        weighted[perm] = sorted_out * flat_w
        out = torch.zeros(T, d, dtype=x.dtype, device=x.device)
        out.index_add_(0, torch.arange(T * self.top_k, device=x.device) // self.top_k,
                       weighted)
        return out.reshape(b, s, d), aux


# ===== 训练时使用 =====
moe_out, aux_loss = moe_layer(x)
loss = ce_loss(model_out, labels) + aux_loss   # 主损失 + α·L_aux
loss.backward()

import torch.distributed as dist

class ExpertParallelMoE(nn.Module):
    """专家并行 MoE：N 个专家均摊到 world_size 张卡，每卡只放 N/ws 个"""
    def __init__(self, d_model, d_ff, num_experts, top_k=2):
        super().__init__()
        ws = dist.get_world_size()
        assert num_experts % ws == 0
        self.top_k = top_k
        self.experts_per_rank = num_experts // ws
        self.router = Router(d_model, num_experts, top_k)   # 各卡持有同一份 Router
        self.local_experts = nn.ModuleList(                 # 本卡只持有自己那一片专家
            [Expert(d_model, d_ff) for _ in range(self.experts_per_rank)])

    def _all_to_all(self, tensor, send_counts, C):
        """按每卡 count 切分、pad 到等长 C 后 all_to_all（简化版；真实系统用变长 all_to_all）"""
        ws = dist.get_world_size()
        buf = torch.zeros(ws * C, *tensor.shape[1:],
                          dtype=tensor.dtype, device=tensor.device)
        off = 0
        for j in range(ws):
            c = send_counts[j].item()
            buf[j * C : j * C + c] = tensor[off : off + c]  # 第 j 块 = 发给卡 j 的数据
            off += c
        out = torch.empty_like(buf)
        dist.all_to_all_single(out, buf)   # out 第 j 块 = 卡 j 发给我的数据
        return out

    def forward(self, x):                     # x: [b, s, d]
        rank, ws = dist.get_rank(), dist.get_world_size()
        b, s, d = x.shape
        T = b * s
        x_flat = x.reshape(T, d)

        # ① Route：本地打分，得到全局专家编号
        topk_val, topk_idx, aux = self.router(x_flat)
        flat_eid = topk_idx.reshape(-1)
        flat_tok = x_flat.repeat_interleave(self.top_k, dim=0)
        flat_w = topk_val.reshape(-1, 1)
        perm = flat_eid.argsort()                          # 按专家编号排序
        send_eid, send_tok = flat_eid[perm], flat_tok[perm]

        # ② Dispatch：先交换每卡 token 数，再交换 token 与专家编号
        send_counts = torch.bincount(send_eid // self.experts_per_rank, minlength=ws)
        recv_counts = torch.empty_like(send_counts)
        dist.all_to_all_single(recv_counts, send_counts)
        C = int(torch.max(send_counts.max(), recv_counts.max()))   # pad 到等长
        recv_tok = self._all_to_all(send_tok, send_counts, C)
        recv_eid = self._all_to_all(send_eid, send_counts, C)

        # ③ Compute：只算本卡持有的专家（grouped GEMM）
        valid_tok = torch.cat([recv_tok[j*C : j*C + recv_counts[j]] for j in range(ws)])
        valid_eid = torch.cat([recv_eid[j*C : j*C + recv_counts[j]] for j in range(ws)])
        local_out = torch.zeros_like(valid_tok)
        base = rank * self.experts_per_rank
        for li, expert in enumerate(self.local_experts):
            mask = valid_eid == (base + li)
            if mask.any():
                local_out[mask] = expert(valid_tok[mask])

        # ④ Combine：结果原路送回（分块与 Dispatch 完全对称），unpermute 后加权合并
        out_buf = torch.zeros(ws * C, d, dtype=x.dtype, device=x.device)
        off = 0
        for j in range(ws):                                # 第 j 块送回卡 j
            c = recv_counts[j].item()
            out_buf[j*C : j*C + c] = local_out[off : off + c]
            off += c
        back = torch.empty_like(out_buf)
        dist.all_to_all_single(back, out_buf)

        outs = torch.empty_like(send_tok)                  # back 第 j 块 = 我发往卡 j 的输出
        off = 0
        for j in range(ws):
            c = send_counts[j].item()
            outs[off : off + c] = back[j*C : j*C + c]
            off += c

        weighted = torch.empty_like(outs)                  # 与单机版相同的 unpermute + 加权
        weighted[perm] = outs * flat_w
        out = torch.zeros(T, d, dtype=x.dtype, device=x.device)
        out.index_add_(0, torch.arange(T * self.top_k, device=x.device) // self.top_k,
                       weighted)
        return out.reshape(b, s, d), aux