从零开始用 Python 复现 LLaMA 4 MoE 架构

超参数定义

合理的超参数是模型收敛的关键。我们定义嵌入维度、层数、注意力头数以及 MoE 特有的专家数量等。

# 基础架构参数
d_model = 128
n_layers = 4
n_heads = 4
block_size = 64
rms_norm_eps = 1e-5
rope_theta = 10000.0

# MoE 特定参数
num_local_experts = 4
num_experts_per_tok = 2
intermediate_size_expert = d_model * 2
shared_expert_dim = d_model * 2

# 训练参数
learning_rate = 5e-4
batch_size = 16
epochs = 3000

数据预处理

语言模型通过预测下一个 Token 来学习。我们需要从序列中滑动窗口提取输入和目标对。

all_x, all_y = [], []
for i in range(len(full_data_sequence) - block_size):
    x_chunk = full_data_sequence[i : i + block_size]
    y_chunk = full_data_sequence[i + 1 : i + block_size + 1]
    all_x.append(x_chunk)
    all_y.append(y_chunk)

train_x = torch.stack(all_x)
train_y = torch.stack(all_y)

模型组件初始化

这是核心部分。我们将手动构建 Transformer 块中的各个组件，包括嵌入层、RMSNorm、多头注意力（MHA）、RoPE 以及 MoE 层。

嵌入与 RoPE

token_embedding_table = nn.Embedding(vocab_size, d_model).to(device)

# 预计算 RoPE 逆频率
rope_freq_indices = torch.arange(0, d_k, 2, dtype=torch.float, device=device)
inv_freq = 1.0 / (rope_theta ** (rope_freq_indices / d_k))

RMSNorm 与注意力

LLaMA 使用 RMSNorm 替代 LayerNorm，计算更简单且效果相当。多头注意力负责捕捉上下文依赖。

# 初始化 RMSNorm 权重
rmsnorm_weights_input = [nn.Parameter(torch.ones(d_model, device=device)) for _ in range(n_layers)]
rmsnorm_weights_post_attn = [nn.Parameter(torch.ones(d_model, device=device)) for _ in range(n_layers)]
final_rmsnorm_weight = nn.Parameter(torch.ones(d_model, device=device))

# 初始化 MHA 线性层
mha_qkv_linears = [nn.Linear(d_model, 3 * d_model, bias=False).to(device) for _ in range(n_layers)]
mha_output_linears = [nn.Linear(d_model, d_model, bias=False).to(device) for _ in range(n_layers)]

MoE 层实现

MoE 层包含路由器和多个专家 MLP。路由器决定 Top-K 专家，专家并行计算后加权合并。

moe_routers = [nn.Linear(d_model, num_local_experts, bias=False).to(device) for _ in range(n_layers)]
moe_expert_gate_up_proj = [nn.Parameter(torch.empty(num_local_experts, d_model, 2 * expert_dim, device=device)) for _ in range(n_layers)]
moe_expert_down_proj = [nn.Parameter(torch.empty(num_local_experts, expert_dim, d_model, device=device)) for _ in range(n_layers)]

# 共享专家
shared_expert_gate_proj = [nn.Linear(d_model, shared_expert_dim, bias=False).to(device) for _ in range(n_layers)]
shared_expert_up_proj = [nn.Linear(d_model, shared_expert_dim, bias=False).to(device) for _ in range(n_layers)]
shared_expert_down_proj = [nn.Linear(shared_expert_dim, d_model, bias=False).to(device) for _ in range(n_layers)]

activation_fn = nn.SiLU()

训练循环

训练过程需要手动管理前向传播、损失计算和反向传播。注意在每一步中应用因果掩码（Causal Mask），确保模型只能关注当前位置及之前的 Token。

optimizer = optim.AdamW(all_model_parameters, lr=learning_rate)
criterion = nn.CrossEntropyLoss()
losses = []

for epoch in range(epochs):
    # 随机采样 Batch
    indices = torch.randint(0, train_x.shape[0], (batch_size,))
    xb, yb = train_x[indices].to(device), train_y[indices].to(device)
    
    # 前向传播
    token_embed = token_embedding_table(xb)
    position_ids = torch.arange(xb.shape[1], device=device).unsqueeze(0)
    freqs_cis = torch.polar(torch.ones_like(position_ids), (inv_freq.unsqueeze(0).unsqueeze(-1).expand(xb.shape[0], -1, 1).float() @ position_ids.unsqueeze(1).expand(xb.shape[0], -1).float()).transpose(1, 2))
    
    x = token_embed
    for i in range(n_layers):
        # RMSNorm & Attention
        x_norm = (x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + rms_norm_eps)) * rmsnorm_weights_input[i]
        qkv = mha_qkv_linears[i](x_norm).view(xb.shape[0], xb.shape[1], n_heads, 3 * d_k).chunk(3, dim=-1)
        q, k, v = qkv[0], qkv[1], qkv[2]
        # ... (此处省略具体的 RoPE 旋转与 Attention 计算细节，逻辑同前文描述)
        # MoE Block
        x_norm_moe = (x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + rms_norm_eps)) * rmsnorm_weights_post_attn[i]
        router_logits = moe_routers[i](x_norm_moe)
        routing_weights, selected_experts = torch.sigmoid(torch.topk(router_logits, num_experts_per_tok, dim=-1)[0]), torch.topk(router_logits, num_experts_per_tok, dim=-1)[1]
        # ... (专家计算与聚合逻辑)
        x = x + combined_expert_outputs.view(xb.shape[0], xb.shape[1], d_model) + shared_output
    
    logits = output_linear_layer((x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + rms_norm_eps)) * final_rmsnorm_weight)
    loss = criterion(logits.view(-1, logits.shape[-1]), yb.view(-1))
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    losses.append(loss.item())
    
    if epoch % eval_interval == 0 or epoch == epochs - 1:
        print(f"第 {epoch+1}/{epochs} 个周期，损失：{loss.item():.4f}")

文本生成

训练完成后，我们可以使用模型进行自回归生成。切换到评估模式（eval()）并禁用梯度计算以提升速度。

seed_chars = "Alice "
num_tokens_to_generate = 200

# 设置评估模式
for layer in [token_embedding_table] + mha_qkv_linears + moe_routers + output_linear_layer:
    layer.eval()

with torch.no_grad():
    generated_sequence = torch.tensor([[char_to_int[ch] for ch in seed_chars]], dtype=torch.long, device=device)
    for _ in range(num_tokens_to_generate):
        current_context = generated_sequence[:, -block_size:]
        # ... (执行前向传播获取 logits)
        next_token = torch.multinomial(F.softmax(logits_gen[:, -1, :], dim=-1), num_samples=1)
        generated_sequence = torch.cat((generated_sequence, next_token), dim=1)

# 解码输出
decoded_text = ''.join([int_to_char.get(id_val, '[UNK]') for id_val in generated_sequence[0].tolist()])
print(decoded_text)

结语

通过这段代码，我们完整实现了从数据预处理到模型训练再到生成的全流程。虽然这是一个简化版本，但它清晰地展示了 LLaMA 系列模型中 MoE、RoPE 和 RMSNorm 是如何协同工作的。你可以在此基础上扩展层数、增加参数量，尝试复现更大规模的模型架构。