用PyTorch从零开始编写DeepSeek-V2

DeepSeek-V2是一个强大的开源混合专家（MoE）语言模型，通过创新的Transformer架构实现了经济高效的训练和推理。该模型总共拥有2360亿参数，其中每个令牌激活21亿参数，支持最大128K令牌的上下文长度。

在开源模型中，DeepSeek-V2实现了顶级性能，成为最强大的开源MoE语言模型。在MMLU（多模态机器学习）上，DeepSeek-V2以较少的激活参数实现了顶尖的性能。与DeepSeek 67B相比，DeepSeek-V2显著提升了性能，降低了42.5%的训练成本，减少了93.3%的KV缓存，并将最大生成吞吐量提高了5.76倍。

我们这里主要实现实现DeepSeek的主要改进：多头隐性注意力、细粒度专家分割和共享的专家隔离

架构细节

DeepSeek-V2整合了两种创新架构，我们将详细讨论：

1. 用于前馈网络（FFNs）的DeepSeekMoE架构。
2. 用于注意力机制的多头隐性注意力（MLA）。

DeepSeekMoE

在标准的MoE架构中，每个令牌被分配给一个（或两个）专家，每个MoE层都有多个在结构上与标准前馈网络（FFN）相同的专家。这种设置带来了两个问题：指定给令牌的专家将试图在其参数中聚集不同类型的知识，但这些知识很难同时利用；其次，被分配给不同专家的令牌可能需要共同的知识，导致多个专家在各自的参数中趋向于收敛，获取共享知识。

为了应对这两个问题，DeepSeekMoE引入了两种策略来增强专家的专业化：

1. 细粒度专家分割：为了在每个专家中更有针对性地获取知识，通过切分FFN中的中间隐藏维度，将所有专家分割成更细的粒度。
2. 共享专家隔离：隔离某些专家作为始终被激活的共享专家，旨在捕获不同上下文中的共同知识，并通过将共同知识压缩到这些共享专家中，减少其他路由专家之间的冗余。

让我们来定义DeepSeekMoE中第t个令牌的专家分配。如果u_t是该令牌的FFN输入，其输出h`_t将会是：

其中和分别是共享专家和路由专家的数量；FFN()和FFN()分别表示-th共享专家和-th路由专家。

对于路由专家而言，g_i,t 是第i个路由专家的门控值，s_i,t 是令牌到专家的亲和分数，Topk(., Kr) 包含了Kr个最高的亲和分数，其中Kr是活跃的路由专家的数量。

有了以上的公式，我们就来使用代码实现

门控模型实现：

class MoEGate(torch.nn.Module):
def __init__(self, num_experts_per_tok: int, n_routed_experts: int, routed_scaling_factor: int, topk_method: str, n_group: int, topk_group: int, hidden_size: int):
super().__init__()
self.top_k = num_experts_per_tok
self.n_routed_experts = n_routed_experts
self.routed_scaling_factor = routed_scaling_factor
self.topk_method = topk_method
self.n_group = n_group
self.topk_group = topk_group
self.weight = torch.nn.Parameter(torch.empty((self.n_routed_experts, hidden_size)))
torch.nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
def forward(self, x: torch.Tensor):
batch, seq_len, h = x.shape
hidden_states = x.view(-1, h)
logits = torch.nn.functional.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32), None)
scores = logits.softmax(dim=-1, dtype=torch.float32)
if self.topk_method == "greedy":
topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)
elif self.topk_method == "group_limited_greedy":
group_scores = (scores.view(batch * seq_len, self.n_group, -1).max(dim=-1).values)
group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1] # [n, top_k_group]
group_mask = torch.zeros_like(group_scores) # [n, n_group]
group_mask.scatter_(1, group_idx, 1) # [n, n_group]
score_mask = (
group_mask.unsqueeze(-1)
.expand(
batch * seq_len, self.n_group, self.n_routed_experts // self.n_group
)
.reshape(batch * seq_len, -1)
) # [n, e]
tmp_scores = scores.masked_fill(~score_mask.bool(), 0.0) # [n, e]
topk_weight, topk_idx = torch.topk(
tmp_scores, k=self.top_k, dim=-1, sorted=False
)
return topk_idx, topk_weight

MoE

class MoE(torch.nn.Module):
def __init__(self, dim: int, routed_scaling_factor: int, topk_method: str, n_group: int, topk_group: int, hidden_dim: int | None = None, n_routed_experts: int = 12, num_experts_per_tok: int = 4, n_shared_experts: int = 2, mlp: str = "swiglu"):
super().__init__()
self.experts_per_rank = n_routed_experts
self.num_experts_per_tok = num_experts_per_tok
self.n_shared_experts = n_shared_experts
mlp_block = SwiGLU
self.experts = torch.nn.ModuleList([mlp_block(dim, hidden_dim) for i in range(n_routed_experts)])
self.gate = MoEGate(num_experts_per_tok, n_routed_experts, routed_scaling_factor, topk_method, n_group, topk_group, dim)
self.shared_experts = mlp_block(dim, hidden_dim * n_shared_experts)
def forward(self, x: torch.Tensor):
identity = x
orig_shape = x.shape
topk_idx, topk_weight = self.gate(x)
x = x.view(-1, x.shape[-1])
flat_topk_idx = topk_idx.view(-1)
x = x.repeat_interleave(self.num_experts_per_tok, dim=0)
y = torch.empty_like(x)
y = y.type(x.dtype)
for i, expert in enumerate(self.experts):
y[flat_topk_idx == i] = expert(x[flat_topk_idx == i]).to(dtype=x.dtype)
y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
y = y.view(*orig_shape)
output = y + self.shared_experts(identity)
return output

多头隐性注意力（MLA）

多头隐性注意力（MLA）相较于标准的多头注意力（MHA）实现了更优的性能，并且显著减少了KV缓存，提高了推理效率。与多查询注意力（MQA）和分组查询注意力（GQA）中减少KV头的方法不同，MLA将键（Key）和值（Value）共同压缩成一个潜在向量。

MLA不是缓存键（Key）和值（Value）矩阵，而是将它们联合压缩成一个低秩向量，这使得缓存的项目数量更少，因为压缩维度远小于多头注意力（MHA）中输出投影矩阵的维度。

标准的RoPE（旋转位置嵌入）与上述的低秩KV压缩不兼容。解耦RoPE策略使用额外的多头查询q_t和共享键k_t来实现RoPE。

下面总结了完整的MLA计算过程：

MLA实现

class MLA(torch.nn.Module):
def __init__(self, model_args: DeepseekConfig):
super().__init__()
d_model = model_args.d_model
self.num_heads = model_args.num_heads
self.head_dim = model_args.d_model // model_args.num_heads
self.attn_dropout = torch.nn.Dropout(model_args.dropout)
self.res_dropout = torch.nn.Dropout(model_args.dropout)
self.flash_attn = hasattr(torch.nn.functional, "scaled_dot_product_attention")
self.q_lora_rank = model_args.q_lora_rank
self.qk_rope_head_dim = model_args.qk_rope_head_dim
self.kv_lora_rank = model_args.kv_lora_rank
self.v_head_dim = model_args.v_head_dim
self.qk_nope_head_dim = model_args.qk_nope_head_dim
self.q_head_dim = model_args.qk_nope_head_dim + model_args.qk_rope_head_dim
self.q_a_proj = torch.nn.Linear(d_model, model_args.q_lora_rank, bias=False)
self.q_a_layernorm = RMSNorm(model_args.q_lora_rank)
self.q_b_proj = torch.nn.Linear(model_args.q_lora_rank, self.num_heads * self.q_head_dim, bias=False)
self.kv_a_proj_with_mqa = torch.nn.Linear(d_model,model_args.kv_lora_rank + model_args.qk_rope_head_dim,bias=False,)
self.kv_a_layernorm = RMSNorm(model_args.kv_lora_rank)
self.kv_b_proj = torch.nn.Linear(model_args.kv_lora_rank,self.num_heads * (self.q_head_dim - self.qk_rope_head_dim +
self.v_head_dim),bias=False,)
self.o_proj = torch.nn.Linear(self.num_heads * self.v_head_dim,d_model, bias=False,)
def forward(self, x: torch.Tensor, mask: torch.Tensor, freqs_cis) -> torch.Tensor:
batch, seq_len, d_model = x.shape
q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(x)))
q = q.view(batch, seq_len, self.num_heads, self.q_head_dim).transpose(1, 2)
q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
compressed_kv = self.kv_a_proj_with_mqa(x)
compressed_kv, k_pe = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
k_pe = k_pe.view(batch, seq_len, 1, self.qk_rope_head_dim).transpose(1, 2)
kv = (self.kv_b_proj(self.kv_a_layernorm(compressed_kv))
.view(batch, seq_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
.transpose(1, 2))
k_nope, value_states = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
q_pe, k_pe = apply_rope(q_pe, k_pe, freqs_cis)
k_pe = k_pe.transpose(2, 1)
q_pe = q_pe.transpose(2, 1)
query_states = k_pe.new_empty(batch, self.num_heads, seq_len, self.q_head_dim)
query_states[:, :, :, : self.qk_nope_head_dim] = q_nope
query_states[:, :, :, self.qk_nope_head_dim :] = q_pe
key_states = k_pe.new_empty(batch, self.num_heads, seq_len, self.q_head_dim)
key_states[:, :, :, : self.qk_nope_head_dim] = k_nope
key_states[:, :, :, self.qk_nope_head_dim :] = k_pe
attn_mtx = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
attn_mtx = attn_mtx + mask[:, :, :seq_len, :seq_len]
attn_mtx = torch.nn.functional.softmax(attn_mtx.float(), dim=-1).type_as(key_states)
attn_mtx = self.attn_dropout(attn_mtx)
output = torch.matmul(attn_mtx, value_states) # (batch, n_head, seq_len, head_dim)
output = output.transpose(1, 2).contiguous().view(batch, seq_len, self.num_heads * self.v_head_dim)
output = self.o_proj(output)
output = self.res_dropout(output)
return output

总结

本文详细介绍了DeepSeek-V2语言模型，这是一个强大的开源混合专家（MoE）语言模型，采用创新的架构来提高训练和推理的经济性和效率。DeepSeek-V2采用了两种核心技术：细粒度专家分割和共享专家隔离，这两种策略显著提高了专家的专业化水平。此外，文章还介绍了多头隐性注意力（MLA），这是一种改进的注意力机制，通过低秩键值联合压缩和解耦旋转位置嵌入，优化了模型的存储和计算效率。

除了理论探讨，我们通过编写代码实现DeepSeek-V2，可以更深入地理解其架构和工作原理。可以帮助你账务如何实现先进的混合专家（MoE）模型，还能深化对多头隐性注意力（MLA）和低秩键值压缩等关键技术的理解。通过实践，读者将能够验证理论的有效性，并对模型的性能和效率有直观的认识。

https://avoid.overfit.cn/post/317a967c8dac42ee98f96d8390851476

作者：Zain ul Abideen