0%
💡 왼쪽 원문을 읽으면서 오른쪽에 따라 써보세요. Tab 키로 힌트를 받을 수 있습니다.
1. 什么是 MoE？
Mixture of Experts（MoE）是一种只激活模型全部参数中一部分、从而提升运算效率的架构。与对所有输入都使用完整模型的 Dense 模型不同，MoE 会根据输入只选择最优的专家（Expert）来处理。
Dense vs Sparse 模型
- Dense 模型：每次都会激活全部参数（例如 LLaMA、GPT-4）
- Sparse MoE：只激活全部参数中的一部分（例如 Mixtral、DeepSeek-V3）
核心优势在于参数量大、但运算量小。Mixtral 8x7B 总共拥有 46.7B 参数，但推理时实际激活的参数只有约 12.9B。
2. MoE 架构的核心组成部分
Expert Network
每个 Expert 都是一个独立的 FFN（Feed-Forward Network）：
import torch
import torch.nn as nn
class Expert(nn.Module):
def __init__(self, d_model: int, d_ff: int):
super().__init__()
self.w1 = nn.Linear(d_model, d_ff, bias=False)
self.w2 = nn.Linear(d_ff, d_model, bias=False)
self.w3 = nn.Linear(d_model, d_ff, bias=False) # SwiGLU gate
def forward(self, x: torch.Tensor) -> torch.Tensor:
# SwiGLU activation
return self.w2(nn.functional.silu(self.w1(x)) * self.w3(x))
Router（Gating Network）
Router 负责决定把每个 token 发送给哪个 Expert：
class TopKRouter(nn.Module):
def __init__(self, d_model: int, num_experts: int, top_k: int = 2):
super().__init__()
self.gate = nn.Linear(d_model, num_experts, bias=False)
self.top_k = top_k
def forward(self, x: torch.Tensor):
# x shape: (batch, seq_len, d_model)
logits = self.gate(x) # (batch, seq_len, num_experts)
top_k_logits, top_k_indices = logits.topk(self.top_k, dim=-1)
top_k_weights = torch.softmax(top_k_logits, dim=-1)
return top_k_weights, top_k_indices
MoE Layer 完整实现
class MoELayer(nn.Module):
def __init__(self, d_model: int, d_ff: int,
num_experts: int = 8, top_k: int = 2):
super().__init__()
self.experts = nn.ModuleList([
Expert(d_model, d_ff) for _ in range(num_experts)
])
self.router = TopKRouter(d_model, num_experts, top_k)
self.num_experts = num_experts
def forward(self, x: torch.Tensor) -> torch.Tensor:
batch_size, seq_len, d_model = x.shape
weights, indices = self.router(x)
# Reshape for expert processing
flat_x = x.view(-1, d_model)
flat_weights = weights.view(-1, weights.shape[-1])
flat_indices = indices.view(-1, indices.shape[-1])
output = torch.zeros_like(flat_x)
for i, expert in enumerate(self.experts):
# Find tokens routed to this expert
mask = (flat_indices == i).any(dim=-1)
if mask.any():
expert_input = flat_x[mask]
expert_output = expert(expert_input)
# Weight by router probability
idx = (flat_indices[mask] == i).float()
w = (flat_weights[mask] * idx).sum(dim=-1, keepdim=True)
output[mask] += w * expert_output
return output.view(batch_size, seq_len, d_model)
3. 主要 MoE 模型分析
Mixtral 8x7B（Mistral AI）
- 8 个 Expert，Top-2 路由
- 总计 46.7B 参数，激活 12.9B
- Attention 层共享，只有 FFN 被拆分为 Expert
DeepSeek-V3 MoE
DeepSeek-V3 采用了更精巧的 MoE 设计：
class DeepSeekMoE(nn.Module):
"""DeepSeek-V3 风格：共享 Expert + Routed Expert"""
def __init__(self, d_model, d_ff, num_shared=1,
num_routed=256, top_k=8):
super().__init__()
# 所有 token 都会经过的共享 Expert
self.shared_experts = nn.ModuleList([
Expert(d_model, d_ff) for _ in range(num_shared)
])
# 按 token 选择的 Routed Expert
self.routed_experts = nn.ModuleList([
Expert(d_model, d_ff // 4) for _ in range(num_routed)
])
self.router = TopKRouter(d_model, num_routed, top_k)
def forward(self, x):
# 共享 Expert 的结果
shared_out = sum(e(x) for e in self.shared_experts)
# Routed Expert 的结果
weights, indices = self.router(x)
routed_out = self._route_tokens(x, weights, indices)
return shared_out + routed_out
- 1 个共享 Expert + 256 个 Routed Expert（Top-8 选择）
- 总计 671B 参数，激活 37B
- 引入了 Auxiliary-loss-free 负载均衡
模型对比
| 模型 | 总参数 | 激活参数 | Expert 数 | Top-K |
|---|---|---|---|---|
| Mixtral 8x7B | 46.7B | 12.9B | 8 | 2 |
| Mixtral 8x22B | 141B | 39B | 8 | 2 |
| DeepSeek-V3 | 671B | 37B | 256+1 | 8+1 |
| Qwen2.5-MoE | 14.3B | 2.7B | 60+4 | 4+4 |
4. 路由策略
Token Choice vs Expert Choice
# Token Choice：每个 token 选择自己的 Expert
def token_choice_routing(logits, top_k=2):
top_k_vals, top_k_idx = logits.topk(top_k, dim=-1)
weights = torch.softmax(top_k_vals, dim=-1)
return weights, top_k_idx
# Expert Choice：每个 Expert 选择自己的 token
def expert_choice_routing(logits, capacity_factor=1.25):
num_tokens = logits.shape[0]
num_experts = logits.shape[1]
capacity = int(num_tokens * capacity_factor / num_experts)
expert_scores = logits.T # (num_experts, num_tokens)
top_k_vals, top_k_idx = expert_scores.topk(capacity, dim=-1)
return top_k_vals, top_k_idx
5. 负载均衡
Expert 之间的负载不均衡是 MoE 的核心问题：
def load_balancing_loss(router_logits, top_k_indices, num_experts):
"""Auxiliary load balancing loss（Switch Transformer 方式）"""
# 每个 Expert 的 token 比例
mask = torch.zeros_like(router_logits)
mask.scatter_(-1, top_k_indices, 1.0)
tokens_per_expert = mask.float().mean(dim=0) # (num_experts,)
# 每个 Expert 的平均路由概率
router_probs = torch.softmax(router_logits, dim=-1)
router_prob_per_expert = router_probs.mean(dim=0)
# 两个分布的点积 = 不均衡程度的度量
loss = num_experts * (tokens_per_expert * router_prob_per_expert).sum()
return loss
DeepSeek-V3 的 Auxiliary-loss-free 方式通过动态调整每个 Expert 的 bias term，在不引入额外 loss 的情况下维持均衡。
6. 推理优化
# Expert Parallelism：把 Expert 分散到多个 GPU
# GPU 0: Expert 0-3, GPU 1: Expert 4-7
class ExpertParallel(nn.Module):
def __init__(self, experts_per_gpu, rank, world_size):
super().__init__()
self.local_experts = nn.ModuleList([
Expert(d_model, d_ff)
for _ in range(experts_per_gpu)
])
self.rank = rank
self.world_size = world_size
def forward(self, x, indices):
# 通过 All-to-All 通信重新分配 token
dispatched = all_to_all(x, indices, self.world_size)
# 用本地 Expert 处理
output = self._process_local(dispatched)
# 重新合并结果
return all_to_all(output, indices, self.world_size)
7. 小测验
Q1：在 Mixtral 8x7B 中，一个 token 在推理时使用的参数量大约是多少？
大约 12.9B 参数。因为 8 个 Expert 中只有 Top-2 被激活，所以只会用到 2 个 Expert 的 FFN 参数，加上共享的 Attention 参数。这大约是全部 46.7B 的 28%。
Q2：DeepSeek-V3 引入的 Auxiliary-loss-free load balancing 的核心思路是什么？
传统 MoE 会为了 load balancing 额外加入 auxiliary loss，而这可能会降低模型性能。DeepSeek-V3 为每个 Expert 设置了动态的 bias term，把接收 token 较多的 Expert 的 bias 调低，把接收较少的 Expert 的 bias 调高，从而自然达成均衡。这样一来，不需要额外的 loss 也能实现稳定的 load balancing。
Q3：Token Choice 与 Expert Choice routing 的区别，以及各自的优缺点是什么？
- Token Choice：每个 token 选择 Top-K 个 Expert。实现简单，但可能出现 token 集中到部分 Expert 上的不均衡问题。
- Expert Choice：每个 Expert 选择自己要处理的 token。能保证完美的 load balancing，但部分 token 可能不会被任何 Expert 选中。
实际中最常用的组合是 Token Choice + load balancing loss。
현재 단락 (1/119)
**Mixture of Experts**（MoE）是一种只激活模型全部参数中一部分、从而提升运算效率的架构。与对所有输入都使用完整模型的 Dense 模型不同，MoE 会根据输入只选择**最优的专...
작성 글자: 0자원문 글자: 5,297자작성 단락: 0/119