CodeLab-LLaMA2 / 星语 MoE：大模型训练实践

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from transformers import PretrainedConfig


# ============================================================
# 模型配置类
# ============================================================
class ModelConfig(PretrainedConfig):
    *"""*
*    模型的全局配置参数类，用于统一管理 Transformer + MoE 层的结构参数。*
*    继承自 HuggingFace 的 PretrainedConfig，方便保存与加载。*
*    """*
*    *model_type = "XY-MoE"

    def __init__(self,
                 dim: int = 1024,           # 模型主维度（token embedding 的维度）
                 n_layers: int = 16,        # Transformer 层数
                 n_heads: int = 16,         # 多头注意力的头数
                 n_kv_heads: int = 8,       # key/value 头数
                 vocab_size: int = 6144,    # 词表大小
                 hidden_dim: int = None,    # FFN 隐藏层维度（若为 None 会自动计算）
                 multiple_of: int = 64,     # FFN 隐藏层对齐倍数
                 norm_eps: float = 1e-5,    # LayerNorm 的 epsilon
                 max_seq_len: int = 256,    # 最大序列长度
                 dropout: float = 0.0,      # Dropout 概率
                 flash_attn: bool = True,   # 是否启用 Flash Attention

                 ########################################################
                 # 下方为 MoE（Mixture of Experts）配置参数
                 ########################################################
                 num_experts_per_tok: int = 1,   # 每个 token 选择的专家数（top-k）
                 n_routed_experts: int = 4,      # MoE 专家总数
                 n_shared_experts: int = 1,      # “共享”专家数（dense FFN，不经过 gating）
                 gate_hidden_size: int = None,   # gate 输入维度（若 None 默认等于 dim）
                 aux_loss_alpha: float = 1e-2,   # MoE 辅助损失权重，用于均衡专家负载
                 seq_aux: bool = False,          # 辅助损失是否按序列统计
                 noise_std: float = 1.0,         # gating logits 噪声强度
                 norm_topk_prob: bool = True,    # 是否归一化 top-k 概率
                 capacity_factor: float = 1.25,  # 每个专家的 token 容量上限因子
                 **kwargs,
                 ):
        *"""*
*        初始化模型参数。这里主要保存为实例属性，供下游模块使用。*
*        """*
*        *# Transformer 主体参数
        self.dim = dim
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.vocab_size = vocab_size
        self.hidden_dim = hidden_dim
        self.multiple_of = multiple_of
        self.norm_eps = norm_eps
        self.max_seq_len = max_seq_len
        self.dropout = dropout
        self.flash_attn = flash_attn

        # MoE 参数
        self.num_experts_per_tok = int(num_experts_per_tok)
        self.n_routed_experts = int(n_routed_experts)
        self.n_shared_experts = int(n_shared_experts)
        self.gate_hidden_size = gate_hidden_size if gate_hidden_size is not None else dim
        self.aux_loss_alpha = aux_loss_alpha
        self.seq_aux = seq_aux
        self.noise_std = noise_std if noise_std is not None else 0.0
        self.norm_topk_prob = bool(norm_topk_prob)
        self.capacity_factor = float(capacity_factor)

        super().__init__(**kwargs)


# ============================================================
# SwiGLU 前馈网络（Feed-Forward Network）
# ============================================================
class FFN(nn.Module):
    *"""*
*    SwiGLU 变体的前馈网络层。*
*    结构：*
*        x -> fc_gate -> SiLU() 激活*
*        x -> fc_up -> 线性映射*
*        逐元素相乘 -> fc_down -> Dropout*
*    输出维度与输入相同（残差连接用）。*
*    """*

*    *def __init__(self, dim: int, hidden_dim: int = None, multiple_of: int = 256, dropout: float = 0.0):
        super().__init__()

        # 若未显式指定 hidden_dim，则自动计算（与 LLaMA 一致）
        if hidden_dim is None:
            hidden_dim = 4 * dim
            hidden_dim = int(2 * hidden_dim / 3)  # 约 2.67x dim
            hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)

        # 线性层定义
        self.fc_gate = nn.Linear(dim, hidden_dim, bias=False)  # 激活分支
        self.fc_up = nn.Linear(dim, hidden_dim, bias=False)    # 门控分支
        self.fc_down = nn.Linear(hidden_dim, dim, bias=False)  # 输出投影
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        *"""*
*        x: [batch, seq_len, dim]*
*        return: [batch, seq_len, dim]*
*        """*
*        *gated = F.silu(self.fc_gate(x))  # 激活通道
        up = self.fc_up(x)               # 线性通道
        hidden = gated * up              # SwiGLU 核心操作
        out = self.fc_down(hidden)       # 映射回原维度
        out = self.dropout(out)
        return out


# ============================================================
# MoE 的 Gate 模块
# ============================================================
class MoEGate(nn.Module):
    *"""*
*    Gating 网络：负责为每个 token 计算其对应的专家分配概率。*

*    输入：*
*        x: [B, L, D] （token embedding）*

*    输出：*
*        topk_idx:    [B*L, K]  每个 token 对应的 top-k 专家索引*
*        topk_weight: [B*L, K]  对应的 top-k 概率权重*
*        aux_loss:    标量，用于负载均衡的辅助损失*
*    """*

*    *def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config

        # ---------------- 基本参数 ----------------
        self.top_k = config.num_experts_per_tok         # 每个 token 选择几个专家
        self.n_experts = config.n_routed_experts        # 专家总数
        self.gating_dim = config.gate_hidden_size       # gate 的输入维度
        self.alpha = config.aux_loss_alpha              # 辅助损失权重
        self.seq_aux = config.seq_aux
        self.noise_std = config.noise_std
        self.norm_topk_prob = config.norm_topk_prob

        # ---------------- 权重定义 ----------------
        # 线性层：输入 gating_dim → 输出 n_experts
        # 输出即每个专家的 logits（未归一化分数）
        self.linear = nn.Linear(self.gating_dim, self.n_experts, bias=False)

        # 初始化参数
        self.reset_parameters()

    def reset_parameters(self):
        *"""使用 Kaiming 初始化线性层权重。"""*
*        *nn.init.kaiming_uniform_(self.linear.weight, a=math.sqrt(5))

    def forward(self, x: torch.Tensor):
        *"""*
*        Args:*
*            x: [B, L, D] token 表示*
*        Returns:*
*            topk_idx, topk_weight, aux_loss*
*        """*
*        *B, L, D = x.shape

        # 若输入维度与 gate_hidden_size 不一致，可以添加线性投影
        if D != self.gating_dim:
            # 临时线性投影（非参数化版，为演示方便）
            proj = nn.Linear(D, self.gating_dim, bias=False).to(x.device)
            nn.init.kaiming_uniform_(proj.weight, a=math.sqrt(5))
            x_flat = proj(x.view(-1, D))  # [B*L, gating_dim]
        else:
            x_flat = x.view(-1, D)

        # 线性层得到 logits
        logits = self.linear(x_flat)  # [B*L, n_experts]

        # 训练时加噪声（Switch Transformer trick）
        if self.training and self.noise_std > 0:
            logits = logits + torch.randn_like(logits) * self.noise_std

        # softmax 得到每个专家的选择概率
        scores = F.softmax(logits, dim=-1)  # [B*L, n_experts]

        # 取 top-k 专家
        topk_vals, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)  # [B*L, K]

        # 对 top-k 权重归一化（保持总和=1）
        if self.top_k > 1 and self.norm_topk_prob:
            denom = topk_vals.sum(dim=-1, keepdim=True) + 1e-20
            topk_vals = topk_vals / denom

        # ---------------- 辅助损失（aux loss） ----------------
        # 用于防止部分专家被频繁选中，保持负载均衡
        if self.training and (self.alpha is not None) and (self.alpha > 0):
            # 各专家平均被选中的概率
            expert_prob = scores.mean(dim=0)  # 每个专家的平均 softmax 概率
            # 每个专家实际被选中的比例
            flat_topk_idx = topk_idx.view(-1)
            mask = F.one_hot(flat_topk_idx, num_classes=self.n_experts).float()
            expert_usage = mask.mean(dim=0)
            # 损失：鼓励 expert_prob 与 expert_usage 均匀
            aux_loss = self.alpha * torch.sum(expert_prob * expert_usage * float(self.n_experts))
        else:
            aux_loss = torch.tensor(0.0, device=x.device, dtype=x.dtype)

        return topk_idx, topk_vals, aux_loss


# ============================================================
# MoE FeedForward（核心层）
# ============================================================
class MOEFeedForward(nn.Module):
    *"""*
*    MoE 前馈层，将多个专家的 FFN 与 Gate 结合。*
*    流程：*
*      1. gate 计算每个 token 的 top-k 专家及权重；*
*      2. 将 token 分派到各自专家（dispatch）；*
*      3. 专家独立计算；*
*      4. 按权重加权求和回原序列；*
*      5. 若存在 shared expert，则叠加输出。*
*    """*

*    *def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config
        n_experts = config.n_routed_experts

        # 1创建专家列表（每个专家一个 FFN）
        self.experts = nn.ModuleList([
            FFN(config.dim, config.hidden_dim, config.multiple_of, config.dropout)
            for _ in range(n_experts)
        ])

        # 2创建 gate 模块
        self.gate = MoEGate(config)

        # 3可选的 shared experts（dense 路径）
        if getattr(config, 'n_shared_experts', 0) > 0:
            self.shared_experts = nn.ModuleList([
                FFN(config.dim, config.hidden_dim, config.multiple_of, config.dropout)
                for _ in range(config.n_shared_experts)
            ])
        else:
            self.shared_experts = None

        # 4capacity 系数
        self.capacity_factor = getattr(config, 'capacity_factor', 1.25)
        self.aux_loss = None  # 存储 aux loss

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        *"""*
*        Args:*
*            x: [B, L, D]*
*        Returns:*
*            y_out: [B, L, D]*
*        """*
*        *B, L, D = x.shape
        T0 = B * L  # flatten token 数量
        K = self.config.num_experts_per_tok

        # ========== Step 1: gating ==========
        topk_idx, topk_weight, aux_loss = self.gate(x)
        self.aux_loss = aux_loss  # 暴露给外部使用（训练时可累加）

        # 展平成 (T0, D)
        x_flat = x.view(-1, D)

        # ========== Step 2: 复制 token ==========
        # 每个 token 对应 K 条路径（用于 top-k 专家）
        x_expanded = x_flat.repeat_interleave(K, dim=0)   # (T0*K, D)
        topk_idx_flat = topk_idx.view(-1)                 # (T0*K,)
        topk_weight_flat = topk_weight.view(-1)           # (T0*K,)

        device = x.device

        # ========== Step 3: capacity 控制 ==========
        # 每个专家最多处理的 token 数
        tokens_per_expert = float(T0) / float(self.config.n_routed_experts)
        capacity = max(1, int(math.ceil(self.capacity_factor * tokens_per_expert)))

        # 记录每个专家当前分配的 token 数
        counts = torch.zeros(self.config.n_routed_experts, dtype=torch.long, device=device)
        keep_mask = torch.zeros_like(topk_idx_flat, dtype=torch.bool, device=device)

        # 简单实现：顺序遍历 token，若专家未满则保留
        for i in range(topk_idx_flat.size(0)):
            e = int(topk_idx_flat[i].item())
            if counts[e] < capacity:
                keep_mask[i] = True
                counts[e] += 1
            else:
                keep_mask[i] = False  # 超出容量的 token 被丢弃

        # 仅保留有效 token
        selected_idx = topk_idx_flat[keep_mask]
        selected_x = x_expanded[keep_mask]
        selected_weight = topk_weight_flat[keep_mask]

        # ========== Step 4: dispatch 到专家 ==========
        y_expanded = torch.zeros_like(x_expanded, device=device)
        keep_indices = torch.nonzero(keep_mask, as_tuple=False).view(-1)  # 对应位置

        # 按专家分组计算（易读实现）
        for expert_id in range(self.config.n_routed_experts):
            mask_expert = (selected_idx == expert_id)
            if mask_expert.any():
                expert_in = selected_x[mask_expert]          # 输入 token
                expert_out = self.experts[expert_id](expert_in)  # 专家输出
                expert_positions = keep_indices[mask_expert]  # 写回位置
                y_expanded[expert_positions] = expert_out

        # ========== Step 5: 合并 top-k 输出 ==========
        # 将 K 路专家输出按权重加权求和
        y_expanded_view = y_expanded.view(T0, K, D)
        topk_weight_view = topk_weight.view(T0, K, 1)
        y_combined = torch.sum(y_expanded_view * topk_weight_view.to(y_expanded_view.dtype), dim=1)
        y_out = y_combined.view(B, L, D)

        # ========== Step 6: 加上 shared experts ==========
        if self.shared_experts is not None:
            for expert in self.shared_experts:
                y_out = y_out + expert(x)

        return y_out


if __name__ == "__main__":
    cfg = ModelConfig(dim=64, n_routed_experts=4, num_experts_per_tok=2)
    moe = MOEFeedForward(cfg)

    x = torch.randn(2, 8, 64)  # batch=2, seq=8
    y = moe(x)
    print("输入形状:", x.shape)
    print("输出形状:", y.shape)
    print("辅助损失:", moe.aux_loss.item())