大模型学习 （Datawhale_Happy-LLM）笔记5: 搭建一个 Transformer

大模型学习 （Datawhale_Happy-LLM）笔记5: 搭建一个 Transformer

搭建 Transformer 的核心组件总结

1. 基础功能模块

• 自注意力机制：通过 QKV 矩阵计算序列内依赖关系，公式为：
  $$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$$
• 多头注意力：将特征拆分为多个子空间并行计算注意力，增强模型表达能力
• 位置编码：通过正弦余弦函数为序列添加位置信息，解决 Transformer 无序列感知问题

2. 核心网络层

• 层归一化 (LayerNorm)：对每个样本的所有特征维度归一化，公式为：
  $$\text{LayerNorm}(x)=\alpha \odot \frac{x-\mu }{\sqrt{{\sigma }^2+ϵ}}+\beta$$
• 前馈神经网络 (MLP)：由 Linear+ReLU+Linear 构成，增强模型非线性表达能力

3. 编码器与解码器架构

• 编码器层：由 LayerNorm+MultiHeadAttention+残差连接 和 LayerNorm+MLP+残差连接 组成
• 解码器层：比编码器多一个 掩码多头注意力，用于避免预测时看到未来信息
• 堆叠结构：多个编码器/解码器层堆叠形成深度网络

4. 输入输出处理

• 嵌入层：将离散token转换为连续向量，与位置编码相加后输入网络
• 输出层：通过线性层将特征映射到词表空间，用于生成概率分布

5. 关键技术点

• 残差连接：解决深度网络训练梯度消失问题
• Dropout：随机丢弃神经元，防止过拟合
• 掩码机制：在解码器中屏蔽未来位置，确保预测时的因果关系

完整 Transformer 架构流程图

输入序列 ──→ 嵌入层 + 位置编码 ──→ 编码器序列 ──→ 编码器输出
                                        │
                                        ↓
                                  解码器序列（含掩码）
                                        │
                                        ↓
                                  线性层 + Softmax ─→ 输出序列

多头自注意力模块

class ModelArgs:
    def __init__(self):
        self.n_embed = 256
        self.n_head = 4
        self.head_dim = self.n_embed // self.n_heads  # 显式定义head维度
        self.dropout = 0.1
        self.max_seq_len = 512
        self.n_layers = 6
        self.vocab_size = None  # 添加词表大小
        self.block_size = None  # 添加最大序列长度

class MultiHeadAttention(nn.Module):
    
    def __init__(self, args: ModelArgs, is_causal=False):
        # 构造函数
        # args: 配置对象
        super().__init__()
        # 隐藏层维度必须是头数的整数倍，因为后面我们会将输入拆成头数个矩阵
        assert args.n_embed % args.n_head == 0
        # 模型并行处理大小，默认为 1
        model_parallel_size = 1
        # 本地计算头数， 等于总头数除以模型并行处理大小
        self.n_local_heads = args.n_heads // model_parallel_size
        # 每个头的维度，等于模型维度除以头的总数
        self.head_dim = args.dim // args.n_heads

        # Wq, Wk, Wv 参数矩阵，每个参数矩阵为 n_embd x n_embd
        # 这⾥通过三个组合矩阵来代替了n个参数矩阵的组合，其逻辑在于矩阵内积再拼接其实等同于拼接矩阵再内积，
        self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
        # 输出权重矩阵，维度为 n_embd x n_embd（head_dim = n_embeds / n_heads）
        self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
        # 注意力的dropout
        self.attn_dropout = nn.Dropout(args.dropout)
        # 残差连接的 dropout
        self.resid_dropout = nn.Dropout(args.dropout)
        # 创建⼀个上三⻆矩阵，⽤于遮蔽未来信息
        # 注意，因为是多头注意⼒，Mask 矩阵⽐之前我们定义的多⼀个维度
        if is_causal:
            mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))
            mask = torch.triu(mask, diagonal=1)
            # 注册为模型的缓冲区
            self.register_buffer("mask", mask)
            
    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
        # 获取批次⼤⼩和序列⻓度，[batch_size, seq_len, dim]
        bsz, seqlen, _ = q.shape
        # 计算查询（Q）、键（K）、值（V）,输⼊通过参数矩阵层，维度为 (B, T, n_embed) x (n_embed, n_embed) -> (B, T, n_embed)
        xq, xk, xv = self.wq(q), self.wk(k), self.wv(v)
        # 将 Q、K、V 拆分成多头，维度为 (B, T, n_head, C // n_head)，然后交换维度，变成 (B, n_head, T, C // n_head)
        # 因为在注意⼒计算中我们是取了后两个维度参与计算
        # 为什么要先按B*T*n_head*C//n_head展开再互换1、2维度⽽不是直接按注意⼒输⼊展开，是因为view的展开⽅式是直接把输⼊全部排开，
        # 然后按要求构造，可以发现只有上述操作能够实现我们将每个头对应部分取出来的⽬标
        xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
        xk = xk.view(bsz, seqlen, self.n_local_heads, self.head_dim)
        xv = xv.view(bsz, seqlen, self.n_local_heads, self.head_dim)
        xq = xq.transpose(1, 2)
        xk = xk.transpose(1, 2)
        xv = xv.transpose(1, 2)
        # 注意⼒计算
        # 计算 QK^T / sqrt(d_k)，维度为 (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
        scores = torch.matmul(xq, xk.transpose(2, 3)) / math.sqrt(self.head_dim)
        # 掩码⾃注意⼒必须有注意⼒掩码
        if self.is_causal:
            assert hasattr(self, 'mask')
            # 这⾥截取到序列⻓度，因为有些序列可能⽐ max_seq_len 短
            scores = scores + self.mask[:, :, :seqlen, :seqlen]
        # 计算 softmax，维度为 (B, nh, T, T)
        scores = F.softmax(scores.float(), dim=-1).type_as(xq)
        # 做 Dropout
        scores = self.attn_dropout(scores)
        # V * Score，维度为(B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
        output = torch.matmul(scores, xv)
        # 恢复时间维度并合并头。
        # 将多头的结果拼接起来, 先交换维度为 (B, T, n_head, C // n_head)，再拼接成 (B, T, n_head * C // n_head)
        # contiguous 函数⽤于重新开辟⼀块新内存存储，因为Pytorch设置先transpose再view会报错，
        # 因为view直接基于底层存储得到，然⽽transpose并不会改变底层存储，因此需要额外存储
        output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
        # 最终投影回残差流。
        output = self.wo(output)
        output = self.resid_dropout(output)
        return output

前馈神经网络

class MLP(nn.Module):
    '''前馈神经网络
    MLP, (Multi-Layer Perceptron) 多层感知机
    用以构建前馈神经网络
    '''
    def __init__(self, dim:int, hidden_dim: int, dropout: float):
        super().__init__()
        # 定义第一层线性变换，从输入维度到隐藏维度
        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
        # 定义第二层线性变换，从隐藏维度到输入维度
        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
        # 定义 dropout 层，用于防止过拟合
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # 前向传播函数
        # 首先， 输入x通过第一层线性变换和 RELU 激活函数
        # 然后， 结果通过第二层线性变换
        # 最后， 通过 dropout 层
        return self.dropout(self.w2(F.relu(self.w1(x))))

LayerNorm 层

归一化操作的数学公式

1. 计算样本均值 (其中，$Z_j^i$ 是样本i在第j个维度上的值，m就是mini-batch的大小)
   ${\mu }_j=\frac{1}{m}\sum _{i=1}^m{Z}_j^i$
2. 再计算样本的方差
   ${\sigma }^2=\frac{1}{m}\sum _{i=1}m(Z_j^i-{\mu }_j{)}^2$
3. 最后对每个样本的值减去均值再除以标准差来将这个mini-batch的样本分布转化为标准正态分布
   ${\tilde{Z}}_j=\frac{Z_j-{\mu }_j}{\sqrt{{\sigma }^2+ϵ}}$
   
   (此处$ϵ$这一极小量是为了避免分母为零)

class LayerNorm(nn.Module):
    '''基于上述归一化的式子，实现一个简单的 Layer Norm 层'''
    def __init__(self, features, eps = 1e-6):
        super(LayerNorm, self).__init__()
        # 线性矩阵做映射
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self, x):
        # 在统计每个样本所有维度的值，求均值和方差
        mean = x.mean(-1, keepdim=True) # mean: [batch_size, max_len, 1]
        # std: [bsz, max_len, 1]   
        # 这里等价于先求方差，再开根号是标准差, var = x.var(-1, keepdim=True) std = torch.sqrt(var + self.eps)  
        std = x.std(-1, keepdim=True) 
        # 注意这里在最后一个维度发生了广播
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

Encoder

class EncoderLayer(nn.Module):
    '''Encoder 层'''
    def __init__(self, args):
        super().__init__()
        # 一个layer中有两个 layer norm 分别在 attention之前和 MLP 之前
        self.attention_norm = LayerNorm(args.n_embed)
        # Encoder 不需要掩码， 传入 is_causal=False
        self.attention = MultiHeadAttention(args, is_causal=False)
        self.fnn_norm = LayerNorm(args.n_embed)
        self.feed_forward = MLP(dim=args.n_embed, hidden_dim=args.n_embed*4, dropout=args.dropout)

    def forward(self, x):
        # Layer Norm
        x = self.attention_norm(x)
        # 自注意力
        h = x + self.attention.forward(x,x,x)
        # 经过前馈神经网络
        out = h + self.feed_forward.forward(self.fnn_norm(h))
        return out

# 然后我们搭建一个 Encoder
class Encoder(nn.Module):
    '''Encoder 块'''
    def __init__(self, args):
        super(Encoder, self).__init__()
        # 一个 Encoder 由 N 个 Encoder Layer 组成
        self.layers = nn.ModuleList([EncoderLayer(args) for _ in range(args.n_layer)])
        self.norm = LayerNorm(args.n_embed)

    def forward(self, x):
        # 分别通过 N 层，Encoder Layer
        for layer in self.layers:
            x = layer(x)
        return self.norm(x)
        

Decoder

class DecoderLayer(nn.Module):
    '''解码层'''
    def __init__(self, args):
        super().__init__()
       # 一个 Layer 中有三个layernorm， 分别在 Mask Attention 之前、Self Attention 之前和 MLP 前
        self.attention_norm_1 = LayerNorm(args.n_embed)
        # Decoder 的第一个部分是 Mask Attention, 传入  is_causal =True
        self.mask_attention = MultiHeadAttention(args, is_causal=True)
        self.attention_norm2 = LayerNorm(args.n_embed)
        # Decoder 的第二个部分是类似于 Encoder 的 Attention, 传入 is_causal=False
        self.attention = MultiHeadAttention(args, is_causal=False)
        self.ffn_norm = LayerNorm(args.n_embed)
        # 第三个部分是 MLP
        self.feed_forward = MLP(args)

    def forward(self, x, enc_out):
        # Layer Norm
        x = self.attention_norm_1(x)
        # 掩码自注意力
        x = x + self.mask_attention.forward(x,x,x)
        # 多头注意力
        x = self.attention_norm_2(x)
        h = x + self.attention.forward(x, enc_out, enc_out)
        # 经过前馈神经网络
        out = h + self.feed_forward.forward(self.ffn_norm(h))
        return out

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_embed)

    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)

搭建一个 Transformer

# 位置编码层
class PositionalEncoding(nn.Module):
    def __init__(self, args):
        super(PositionalEncoding, self).__init__()
        # Dropout 层
        self.dropout = nn.Dropout(p=args.dropout)

        # block size 是序列最大长度
        pe = torch.zeros(args.block_size, args.n_embed)
        position = torch.arange(0, args.block_size).unsqueeze(1)
        # 计算 theta
        div_term = torch.exp(
            torch.arange(0, args.n_embed,2) * -(math.log(10000.0) / args.n_embed)
        )
        # 分别计算 sin, cos 结果
        pe[:,0::2] = torch.sin(position * div_term)
        pe[:,1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer("pe", pe)

    def forward(self, x):
        # 将位置编码加到 Embedding 结果上
        x = x + self.pe[:, : x.size(1)].requires_grad_(False)
        return self.dropout(x)

一个完整的 Transformer

class Transformer(nn.Module):
    """整体模型"""
    def __init__(self, args):
        super().__init__()
        # 必须输入词表大小和 block size
        assert args.vocab_size is not None
        assert args.block_size is not None
        self.args = args
        self.transformer = nn.ModuleDict(dict(
            wte = nn.Embedding(args.vocab_size, args.n_embed),
            wpe = PositionalEncoding(args),
            drop = nn.Dropout(args.dropout),
            encoder = Encoder(args),
            decoder = Decoder(args),
        ))
        # 最后的线性层，输入是 n_embed, 输出是词表大小
        self.lm_head = nn.Linear(args.n_embed, args.vocab_size, bias=False)

        # 初始化所有的权重
        self.apply(self._init_weights)

        # 查看所有参数数量
        print(f'number of parameters: {self.get_num_params()/1e6:.2fM}')
    
    def get_num_params(self, non_embedding=False):
        """统计所有参数的数量"""
        # non_embedding： 是否统计 embedding 的参数
        n_params = sum(p.numel() for p in self.parameters())
        # 如果不统计 embedding 的参数就减去
        if non_embedding:
            n_params -= self.transformer.wpe.weight.numel()
        return n_params

    def _init_weights(self, module):
        """初始化权重"""
        # 线性层和 embedding 层初始化为正则分布
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        # 输入为 idx,  维度为 (batch_size, sequence len); targets为目标序列，用于计算 Loss
        device = idx.device
        b, t = idx.size()
        assert t <= self.args.block_size, f'不能计算该序列，该序列长度为{t}, 最大序列长度只有 {self.args.block_size}'

        # 通过 self.transformer
        # 首先将输入 idx 通过 embedding 层， 得到维度为 (batch_size, sequence len, n_embed)
        print(f'idx: {idx.size()}')
        # 通过 embedding 层
        tok_emb = self.transformer.wte(idx)
        print(f"tok_emb: {tok_emb.size()}")
        # 然后通过位置编码
        pos_emb = self.transformer.wpe(tok_emb)
        # 然后再进行 dropout
        x = self.transformer.drop(pos_emb)
        # 然后通过 Encoder
        print(f'x after wpe: {x.size()}')

        if targets is not None:
            # 训练阶段，如果我们给了 targets，就计算 loss
            # 先通过最后的 Linear 层，得到维度为（batch_size, seq len, vocab size）
            logits = self.lm_head(x)
            # 再跟 targets 计算交叉熵
            loss = F.cross_entropy(
            	logits.view(-1, logits.size(-1)),
            	targets.view(-1),
            	ignore_index=-1)
        else:
            # 推理阶段，我们只需要 logits, loss 为 None
            # 取 -1 是只取序列中的最后一个作为输出
            logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim
            loss = None
        return logits, loss