混合专家模型架构

y = Σ_{i ∈ TopK(g(x))} g(x)_i · E_i(x)

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

class TopKGating(nn.Module):
    def __init__(self, d_model, n_experts, top_k=2):
        super().__init__()
        self.top_k = top_k
        self.n_experts = n_experts
        self.gate = nn.Linear(d_model, n_experts, bias=False)

    def forward(self, x):
        # x: (B*T, D)
        logits = self.gate(x)                      # (N, E)
        scores = F.softmax(logits, dim=-1)
        topk_val, topk_idx = scores.topk(self.top_k, dim=-1)
        # 归一化，让 Top-K 权重和为 1
        topk_val = topk_val / topk_val.sum(dim=-1, keepdim=True)
        return topk_idx, topk_val

y = E_shared(x) + Σ_{i ∈ TopK} g_i · E_i(x)

def load_balancing_loss(gate_logits, topk_idx, n_experts):
    """Switch Transformer 风格的辅助损失"""
    # f_i: 分配给专家 i 的 token 比例
    one_hot = F.one_hot(topk_idx, n_experts).float()      # (N, K, E)
    f = one_hot.sum(dim=(0, 1)) / (topk_idx.shape[0] * topk_idx.shape[1])
    # P_i: 专家 i 的平均门控概率
    P = F.softmax(gate_logits, dim=-1).mean(dim=0)        # (E,)
    return n_experts * (f * P).sum()

capacity = ceil(capacity_factor × tokens_per_batch / n_experts)

def bias_update(bias, usage, gamma=0.001):
    """usage 过高则降低偏置，过低则提高，目标是均衡"""
    target = usage.mean()
    return bias + gamma * (target - usage)

def routing_stats(topk_idx, n_experts):
    counts = torch.bincount(topk_idx.flatten(), minlength=n_experts).float()
    usage = counts / counts.sum()
    entropy = -(usage * (usage + 1e-9).log()).sum()
    return {
        "usage": usage.tolist(),
        "max_share": usage.max().item(),
        "entropy": entropy.item(),
        "ideal_share": 1.0 / n_experts,
    }

import torch.distributed as dist

def expert_parallel_forward(x, topk_idx, topk_val, experts_per_rank):
    rank = dist.get_rank()
    world = dist.get_world_size()
    # 1. 按目标 rank 分桶
    target_rank = topk_idx // experts_per_rank
    send_buckets = [x[target_rank == r] for r in range(world)]
    # 2. All-to-All 交换 token
    recv_buckets = [torch.empty_like(send_buckets[0]) for _ in range(world)]
    dist.all_to_all(recv_buckets, send_buckets)
    # 3. 本地专家计算
    local_out = run_local_experts(recv_buckets)
    # 4. 反向 All-to-All 回收
    back_buckets = [torch.empty_like(local_out[0]) for _ in range(world)]
    dist.all_to_all(back_buckets, local_out)
    return torch.cat(back_buckets, dim=0)

def expert_sorted_batch(tokens, topk_idx):
    flat_experts = topk_idx.flatten()
    order = torch.argsort(flat_experts, stable=True)
    sorted_tokens = tokens.repeat_interleave(topk_idx.shape[1], dim=0)[order]
    sorted_experts = flat_experts[order]
    return sorted_tokens, sorted_experts, order

def moe_vs_dense_cost(n_experts, top_k, d_model, d_ff, n_layers):
    dense_ffn = 2 * d_model * d_ff
    moe_total = n_experts * dense_ffn
    moe_active = top_k * dense_ffn + 2 * d_model * n_experts   # 含门控
    return {
        "dense_equivalent_params": n_layers * dense_ffn,
        "moe_total_params": n_layers * moe_total,
        "moe_active_params": n_layers * moe_active,
        "activation_ratio": moe_active / moe_total,
    }

class MoELayer(nn.Module):
    def __init__(self, d_model, d_ff, n_experts=8, top_k=2, capacity_factor=1.25):
        super().__init__()
        self.n_experts = n_experts
        self.top_k = top_k
        self.capacity_factor = capacity_factor
        self.gating = TopKGating(d_model, n_experts, top_k)
        self.experts = nn.ModuleList([
            nn.Sequential(nn.Linear(d_model, d_ff), nn.GELU(),
                          nn.Linear(d_ff, d_model))
            for _ in range(n_experts)
        ])

    def forward(self, x):
        B, T, D = x.shape
        flat = x.reshape(-1, D)
        idx, val = self.gating(flat)                      # (N,K),(N,K)
        out = torch.zeros_like(flat)
        capacity = int(self.capacity_factor * flat.shape[0] / self.n_experts)

        for e in range(self.n_experts):
            # 找出路由到专家 e 的 token
            mask = (idx == e)
            token_pos, slot = mask.nonzero(as_tuple=True)
            if token_pos.numel() == 0:
                continue
            if token_pos.numel() > capacity:               # 容量溢出，丢弃
                token_pos = token_pos[:capacity]
                slot = slot[:capacity]
            expert_in = flat[token_pos]
            expert_out = self.experts[e](expert_in)
            out.index_add_(0, token_pos, expert_out * val[token_pos, slot].unsqueeze(-1))

        return out.reshape(B, T, D)