makeMoE: 从零开始实现稀疏混合专家语言模型
稀疏混合专家(Sparse Mixture of Experts, MoE)模型是近年来大型语言模型领域的一个重要发展方向。本文将详细介绍如何从头实现一个基于MoE架构的语言模型,该项目名为makeMoE,受到了Andrej Karpathy的makemore项目的启发。
项目背景
makeMoE是一个自回归字符级语言模型,采用了稀疏混合专家架构。与传统Transformer模型不同,MoE模型在每个Transformer块中使用多个专家网络替代了单一的前馈网络。这种设计可以显著增加模型参数量,同时保持计算效率。
本项目的主要目标是:
- 深入理解MoE模型的工作原理
- 提供一个可读性强、易于修改的PyTorch实现
- 为研究人员和开发者提供一个实验平台
模型架构
makeMoE模型的核心组件包括:
- 多头自注意力机制
- 专家模块(Expert)
- Top-k门控网络
- 带噪声的Top-k门控(用于负载均衡)
让我们逐一深入了解这些组件。
多头自注意力机制
自注意力是Transformer模型的核心,MoE模型保留了这一机制。以下是因果自注意力头的实现:
class Head(nn.Module):
def __init__(self, head_size):
super().__init__()
self.key = nn.Linear(n_embd, head_size, bias=False)
self.query = nn.Linear(n_embd, head_size, bias=False)
self.value = nn.Linear(n_embd, head_size, bias=False)
self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))
self.dropout = nn.Dropout(dropout)
def forward(self, x):
B,T,C = x.shape
k = self.key(x)
q = self.query(x)
wei = q @ k.transpose(-2,-1) * C**-0.5
wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
wei = F.softmax(wei, dim=-1)
wei = self.dropout(wei)
v = self.value(x)
out = wei @ v
return out
多头注意力机制通过并行应用多个注意力头来提高模型的表达能力:
class MultiHeadAttention(nn.Module):
def __init__(self, num_heads, head_size):
super().__init__()
self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
self.proj = nn.Linear(n_embd, n_embd)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
out = torch.cat([h(x) for h in self.heads], dim=-1)
out = self.dropout(self.proj(out))
return out
专家模块
专家模块本质上是一个简单的多层感知机(MLP):
class Expert(nn.Module):
def __init__(self, n_embd):
super().__init__()
self.net = nn.Sequential(
nn.Linear(n_embd, 4 * n_embd),
nn.ReLU(),
nn.Linear(4 * n_embd, n_embd),
nn.Dropout(dropout),
)
def forward(self, x):
return self.net(x)
Top-k门控网络
门控网络决定了每个token应该由哪些专家处理。Top-k门控选择k个最相关的专家:
class TopkRouter(nn.Module):
def __init__(self, n_embed, num_experts, top_k):
super(TopkRouter, self).__init__()
self.top_k = top_k
self.linear = nn.Linear(n_embed, num_experts)
def forward(self, mh_output):
logits = self.linear(mh_output)
top_k_logits, indices = logits.topk(self.top_k, dim=-1)
zeros = torch.full_like(logits, float('-inf'))
sparse_logits = zeros.scatter(-1, indices, top_k_logits)
router_output = F.softmax(sparse_logits, dim=-1)
return router_output, indices
带噪声的Top-k门控
为了实现更好的负载均衡,我们可以在门控网络中引入噪声:
class NoisyTopkRouter(nn.Module):
def __init__(self, n_embed, num_experts, top_k):
super(NoisyTopkRouter, self).__init__()
self.top_k = top_k
self.topkroute_linear = nn.Linear(n_embed, num_experts)
self.noise_linear = nn.Linear(n_embed, num_experts)
def forward(self, mh_output):
logits = self.topkroute_linear(mh_output)
noise_logits = self.noise_linear(mh_output)
noise = torch.randn_like(logits) * F.softplus(noise_logits)
noisy_logits = logits + noise
top_k_logits, indices = noisy_logits.topk(self.top_k, dim=-1)
zeros = torch.full_like(noisy_logits, float('-inf'))
sparse_logits = zeros.scatter(-1, indices, top_k_logits)
router_output = F.softmax(sparse_logits, dim=-1)
return router_output, indices
稀疏混合专家模块
将以上组件组合起来,我们得到了完整的稀疏混合专家模块:
class SparseMoE(nn.Module):
def __init__(self, n_embed, num_experts, top_k):
super(SparseMoE, self).__init__()
self.router = NoisyTopkRouter(n_embed, num_experts, top_k)
self.experts = nn.ModuleList([Expert(n_embed) for _ in range(num_experts)])
self.top_k = top_k
def forward(self, x):
gating_output, indices = self.router(x)
final_output = torch.zeros_like(x)
flat_x = x.view(-1, x.size(-1))
flat_gating_output = gating_output.view(-1, gating_output.size(-1))
for i, expert in enumerate(self.experts):
expert_mask = (indices == i).any(dim=-1)
flat_mask = expert_mask.view(-1)
if flat_mask.any():
expert_input = flat_x[flat_mask]
expert_output = expert(expert_input)
gating_scores = flat_gating_output[flat_mask, i].unsqueeze(1)
weighted_output = expert_output * gating_scores
final_output[expert_mask] += weighted_output.squeeze(1)
return final_output
完整的MoE语言模型
最后,我们将所有组件组合成一个完整的MoE语言模型:
class SparseMoELanguageModel(nn.Module):
def __init__(self):
super().__init__()
self.token_embedding_table = nn.Embedding(vocab_size, n_embed)
self.position_embedding_table = nn.Embedding(block_size, n_embed)
self.blocks = nn.Sequential(*[Block(n_embed, n_head=n_head, num_experts=num_experts, top_k=top_k) for _ in range(n_layer)])
self.ln_f = nn.LayerNorm(n_embed)
self.lm_head = nn.Linear(n_embed, vocab_size)
def forward(self, idx, targets=None):
B, T = idx.shape
tok_emb = self.token_embedding_table(idx)
pos_emb = self.position_embedding_table(torch.arange(T, device=device))
x = tok_emb + pos_emb
x = self.blocks(x)
x = self.ln_f(x)
logits = self.lm_head(x)
if targets is None:
loss = None
else:
B, T, C = logits.shape
logits = logits.view(B*T, C)
targets = targets.view(B*T)
loss = F.cross_entropy(logits, targets)
return logits, loss
def generate(self, idx, max_new_tokens):
for _ in range(max_new_tokens):
idx_cond = idx[:, -block_size:]
logits, loss = self(idx_cond)
logits = logits[:, -1, :]
probs = F.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1)
idx = torch.cat((idx, idx_next), dim=1)
return idx
模型训练与评估
模型的训练过程采用了标准的语言模型训练方法。我们使用AdamW优化器,并定期评估模型在训练集和验证集上的损失。
model = SparseMoELanguageModel()
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
for iter in range(max_iters):
if iter % eval_interval == 0:
losses = estimate_loss()
print(f"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")
xb, yb = get_batch('train')
logits, loss = model(xb, yb)
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
结论与未来方向
makeMoE项目展示了如何从零开始实现一个稀疏混合专家语言模型。尽管这个实现更注重可读性和可修改性,而非性能优化,但它为理解MoE模型的工作原理提供了宝贵的洞察。
未来的改进方向包括:
- 优化稀疏激活的实现,提高计算效率
- 尝试不同的初始化策略,如Xavier/Glorot初始化
- 从字符级扩展到子词级别的分词
- 实现专家容量(Expert Capacity)机制
- 探索MoE与多模态模型的结合
通过深入理解和实现MoE模型,我们不仅能够掌握这一先进的模型架构,还能为未来的语言模型研究和开发奠定基础。makeMoE项目为研究人员和开发者提供了一个灵活的实验平台,期待看到更多基于此的创新和突破。
