moe_ranking.py

1"""2MoE 精排模型（第12章 §12.4.3.7、§12.4.3.8）34用稀疏专家混合替换 Transformer 中的前馈网络（FFN）：56    MoE(x) = Σ_i g_i(x) · E_i(x)78  1. 稠密 MoE（所有专家都参与）：9       g_i(x) = softmax(W_r x)_i10     计算量随专家数线性增长，论文指出这不是稀疏化的初衷。1112  2. Top-K 稀疏路由（§12.4.3.7）：13       r   = W_r x14       p   = softmax(r)15       𝒯_K = TopK(p, K)                    # 只激活 K 个专家16       MoE(x) = Σ_{i∈𝒯_K} g̃_i · E_i(x),  g̃_i = p_i / Σ_{j∈𝒯_K} p_j17     计算量只与 K 有关，与专家总数解耦。1819  3. 共享专家（§12.4.3.8，DeepSeekMoE）：20       MoE(x) = Σ_{i=1}^{K_s} E_i^shared(x) + Σ_{j∈ℛ(x)} g_j · E_j^routed(x)21     共享专家对所有 token 都激活，负责捕获通用知识；22     路由专家负责捕获差异化知识，减少专家间的参数冗余。2324  4. 负载均衡损失（§12.4.3.8）：25     专家级：L_ExpBal = α₁ · N · Σ_i f_i · P_i26             f_i = 分配到专家 i 的 token 占比，P_i = 专家 i 的平均路由概率27     设备级：L_DevBal = α₂ · D · Σ_d f'_d · P'_d28             （按设备对专家分组，避免单个设备成为瓶颈）2930  5. 工程约束（§12.4.3.7 正文）：31     MoE 引入的通信开销与负载不均，往往是「稀疏化省下的算力」32     被抵消的主要原因，因此负载均衡损失不是可有可无的正则项。3334本文件的对照实验直接检验上述两点：35  - 去掉负载均衡损失，观察专家利用率是否塌缩到少数专家；36  - 加上专家级 / 设备级均衡损失，观察利用率与设备分配是否变均匀。3738本书依据：《第十二章-精排模型.pdf》§12.4.3.7–§12.4.3.8。39依赖：PyTorch。运行：python3 moe_ranking.py40"""4142import random43import torch44import torch.nn as nn45import torch.nn.functional as F46import torch.optim as optim474849def _auc(scores, labels) -> float:50    pos = [s for s, y in zip(scores, labels) if y == 1]51    neg = [s for s, y in zip(scores, labels) if y == 0]52    if not pos or not neg:53        return 0.554    return sum(1.0 if p > n else 0.5 if p == n else 0.055               for p in pos for n in neg) / (len(pos) * len(neg))565758# ── 路由层 ───────────────────────────────────────────────────────────────────5960class TopKRouter(nn.Module):61    """62    Top-K 稀疏路由：r = W_r x → softmax → 取 Top-K → 在 K 个专家内重新归一化。6364    返回：65      topk_idx  : (B, K)  被选中的专家下标66      gates     : (B, K)  归一化后的门控权重67      probs     : (B, N)  完整路由分布（用于计算负载均衡的 P_i）68    """6970    def __init__(self, d_model: int, n_experts: int, top_k: int):71        super().__init__()72        self.top_k = top_k73        self.gate = nn.Linear(d_model, n_experts, bias=False)7475    def forward(self, x: torch.Tensor):76        probs = F.softmax(self.gate(x), dim=-1)            # (B, N)77        topk_p, topk_idx = torch.topk(probs, self.top_k, dim=-1)78        gates = topk_p / topk_p.sum(dim=-1, keepdim=True)  # 重新归一化79        return topk_idx, gates, probs808182# ── MoE 前馈层 ───────────────────────────────────────────────────────────────8384class MoEFFN(nn.Module):85    """86    Top-K 路由专家 + 共享专家。8788    参数：89      d_model     : 输入/输出维度90      d_ff        : 每个专家的隐藏维度91      n_routed    : 路由专家数（细粒度切分时取较大值，单个专家较窄）92      top_k       : 每个 token 激活的路由专家数93      n_shared    : 共享专家数（DeepSeekMoE 做法，对所有 token 激活）94      n_device    : 设备数，用于设备级负载均衡损失95    """9697    def __init__(self, d_model: int, d_ff: int, n_routed: int = 8,98                 top_k: int = 2, n_shared: int = 1, n_device: int = 2):99        super().__init__()100        self.n_routed = n_routed101        self.n_shared = n_shared102        self.n_device = n_device103104        self.router = TopKRouter(d_model, n_routed, top_k)105        self.routed_experts = nn.ModuleList([106            nn.Sequential(nn.Linear(d_model, d_ff), nn.GELU(),107                          nn.Linear(d_ff, d_model))108            for _ in range(n_routed)])109        self.shared_experts = nn.ModuleList([110            nn.Sequential(nn.Linear(d_model, d_ff), nn.GELU(),111                          nn.Linear(d_ff, d_model))112            for _ in range(n_shared)])113        # 残差投影：与 Transformer FFN 的位置对齐114        self.out_proj = nn.Linear(d_model, d_model)115116    def _device_of(self, expert_idx: torch.Tensor) -> torch.Tensor:117        """把专家下标平均分到 n_device 个设备上。"""118        return expert_idx * self.n_device // self.n_routed119120    def forward(self, x: torch.Tensor):121        """122        x : (B, d_model)123        返回: y (B, d_model), stats(dict) 供计算负载均衡损失与诊断124        """125        B, _ = x.shape126        topk_idx, gates, probs = self.router(x)            # (B,K), (B,K), (B,N)127128        # 路由专家：逐专家处理被分配给它的 token（实现上是掩码 + 累加，129        # 与真实系统中按 token 重排（permutation）后分组计算等价）130        y = torch.zeros_like(x)131        for k in range(topk_idx.size(1)):132            idx_k = topk_idx[:, k]                         # (B,)133            gate_k = gates[:, k].unsqueeze(-1)             # (B,1)134            for e in range(self.n_routed):135                mask = (idx_k == e).unsqueeze(-1)          # (B,1)136                if not mask.any():137                    continue138                out_e = self.routed_experts[e](x)          # (B, d_model)139                y = y + torch.where(mask, gate_k * out_e,140                                    torch.zeros_like(out_e))141142        # 共享专家：所有 token 都激活143        for e in self.shared_experts:144            y = y + e(x)145146        y = self.out_proj(y)147148        # ── 负载均衡统计 ─────────────────────────────────────────────────────149        # f_i：分配到专家 i 的 token 占比150        one_hot = F.one_hot(topk_idx, self.n_routed).float().sum(dim=1)  # (B,N)151        f = one_hot.mean(dim=0)                            # (N,)152        # P_i：专家 i 的平均路由概率153        P = probs.mean(dim=0)                              # (N,)154        # 设备级：把专家按设备分组聚合155        dev = self._device_of(torch.arange(self.n_routed))156        f_dev = torch.zeros(self.n_device, device=x.device)157        P_dev = torch.zeros(self.n_device, device=x.device)158        f_dev.index_add_(0, dev, f)159        P_dev.index_add_(0, dev, P)160161        stats = dict(f=f, P=P, f_dev=f_dev, P_dev=P_dev,162                     top1_share=one_hot[:, topk_idx[:, 0]].mean().item())163        return y, stats164165    # ── 负载均衡损失 ─────────────────────────────────────────────────────────166    def balance_loss(self, stats, alpha_exp: float = 0.01,167                     alpha_dev: float = 0.01) -> torch.Tensor:168        N = self.n_routed169        L_exp = N * (stats["f"] * stats["P"]).sum()170        D = self.n_device171        L_dev = D * (stats["f_dev"] * stats["P_dev"]).sum()172        return alpha_exp * L_exp + alpha_dev * L_dev173174175# ── 单个 MoE 层上的精排模型 ──────────────────────────────────────────────────176177class MoERanker(nn.Module):178    """179    演示用的极简结构：嵌入 → 投影到 d_model → MoE 层 → 输出 logit。180181    只用一层 MoE，是为了让路由行为成为唯一的观察对象；182    真实系统中的 MoE 位于 Transformer 的每个 FFN 位置。183    """184185    def __init__(self, vocab: int, n_fields: int, d_model: int = 16,186                 d_ff: int = 16, n_routed: int = 8, top_k: int = 2,187                 n_shared: int = 1, n_device: int = 2):188        super().__init__()189        self.emb = nn.Embedding(vocab, d_model)190        self.n_fields = n_fields191        self.proj = nn.Linear(d_model * n_fields, d_model)192        self.moe = MoEFFN(d_model, d_ff, n_routed, top_k, n_shared, n_device)193        self.norm = nn.LayerNorm(d_model)194        self.head = nn.Linear(d_model, 1)195196    def forward(self, feat_ids: torch.Tensor):197        e = torch.cat([self.emb(feat_ids[:, i])198                       for i in range(self.n_fields)], dim=1)199        h = F.gelu(self.proj(e))                           # (B, d_model)200        y, stats = self.moe(h)201        h = self.norm(h + y)                               # 残差 + LayerNorm202        return self.head(h).squeeze(-1), stats203204    def fit(self, feat_ids, y, use_balance: bool, epochs: int = 40,205            lr: float = 5e-3):206        opt = optim.Adam(self.parameters(), lr=lr, weight_decay=1e-4)207        for ep in range(epochs):208            opt.zero_grad()209            logit, stats = self.forward(feat_ids)210            loss = F.binary_cross_entropy_with_logits(logit, y.float())211            if use_balance:212                loss = loss + self.moe.balance_loss(stats)213            loss.backward()214            opt.step()215            if (ep + 1) % 20 == 0:216                print(f"  epoch {ep+1:3d}  loss={loss.item():.4f}")217        return loss.item()218219    def predict_proba(self, feat_ids):220        with torch.no_grad():221            logit, _ = self.forward(feat_ids)222            return torch.sigmoid(logit)223224225# ── 合成数据 ─────────────────────────────────────────────────────────────────226227def make_data(n: int = 3000, vocab: int = 16, n_fields: int = 6,228              seed: int = 42):229    """230    特征交互型标签：让不同样本落在特征空间的不同区域，231    从而为「不同专家负责不同区域」提供存在的理由。232233    注意：真实系统中专家分工是被数据自然诱导出来的；234    本文件不人为构造「专家分工」，只观察路由是否出现塌缩。235    """236    torch.manual_seed(seed)237    random.seed(seed)238239    feat = torch.randint(0, vocab, (n, n_fields)).float()240    logit = (0.35 * feat[:, 0] + 0.30 * feat[:, 1]241             - 0.25 * feat[:, 2]242             + 0.20 * (feat[:, 3] * feat[:, 4]).sin()243             + 0.15 * (feat[:, 5] > 8).float()244             - 1.2)245    y = (torch.sigmoid(logit / 3.0 + 0.4 * torch.randn(n)) > 0.5).long()246    return dict(feat=feat.long(), y=y)247248249# ── 主演示 ────────────────────────────────────────────────────────────────────250251if __name__ == "__main__":252    random.seed(42)253    torch.manual_seed(42)254255    VOCAB, N_FIELDS = 16, 6256    data = make_data(3000, vocab=VOCAB, n_fields=N_FIELDS)257    split = 2400258    tr = {k: v[:split] for k, v in data.items()}259    te = {k: v[split:] for k, v in data.items()}260261    def build():262        torch.manual_seed(42)263        return MoERanker(VOCAB, N_FIELDS, d_model=16, d_ff=16,264                         n_routed=8, top_k=2, n_shared=1, n_device=2)265266    def evaluate(model):267        p = model.predict_proba(te["feat"])268        return _auc(p.tolist(), te["y"].tolist())269270    def show_util(model, tag):271        with torch.no_grad():272            _, stats = model.forward(te["feat"])273        f = stats["f"].tolist()274        fd = stats["f_dev"].tolist()275        print(f"  {tag}")276        print("    专家利用率 f_i: " +277              " ".join(f"E{i}={v:.3f}" for i, v in enumerate(f)))278        print(f"    最大/最小利用率比 = {max(f) / max(min(f), 1e-6):.2f}"279              f"   分布标准差 = {torch.tensor(f).std().item():.4f}")280        print("    设备 token 占比 f'_d: " +281              " ".join(f"Dev{d}={v:.3f}" for d, v in enumerate(fd)))282283    # ── 对照组：无负载均衡损失 ───────────────────────────────────────────────284    print("=== 对照组：只优化主任务（无负载均衡损失）===")285    m_free = build()286    m_free.fit(tr["feat"], tr["y"], use_balance=False, epochs=40)287    auc_free = evaluate(m_free)288    show_util(m_free, "路由利用情况：")289290    # ── 实验组：加入专家级 + 设备级负载均衡损失 ───────────────────────────────291    print("\n=== MoE：主任务 + 专家级 / 设备级负载均衡损失 ===")292    m_bal = build()293    m_bal.fit(tr["feat"], tr["y"], use_balance=True, epochs=40)294    auc_bal = evaluate(m_bal)295    show_util(m_bal, "路由利用情况：")296297    # ── 对比结论 ─────────────────────────────────────────────────────────────298    print(f"\n对比（正文 §12.4.3.8）：")299    print(f"  无均衡损失 AUC={auc_free:.4f}")300    print(f"  有均衡损失 AUC={auc_bal:.4f}")301    print("  正文要点：Top-K 路由天然倾向于把 token 集中到少数「强者」专家，")302    print("  这会让其余专家参数闲置，并在分布式训练中造成设备负载倾斜；")303    print("  负载均衡损失是让稀疏化真正兑现「算力与专家数解耦」的必要条件。")304305    # ── 共享专家的作用 ───────────────────────────────────────────────────────306    print("\n共享专家 vs 路由专家（正文 §12.4.3.8）：")307    print("  DeepSeekMoE 让共享专家对所有 token 激活，捕获通用知识；")308    print("  路由专家只对部分 token 激活，捕获差异化知识。")309    n_shared_p = sum(p.numel() for p in m_bal.moe.shared_experts.parameters())310    n_routed_p = sum(p.numel() for p in m_bal.moe.routed_experts.parameters())311    print(f"  共享专家参数={n_shared_p}  路由专家总参数={n_routed_p}"312          f"  每个 token 实际激活的路由专家={m_bal.moe.router.top_k} / "313          f"{m_bal.moe.n_routed}")314315    # ── 精排推理演示 ─────────────────────────────────────────────────────────316    p = m_bal.predict_proba(te["feat"][:20]).tolist()317    top5 = sorted(enumerate(p), key=lambda t: -t[1])[:5]318    print("\n精排 Top-5（候选集 20 条）:")319    for rank_i, (idx, s) in enumerate(top5, 1):320        print(f"  #{rank_i}  item={idx:2d}  p_click={s:.4f}  "321              f"label={te['y'][idx].item()}")