Mixture of Experts（MoE）（二）：大模型的稀疏化架构

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

class Expert(nn.Module):
    """单个专家网络（标准 FFN with SwiGLU）"""
    def __init__(self, dim: int, hidden_dim: int):
        super().__init__()
        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
        self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # SwiGLU 激活: Swish(gate) * up
        gate = F.silu(self.gate_proj(x))
        up = self.up_proj(x)
        return self.down_proj(gate * up)


class MoELayer(nn.Module):
    """Mixture of Experts 层"""
    def __init__(self, dim: int, hidden_dim: int, 
                 num_experts: int, top_k: int = 2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        
        # 创建专家组
        self.experts = nn.ModuleList([
            Expert(dim, hidden_dim) for _ in range(num_experts)
        ])
        
        # 门控路由器
        self.gate = nn.Linear(dim, num_experts, bias=False)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        x: (batch, seq_len, dim)
        返回: (batch, seq_len, dim)
        """
        batch, seq_len, dim = x.shape
        
        # 1. 计算门控分数: (batch, seq_len, num_experts)
        gate_logits = self.gate(x)
        gate_weights = F.softmax(gate_logits, dim=-1)
        
        # 2. 选择 Top-K 专家
        topk_weights, topk_indices = torch.topk(
            gate_weights, self.top_k, dim=-1
        )
        # 归一化 Top-K 权重
        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
        
        # 3. 路由到专家并加权求和
        output = torch.zeros_like(x)
        for expert_idx in range(self.num_experts):
            # 找出哪些 token 选择了这个专家
            mask = (topk_indices == expert_idx)  # (batch, seq_len, top_k)
            token_mask = mask.any(dim=-1)        # (batch, seq_len)
            
            if token_mask.any():
                # 获取权重
                expert_weights = torch.zeros(
                    batch, seq_len, device=x.device
                )
                for k in range(self.top_k):
                    expert_weights += topk_weights[:, :, k] * (
                        topk_indices[:, :, k] == expert_idx
                    ).float()
                
                # 专家处理
                expert_input = x[token_mask]
                expert_output = self.experts[expert_idx](expert_input)
                
                # 加权写回
                weighted_output = expert_output * expert_weights[
                    token_mask
                ].unsqueeze(-1)
                output[token_mask] += weighted_output
        
        return output

def compute_load_balance_loss(
    gate_weights: torch.Tensor,  # (batch, seq, num_experts)
    topk_indices: torch.Tensor,  # (batch, seq, top_k)
    num_experts: int,
) -> torch.Tensor:
    """
    Switch Transformer 负载均衡辅助损失
    
    目标：让每个专家处理大致相同数量的 token
    损失 = num_experts * sum(f_i * P_i)
    其中 f_i = 专家 i 被选中的比例
         P_i = 专家 i 的平均门控权重
    """
    batch, seq_len, top_k = topk_indices.shape
    
    # 1. 计算每个专家被选中的频率 f_i
    expert_counts = torch.zeros(
        num_experts, device=gate_weights.device
    )
    for k in range(top_k):
        indices = topk_indices[:, :, k].flatten()
        expert_counts += torch.bincount(
            indices, minlength=num_experts
        ).float()
    expert_counts = expert_counts / expert_counts.sum()
    
    # 2. 计算每个专家的平均门控权重 P_i
    # 展开所有门控权重
    flat_gate = gate_weights.reshape(-1, num_experts)
    mean_weights = flat_gate.mean(dim=0)
    
    # 3. 计算辅助损失
    aux_loss = num_experts * (expert_counts * mean_weights).sum()
    return aux_loss


# 使用示例：添加到总损失中
# total_loss = task_loss + 0.01 * aux_loss
# aux_loss 的权重通常设为 0.01 量级
print("负载均衡辅助损失函数已定义")

# MoE 模型的参数量计算
def calculate_moe_params(
    num_layers: int,
    dim: int,
    hidden_dim: int,
    num_experts_per_layer: int,
    num_shared_experts: int = 0,
):
    """
    计算 MoE 模型的总参数量和激活参数量
    
    参数:
        num_layers: Transformer 层数
        dim: 隐藏层维度
        hidden_dim: FFN 中间维度
        num_experts_per_layer: 每层专家数
        num_shared_experts: 每层共享专家数
    
    返回:
        (total_params, active_params)
    """
    # 每个专家的参数量（SwiGLU FFN: gate + up + down）
    expert_params = dim * hidden_dim * 2 + hidden_dim * dim
    
    # 路由器的参数量
    router_params = dim * num_experts_per_layer
    
    # 每层的参数量
    # 自注意力 + MoE FFN + 归一化
    attention_params = 4 * dim * dim  # Q, K, V, O 投影
    layer_params = (
        attention_params +
        num_experts_per_layer * expert_params +
        router_params +
        4 * dim  # 归一化层参数
    )
    
    total_params = num_layers * layer_params
    
    # 激活参数量（假设 Top-2 路由 + 共享专家）
    active_experts = 2 + num_shared_experts
    active_layer_params = (
        attention_params +
        active_experts * expert_params +
        router_params +
        4 * dim
    )
    active_params = num_layers * active_layer_params
    
    return total_params, active_params

# Mixtral 8x7B 近似
total, active = calculate_moe_params(
    num_layers=32, dim=4096, hidden_dim=14336,
    num_experts_per_layer=8
)
print(f"Mixtral 8x7B:")
print(f"  总参量: {total/1e9:.1f}B")
print(f"  激活参量: {active/1e9:.1f}B")
print(f"  稀疏率: {(1 - active/total)*100:.1f}%")