手撕 Decoder
Happy-LLM:从零开始的大语言模型原理与实践教程.pdf P24
Decoder Layer
一个Decoder Layer内的数据流动顺序为:
Input (x)↓
LayerNorm 1↓
Masked Self-Attention↓
x + Attention(x)↓
LayerNorm 2↓
Cross-Attention (with enc_out)↓
h = x + Attention(x, enc_out)↓
LayerNorm 3↓
MLP (Feed-Forward Network)↓
out = h + MLP(h)
该实现依然与标准transformer不一样,以下代码里,先归一化再残差连接(Pre-LayerNorm),而标准transformer则相反
代码
class DecoderLayer(nn.Module):def __init__(self, args):super().__init__()self.attention_norm_1 = LayerNorm(args.n_embd)self.mask_attention = MultiHeadAttention(args, is_causal=True)self.attention_norm_2 = LayerNorm(args.n_embd)self.attention = MultiHeadAttention(args, is_causal=False)self.ffn_norm = LayerNorm(args.n_embd)self.feed_forward = MLP(args)def forward(self, x, enc_out):# Layer Normnorm_x = self.attention_norm_1(x)# 掩码自注意力x = x + self.mask_attention.forward(norm_x, norm_x, norm_x)# 多头注意力norm_x = self.attention_norm_2(x)h = x + self.attention.forward(norm_x, enc_out, enc_out)# 经过前馈神经网络out = h + self.feed_forward.forward(self.ffn_norm(h))return out
初始化了三个子层(Masked Multi-Head Attention, Cross Attention, Feed Forward),每个子层都含有一个Layer Norm,三个层归一化的函数相同,均为LayerNorm(args.n_embd),仅名称不同
代码逻辑实质上是三个重复的
搭建Decoder
class Decoder(nn.Module):'''解码器'''def __init__(self, args):super(Decoder, self).__init__() # ⼀个 Decoder 由 N 个 Decoder Layer 组成self.layers = nn.ModuleList([DecoderLayer(args) for _ in range(args.n_layer)])self.norm = LayerNorm(args.n_embd)def forward(self, x, enc_out):"Pass the input (and mask) through each layer in turn."for layer in self.layers:x = layer(x, enc_out)return self.norm(x)
[DecoderLayer(args) for _ in range(args.n_layer)] 通过循环生成 args.n_layer 个 DecoderLayer 实例;nn.ModuleList 将生成的 DecoderLayer 实例列表包装为 nn.ModuleList,动态创建并注册多个解码器层,构建符合 Transformer 架构的解码器
代码末尾的 self.norm(x) 是对所有层处理后的最终输出进行归一化
参考文章
Happy-LLM:从零开始的大语言模型原理与实践教程.pdf P24