1. 为什么“手撕Transformer”不是炫技而是每个想真正吃透模型的人必经的断骨重生我第一次在实验室服务器上跑通自己写的Transformer Encoder层时盯着终端里跳出来的loss: 4.217发了三分钟呆——不是因为成功了而是因为这串数字背后我亲手把那个被无数论文和教程神化的黑箱一寸寸拆开、校准、缝合最后让它开始呼吸。这不是为了证明“我能写”而是当我在调参时发现某个位置编码的sin/cos相位偏移0.001导致收敛变慢当我在调试Multi-Head Attention的QKV投影矩阵维度对不齐引发的广播错误时我才真正明白所有关于“注意力机制很强大”的描述都抵不过你亲手让一个softmax输出的梯度流过自己定义的线性层那一刻的触感。“手撕Transformer”这个说法在2024年已经不是新鲜事但绝大多数人停在了“抄代码—改参数—跑通demo”的循环里。热搜词里反复出现的“pytorch安装”“cuda版本适配”“位置编码代码”恰恰暴露了一个事实我们太习惯站在巨人的肩膀上俯视模型结构图却忘了巨人也是从一行import torch开始写起的。哈佛大学那篇著名的《The Annotated Transformer》之所以成为经典不是因为它画出了多漂亮的矩阵变换图而是它用Python代码把每一步张量形状变化都钉死在注释里——比如当你看到q k.transpose(-2, -1) / math.sqrt(d_k)这行时如果脑子里没有立刻浮现出(batch, head, seq_len, d_k) (batch, head, d_k, seq_len) → (batch, head, seq_len, seq_len)的形状推演那你就还没真正“看见”Self-Attention。这篇博文不提供“一键安装包”不承诺“三天速成”它只做一件事带你回到2017年那篇划时代论文刚发布时的状态——没有现成的nn.MultiheadAttention没有Hugging Face的AutoModel只有纸、笔、PyTorch文档和一颗愿意为每个维度错误花两小时debug的心。你会亲手实现Positional Encoding的sin/cos波形生成逻辑会逐行推导Scaled Dot-Product Attention中masking如何影响梯度回传会发现原来Feed-Forward Network里的两个线性层之间那个ReLU其实在原始论文里是GELU而PyTorch默认用的是ReLU这个细节连很多资深工程师都会忽略。这些不是“知识点”而是你未来在调试Swin Transformer的shifted window attention、或者修改Diffusion Transformer的time-step embedding时能瞬间定位问题的肌肉记忆。所以如果你正被“transformer能记住多少条k线”这类应用问题困扰或者纠结于“anaconda配置pytorch环境”时CUDA版本报错那么请先放下那些工具链焦虑——真正的底层掌控力永远来自对最简结构的绝对理解。接下来我们将从零开始用最朴素的PyTorch原语构建一个可训练、可调试、可逐层inspect的Transformer Encoder Block。所有代码均可直接运行所有张量形状变化都会配上手绘级推演所有“为什么这样设计”的答案都藏在2017年那篇论文第5.1节的字里行间。2. 从白板到第一行代码构建Encoder Block前必须厘清的四个反直觉前提在敲下第一个class EncoderLayer(nn.Module):之前我建议你先合上电脑拿出一张A4纸画出四个关键前提。这些前提在90%的教程里被当作“常识”跳过但它们恰恰是后续所有维度错误、梯度消失、mask失效的根源。2.1 前提一Attention不是“计算相似度”而是“动态构造权重矩阵”初学者最容易陷入的误区是把Self-Attention想象成“计算每个词和其它词的相似度”。这是致命的简化。真实情况是Attention机制的本质是在每个时间步根据当前输入动态生成一个seq_len × seq_len的权重矩阵然后用这个矩阵对Value向量做加权求和。这个权重矩阵不是预设的也不是学习出来的固定参数而是由Query和Key实时计算出的——这意味着它的每一行都对应着“当前词应该关注哪些上下文词”的概率分布。举个具体例子假设输入序列是[I, love, NLP]长度为3。当处理love这个词索引1时Attention机制会生成一个3×3的权重矩阵其中第1行对应love的数值决定了love对I、love自身、NLP的注意力分配比例。这个比例不是靠人工规则设定而是通过Q K^T再经过softmax得到的。因此当你看到q k.transpose(-2, -1)这行代码时脑子里必须立刻反应出这里生成的不是一个标量而是一个(batch, seq_len, seq_len)的矩阵它的每一行之和必须为1softmax保证且该矩阵将被用于对V进行加权。提示很多初学者在实现时忘记对q k.transpose(-2, -1)的结果除以sqrt(d_k)导致softmax的输入值过大梯度趋近于零。这不是数学错误而是对“为什么要缩放”的物理意义理解不足——缩放是为了防止点积结果过大使softmax函数进入梯度极小的饱和区。2.2 前提二Multi-Head不是“并行跑多个Attention”而是“用不同子空间观察同一序列”“Multi-Head Attention”常被解释为“同时运行8个独立的Attention”这又是一个危险的简化。真实情况是Multi-Head的本质是将原始的高维特征空间线性投影到h个低维子空间中让每个子空间独立学习一种特定的依赖模式最后再把它们拼接起来恢复到原始维度。这就像用8个不同焦距的镜头同时拍摄同一场景每个镜头捕捉到的细节长距离依赖、短距离语法、实体关系等不同最终合成一张信息更丰富的图像。关键推演在于维度变化假设原始embedding维度是d_model512head数h8那么每个head的子空间维度d_k d_v d_model // h 64。此时Q、K、V的投影矩阵不再是(d_model, d_model)而是(d_model, d_k)、(d_model, d_k)、(d_model, d_v)。当你执行q self.w_q(x)时x的shape是(batch, seq_len, d_model)w_q的shape是(d_model, d_k)结果q的shape就是(batch, seq_len, d_k)。但注意这只是单个head的输出。为了支持h个head并行计算我们需要把q reshape为(batch, seq_len, h, d_k)再transpose为(batch, h, seq_len, d_k)——这个transpose操作才是Multi-Head能并行计算的物理基础。如果你漏掉这一步后续的q k.transpose(-2, -1)就会因维度不匹配而报错。2.3 前提三Positional Encoding不是“加个向量”而是“注入不可学习的周期性先验”几乎所有教程都会告诉你“把位置编码加到词向量上”。但没人告诉你这个“加法”操作是整个Transformer架构中最精妙的设计之一它用纯数学函数sin/cos构造出一个完全不可学习、但具有强归纳偏置的位置表示其周期性保证了模型能泛化到远超训练长度的序列。sin/cos函数的选择绝非偶然。原始论文中定义的PE(pos, 2i) sin(pos / 10000^(2i/d_model))PE(pos, 2i1) cos(pos / 10000^(2i/d_model))其中i是维度索引。这个公式确保了1不同维度对应不同波长的正弦波2波长形成一个从2π到10000·2π的几何级数3对于任意固定的偏移量kPE(posk)都可以被PE(pos)的线性组合所表示——这正是模型能学习到相对位置关系的数学基础。当你用torch.sin(position * div_term)生成编码时div_term的计算必须严格遵循1 / (10000 ** (2*i / d_model))任何近似如用1e4代替10000都会破坏这种优雅的周期性结构。注意位置编码的shape必须与词嵌入完全一致(1, seq_len, d_model)。这里的1代表batch维度的广播意味着同一个位置编码会被加到batch中所有样本的相同位置上。如果你不小心生成了(seq_len, d_model)在加法时会触发PyTorch的隐式广播导致难以察觉的bug。2.4 前提四Encoder-Decoder架构中的Masking本质是“在计算图中物理切断梯度流”在Encoder-Decoder结构中Decoder的Self-Attention需要使用causal mask也叫look-ahead mask确保每个位置只能看到它之前的token。很多人以为这只是“在softmax前把未来位置的logits设为-inf”但更深层的意义在于这个mask操作是在计算图中物理性地切断了未来token对当前token的梯度依赖路径。这意味着在反向传播时未来token的梯度不会回传到当前token的参数上从而强制模型学习单向的因果关系。实现时一个常见的坑是mask的shape。假设batch_size2seq_len5那么causal mask应该是(1, 1, 5, 5)其中1用于广播到batch和head维度。如果你生成了(5, 5)的mask在attn_scores attn_scores.masked_fill(mask 0, float(-inf))时PyTorch会尝试广播但一旦你的attn_scores是(2, 8, 5, 5)batch, head, seq, seq广播规则就可能出错。正确的做法是mask torch.tril(torch.ones(seq_len, seq_len)).view(1, 1, seq_len, seq_len)然后在计算前确保attn_scores的shape与mask兼容。这四个前提不是理论铺垫而是你后续每一行代码的“宪法”。当你在调试时遇到RuntimeError: The size of tensor a (64) must match the size of tensor b (512) at non-singleton dimension 2你应该立刻想到是不是在Multi-Head的reshape环节忘了把(batch, seq_len, d_model)转成(batch, seq_len, h, d_k)再transpose当你发现attention weights全是0.25序列长度为4时你应该检查是不是忘了除以sqrt(d_k)导致softmax饱和这些debug直觉全部建立在这四个反直觉前提的肌肉记忆之上。3. 逐行解剖从Positional Encoding到Multi-Head Attention的完整实现与深度验证现在让我们把白板上的推演变成可运行、可调试、可逐层inspect的代码。我会像带新人一样写出每一行并解释它存在的唯一理由。所有代码均基于PyTorch 2.0无需额外依赖可直接粘贴到.py文件中运行。3.1 第一块基石PositionalEncoding类——用数学函数编织时空坐标import torch import torch.nn as nn import math class PositionalEncoding(nn.Module): def __init__(self, d_model: int, dropout: float 0.1, max_len: int 5000): super().__init__() self.dropout nn.Dropout(pdropout) # 创建一个足够长的位置编码矩阵 (max_len, d_model) pe torch.zeros(max_len, d_model) # 创建位置索引向量 (max_len, 1) position torch.arange(0, max_len, dtypetorch.float).unsqueeze(1) # 计算div_term: (d_model//2,)对应sin/cos的分母 # 注意这里用的是 10000 ** (2*i / d_model)i从0到d_model//2-1 div_term torch.exp( torch.arange(0, d_model, 2, dtypetorch.float) * (-math.log(10000.0) / d_model) ) # 将sin应用到偶数维度 (0,2,4,...) pe[:, 0::2] torch.sin(position * div_term) # 将cos应用到奇数维度 (1,3,5,...) pe[:, 1::2] torch.cos(position * div_term) # 增加batch维度变为 (1, max_len, d_model)便于广播 pe pe.unsqueeze(0) # 将pe注册为buffer而非parameter因为它不需要学习 self.register_buffer(pe, pe) def forward(self, x: torch.Tensor) - torch.Tensor: x: (batch_size, seq_len, d_model) 返回: (batch_size, seq_len, d_model)位置编码已相加 # 确保x的seq_len不超过pe的最大长度 x x self.pe[:, :x.size(1)] return self.dropout(x)这段代码看似简单但藏着三个必须掌握的细节register_buffervsnn.Parameter位置编码是固定的、不可学习的所以必须用register_buffer。如果误用nn.Parameter(pe)PyTorch会把这个矩阵加入到model.parameters()中导致优化器试图更新它这不仅浪费计算还会破坏精心设计的周期性结构。你可以用list(model.named_buffers())来验证pe是否在buffer列表中。div_term的精确计算torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))这行是核心。它等价于1 / (10000 ** (2*i / d_model))但用exp(log())的方式计算是为了数值稳定性。如果你直接写1 / (10000 ** (2*i / d_model))当d_model很大时指数运算可能导致下溢或上溢。x self.pe[:, :x.size(1)]的广播机制self.pe的shape是(1, max_len, d_model)x是(batch, seq_len, d_model)。PyTorch的广播规则是从右向左对齐维度1可以广播到任意大小。所以self.pe[:, :x.size(1)]切片后是(1, seq_len, d_model)与x相加时1维度自动广播到batch大小。这是PyTorch高效实现的关键但也是初学者最容易困惑的地方。实操验证在你的脚本中添加以下测试代码# 测试PositionalEncoding pe PositionalEncoding(d_model12, dropout0.0) x torch.zeros(2, 5, 12) # batch2, seq_len5, d_model12 out pe(x) print(fInput shape: {x.shape}) print(fOutput shape: {out.shape}) print(fFirst tokens PE (first 4 dims): {out[0, 0, :4]}) print(fSecond tokens PE (first 4 dims): {out[0, 1, :4]})运行后你会看到out[0, 0, :4]和out[0, 1, :4]的值完全不同且呈现明显的正弦/余弦波形特征。这就是时空坐标的注入过程。3.2 第二块基石ScaledDotProductAttention——Attention机制的原子核def scaled_dot_product_attention( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor None, dropout_p: float 0.0 ) - torch.Tensor: 核心Attention计算函数 q, k, v: (batch, head, seq_len, d_k/d_v) mask: (batch, 1, seq_len, seq_len) 或 (1, 1, seq_len, seq_len) 返回: (batch, head, seq_len, d_v) # Step 1: 计算点积得分 (batch, head, seq_len, seq_len) # q k.transpose(-2, -1) - (batch, head, seq_len, d_k) (batch, head, d_k, seq_len) attn_scores torch.matmul(q, k.transpose(-2, -1)) # Step 2: 缩放 (scaled) # d_k 是每个head的维度必须是整数 d_k q.size(-1) attn_scores attn_scores / math.sqrt(d_k) # Step 3: 应用mask如果提供 if mask is not None: # mask为0的位置设置为负无穷使softmax后概率为0 attn_scores attn_scores.masked_fill(mask 0, float(-inf)) # Step 4: softmax得到注意力权重 (batch, head, seq_len, seq_len) # 沿着最后一个维度seq_len做softmax即每个位置对所有位置的权重和为1 attn_weights torch.softmax(attn_scores, dim-1) # Step 5: Dropout可选用于正则化 if dropout_p 0.0: dropout nn.Dropout(pdropout_p) attn_weights dropout(attn_weights) # Step 6: 加权求和 (batch, head, seq_len, seq_len) (batch, head, seq_len, d_v) # - (batch, head, seq_len, d_v) output torch.matmul(attn_weights, v) return output, attn_weights这个函数是整个Transformer的“心脏”它的每一行都值得深究k.transpose(-2, -1)这是张量操作的精髓。-2和-1代表倒数第二和倒数第一维度无论k的shape是(2, 8, 10, 64)还是(1, 1, 5, 64)这个操作都能正确地将最后两个维度交换。这是编写鲁棒代码的关键。attn_scores.masked_fill(mask 0, float(-inf))这里mask 0生成一个布尔张量masked_fill会将所有True位置的值替换为-inf。float(-inf)是PyTorch中表示负无穷的标准方式它能确保torch.softmax在计算时这些位置的输出为0。切记不能用-1e9之类的近似值因为在某些GPU上-1e9可能不会被完全视为负无穷导致微小的泄漏。torch.softmax(attn_scores, dim-1)dim-1指定了softmax沿着最后一个维度即seq_len进行这保证了每个位置的注意力权重之和为1。如果你错误地用了dim1那就会让每个head的权重和为1这是完全错误的。深度验证写一个单元测试强制验证mask的效果# 测试mask功能 q_test torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]]) # (1,1,2,2) k_test torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]]) # (1,1,2,2) v_test torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]]) # (1,1,2,2) # 创建一个causal mask: [[1,0],[1,1]] mask_test torch.tensor([[[[1, 0], [1, 1]]]], dtypetorch.bool) # (1,1,2,2) out, weights scaled_dot_product_attention(q_test, k_test, v_test, maskmask_test) print(Attention weights with causal mask:) print(weights.squeeze().numpy()) # 预期输出: [[1.0, 0.0], [0.5, 0.5]]第二行显示位置1能看到位置0和1但第一行位置0只能看到自己3.3 第三块基石MultiHeadAttention——将原子核组装成引擎class MultiHeadAttention(nn.Module): def __init__(self, d_model: int, num_heads: int, dropout: float 0.0): super().__init__() assert d_model % num_heads 0, d_model must be divisible by num_heads self.d_model d_model self.num_heads num_heads self.d_k d_model // num_heads # 每个head的维度 self.d_v d_model // num_heads # 定义Q, K, V的线性投影层 # w_q, w_k, w_v 的shape都是 (d_model, d_k) 或 (d_model, d_v) self.w_q nn.Linear(d_model, d_model) # 输出d_model后续再split self.w_k nn.Linear(d_model, d_model) self.w_v nn.Linear(d_model, d_model) # 输出投影层将h个head的输出拼接后映射回d_model self.w_o nn.Linear(d_model, d_model) self.dropout nn.Dropout(dropout) def forward( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor None ) - torch.Tensor: query, key, value: (batch, seq_len, d_model) mask: (batch, 1, seq_len, seq_len) or (1, 1, seq_len, seq_len) 返回: (batch, seq_len, d_model) batch_size query.size(0) # Step 1: 线性投影 (batch, seq_len, d_model) - (batch, seq_len, d_model) q self.w_q(query) # (batch, seq_len, d_model) k self.w_k(key) # (batch, seq_len, d_model) v self.w_v(value) # (batch, seq_len, d_model) # Step 2: Reshape and transpose to (batch, num_heads, seq_len, d_k/d_v) # 这是Multi-Head的核心操作将d_model维度拆分为(num_heads, d_k) q q.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) k k.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) v v.view(batch_size, -1, self.num_heads, self.d_v).transpose(1, 2) # 现在 q, k, v 的shape都是 (batch, num_heads, seq_len, d_k/d_v) # Step 3: 调用核心Attention函数 # 注意mask需要适配新的维度 (batch, 1, seq_len, seq_len) # 如果mask是 (batch, seq_len, seq_len)需要unsqueeze(1)变成 (batch, 1, seq_len, seq_len) if mask is not None and len(mask.shape) 3: mask mask.unsqueeze(1) # (batch, 1, seq_len, seq_len) x, self.attn_weights scaled_dot_product_attention( q, k, v, maskmask, dropout_pself.dropout.p ) # Step 4: 把num_heads维度合并回d_model # x: (batch, num_heads, seq_len, d_v) - (batch, seq_len, num_heads, d_v) x x.transpose(1, 2).contiguous() # contiguous() 是必须的因为transpose后内存不连续view会报错 x x.view(batch_size, -1, self.d_model) # (batch, seq_len, d_model) # Step 5: 最终线性投影 return self.w_o(x)这个类的实现集中体现了前面所有前提的落地q.view(...).transpose(1, 2)这是Multi-Head的物理实现。view将(batch, seq_len, d_model)重塑为(batch, seq_len, num_heads, d_k)transpose(1, 2)将其变为(batch, num_heads, seq_len, d_k)为并行计算做好准备。transpose(1, 2)中的1和2分别指代seq_len和num_heads维度的索引。contiguous()的必要性这是PyTorch中一个极其重要但常被忽略的细节。transpose操作会改变张量的内存布局使其变得“不连续”。而view操作要求张量在内存中是连续的。如果不加contiguous()view会抛出RuntimeError: view size is not compatible with input tensors size and stride。这个错误在调试时非常隐蔽因为它只在特定的维度变化下才出现。mask.unsqueeze(1)的智能适配我们的scaled_dot_product_attention函数期望mask的shape是(batch, 1, seq_len, seq_len)但用户传入的mask可能是(batch, seq_len, seq_len)。这个unsqueeze(1)操作就是为了让接口更友好避免用户在每次调用时都手动增加维度。实操心得在实际项目中我习惯在forward函数末尾添加一句self.attn_weights self.attn_weights.detach()。因为attn_weights是一个巨大的中间变量shape(batch, num_heads, seq_len, seq_len)如果它一直保留在计算图中会极大增加显存占用。detach()将其从计算图中剥离只保留数值用于可视化或分析这对训练大型模型至关重要。4. 组装与验证构建完整Encoder Block并用真实数据流进行端到端压力测试有了PositionalEncoding和MultiHeadAttention这两块基石现在我们可以组装出完整的Encoder Layer。但请注意一个标准的Encoder Block不仅仅是Attention它还包括Add Norm和Feed-Forward NetworkFFN。我们将严格按照原始论文的结构实现并在每一步都插入验证点。4.1 EncoderLayerAttention AddNorm FFN AddNorm的完整流水线class FeedForward(nn.Module): Feed-Forward Network原始论文中为两层线性变换 ReLU def __init__(self, d_model: int, d_ff: int, dropout: float 0.1): super().__init__() self.linear1 nn.Linear(d_model, d_ff) self.dropout nn.Dropout(dropout) self.linear2 nn.Linear(d_ff, d_model) # 注意原始论文使用的是ReLU但后来的实现如BERT常用GELU # 这里我们严格遵循原始论文用ReLU self.activation nn.ReLU() def forward(self, x: torch.Tensor) - torch.Tensor: # x: (batch, seq_len, d_model) x self.linear1(x) # (batch, seq_len, d_ff) x self.activation(x) # (batch, seq_len, d_ff) x self.dropout(x) # (batch, seq_len, d_ff) x self.linear2(x) # (batch, seq_len, d_model) return x class EncoderLayer(nn.Module): def __init__( self, d_model: int, num_heads: int, d_ff: int, dropout: float 0.1 ): super().__init__() self.self_attn MultiHeadAttention(d_model, num_heads, dropout) self.feed_forward FeedForward(d_model, d_ff, dropout) # 两个LayerNorm层 self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout1 nn.Dropout(dropout) self.dropout2 nn.Dropout(dropout) def forward(self, x: torch.Tensor, src_mask: torch.Tensor None) - torch.Tensor: x: (batch, seq_len, d_model) src_mask: (batch, 1, seq_len, seq_len) 或 None # Sublayer 1: Self-Attention # 先Norm再Attention再Add残差连接 # 这是Pre-LNPre-LayerNorm结构原始论文用的是Post-LN # 但Pre-LN在训练初期更稳定是现代实现的常见选择 norm_x self.norm1(x) attn_output self.self_attn(norm_x, norm_x, norm_x, masksrc_mask) x x self.dropout1(attn_output) # Sublayer 2: Feed-Forward norm_x self.norm2(x) ff_output self.feed_forward(norm_x) x x self.dropout2(ff_output) return x class Encoder(nn.Module): 堆叠N个EncoderLayer def __init__( self, layer: EncoderLayer, N: int ): super().__init__() self.layers nn.ModuleList([copy.deepcopy(layer) for _ in range(N)]) self.norm nn.LayerNorm(layer.d_model) def forward(self, x: torch.Tensor, mask: torch.Tensor None) - torch.Tensor: for layer in self.layers: x layer(x, mask) return self.norm(x)这个EncoderLayer的实现严格遵循了原始论文的“Add Norm”范式但做了一个关键的现代改进Pre-LNPre-LayerNorm。原始论文Post-LN的顺序是x Attn(LayerNorm(x))而Pre-LN是x Attn(LayerNorm(x))。这个改动看似微小但实测下来Pre-LN能让模型在训练初期收敛得更快、更稳定尤其是在层数较多时。这也是为什么Hugging Face的BertLayer默认使用Pre-LN。4.2 端到端压力测试用真实数据流驱动暴露所有潜在Bug光有代码还不够我们必须用真实的数据流来“压测”它。下面是一个完整的、可运行的端到端测试脚本它会模拟一个真实的训练步骤从随机输入开始经过Embedding、PositionalEncoding、Encoder最后计算一个简单的loss。这个过程会暴露出90%的实现错误。import torch import torch.nn as nn import torch.optim as optim import copy # 1. 构建一个最小的Transformer Encoder模型 class SimpleTransformerEncoder(nn.Module): def __init__(self, vocab_size1000, d_model512, num_heads8, d_ff2048, N2, dropout0.1): super().__init__() self.embedding nn.Embedding(vocab_size, d_model) self.pos_encoding PositionalEncoding(d_model, dropout) encoder_layer EncoderLayer(d_model, num_heads, d_ff, dropout) self.encoder Encoder(encoder_layer, N) # 添加一个简单的分类头用于生成loss self.classifier nn.Linear(d_model, 2) # 二分类 def forward(self, src: torch.Tensor, src_mask: torch.Tensor None) - torch.Tensor: # src: (batch, seq_len) x self.embedding(src) * math.sqrt(self.embedding.embedding_dim) # 缩放 x self.pos_encoding(x) x self.encoder(x, src_mask) # 取第一个token的输出作为句子表示类似[CLS] cls_token x[:, 0, :] return self.classifier(cls_token) # 2. 创建测试数据 batch_size 2 seq_len 10 vocab_size 1000 model SimpleTransformerEncoder(vocab_sizevocab_size, d_model128, num_heads4, d_ff512, N2) # 生成随机输入和标签 src torch.randint(0, vocab_size, (batch_size, seq_len)) # (2, 10) labels torch.randint(0, 2, (batch_size,)) # (2,) # 3. 创建mask一个简单的padding mask假设所有token都是有效的 # 在真实场景中这里会是根据padding token生成的mask src_mask torch.ones(batch_size, 1, seq_len, seq_len) # (2, 1, 10, 10) # 4. 前向传播 output model(src, src_mask) print(fModel output shape: {output.shape}) # 应该是 (2, 2) # 5. 计算loss并反向传播 criterion nn.CrossEntropyLoss() loss criterion(output, labels) loss.backward() # 6. 检查梯度 print(\nGradient check:) for name, param in model.named_parameters(): if param.grad is not None: print(f{name}: grad_norm {param.grad.norm().item():.4f}) # 7. 关键验证检查Attention权重是否合理 # 我们可以访问最后一层EncoderLayer的self_attn.attn_weights # 由于我们封装了多层需要稍微绕一下 last_encoder_layer model.encoder.layers[-1] if hasattr(last_encoder_layer.self_attn, attn_weights): attn_weights last_encoder_layer.self_attn.attn_weights print(f\nLast layer Attention weights shape: {attn_weights.shape}) print(fMean attention weight: {attn_weights.mean().item():.4f}) print(fMax attention weight: {attn_weights.max().item():.4f}) print(fMin attention weight: {attn_weights.min().item():.4f}) # 检查是否满足概率分布每行和应为1 row_sums attn_weights.sum(dim-1) print(fRow sums (should be ~1.0): {row_sums[0, 0].tolist()})运行这个脚本你会得到一系列关键输出。**请务必
1. 为什么“手撕Transformer”不是炫技而是每个想真正吃透模型的人必经的断骨重生我第一次在实验室服务器上跑通自己写的Transformer Encoder层时盯着终端里跳出来的loss: 4.217发了三分钟呆——不是因为成功了而是因为这串数字背后我亲手把那个被无数论文和教程神化的黑箱一寸寸拆开、校准、缝合最后让它开始呼吸。这不是为了证明“我能写”而是当我在调参时发现某个位置编码的sin/cos相位偏移0.001导致收敛变慢当我在调试Multi-Head Attention的QKV投影矩阵维度对不齐引发的广播错误时我才真正明白所有关于“注意力机制很强大”的描述都抵不过你亲手让一个softmax输出的梯度流过自己定义的线性层那一刻的触感。“手撕Transformer”这个说法在2024年已经不是新鲜事但绝大多数人停在了“抄代码—改参数—跑通demo”的循环里。热搜词里反复出现的“pytorch安装”“cuda版本适配”“位置编码代码”恰恰暴露了一个事实我们太习惯站在巨人的肩膀上俯视模型结构图却忘了巨人也是从一行import torch开始写起的。哈佛大学那篇著名的《The Annotated Transformer》之所以成为经典不是因为它画出了多漂亮的矩阵变换图而是它用Python代码把每一步张量形状变化都钉死在注释里——比如当你看到q k.transpose(-2, -1) / math.sqrt(d_k)这行时如果脑子里没有立刻浮现出(batch, head, seq_len, d_k) (batch, head, d_k, seq_len) → (batch, head, seq_len, seq_len)的形状推演那你就还没真正“看见”Self-Attention。这篇博文不提供“一键安装包”不承诺“三天速成”它只做一件事带你回到2017年那篇划时代论文刚发布时的状态——没有现成的nn.MultiheadAttention没有Hugging Face的AutoModel只有纸、笔、PyTorch文档和一颗愿意为每个维度错误花两小时debug的心。你会亲手实现Positional Encoding的sin/cos波形生成逻辑会逐行推导Scaled Dot-Product Attention中masking如何影响梯度回传会发现原来Feed-Forward Network里的两个线性层之间那个ReLU其实在原始论文里是GELU而PyTorch默认用的是ReLU这个细节连很多资深工程师都会忽略。这些不是“知识点”而是你未来在调试Swin Transformer的shifted window attention、或者修改Diffusion Transformer的time-step embedding时能瞬间定位问题的肌肉记忆。所以如果你正被“transformer能记住多少条k线”这类应用问题困扰或者纠结于“anaconda配置pytorch环境”时CUDA版本报错那么请先放下那些工具链焦虑——真正的底层掌控力永远来自对最简结构的绝对理解。接下来我们将从零开始用最朴素的PyTorch原语构建一个可训练、可调试、可逐层inspect的Transformer Encoder Block。所有代码均可直接运行所有张量形状变化都会配上手绘级推演所有“为什么这样设计”的答案都藏在2017年那篇论文第5.1节的字里行间。2. 从白板到第一行代码构建Encoder Block前必须厘清的四个反直觉前提在敲下第一个class EncoderLayer(nn.Module):之前我建议你先合上电脑拿出一张A4纸画出四个关键前提。这些前提在90%的教程里被当作“常识”跳过但它们恰恰是后续所有维度错误、梯度消失、mask失效的根源。2.1 前提一Attention不是“计算相似度”而是“动态构造权重矩阵”初学者最容易陷入的误区是把Self-Attention想象成“计算每个词和其它词的相似度”。这是致命的简化。真实情况是Attention机制的本质是在每个时间步根据当前输入动态生成一个seq_len × seq_len的权重矩阵然后用这个矩阵对Value向量做加权求和。这个权重矩阵不是预设的也不是学习出来的固定参数而是由Query和Key实时计算出的——这意味着它的每一行都对应着“当前词应该关注哪些上下文词”的概率分布。举个具体例子假设输入序列是[I, love, NLP]长度为3。当处理love这个词索引1时Attention机制会生成一个3×3的权重矩阵其中第1行对应love的数值决定了love对I、love自身、NLP的注意力分配比例。这个比例不是靠人工规则设定而是通过Q K^T再经过softmax得到的。因此当你看到q k.transpose(-2, -1)这行代码时脑子里必须立刻反应出这里生成的不是一个标量而是一个(batch, seq_len, seq_len)的矩阵它的每一行之和必须为1softmax保证且该矩阵将被用于对V进行加权。提示很多初学者在实现时忘记对q k.transpose(-2, -1)的结果除以sqrt(d_k)导致softmax的输入值过大梯度趋近于零。这不是数学错误而是对“为什么要缩放”的物理意义理解不足——缩放是为了防止点积结果过大使softmax函数进入梯度极小的饱和区。2.2 前提二Multi-Head不是“并行跑多个Attention”而是“用不同子空间观察同一序列”“Multi-Head Attention”常被解释为“同时运行8个独立的Attention”这又是一个危险的简化。真实情况是Multi-Head的本质是将原始的高维特征空间线性投影到h个低维子空间中让每个子空间独立学习一种特定的依赖模式最后再把它们拼接起来恢复到原始维度。这就像用8个不同焦距的镜头同时拍摄同一场景每个镜头捕捉到的细节长距离依赖、短距离语法、实体关系等不同最终合成一张信息更丰富的图像。关键推演在于维度变化假设原始embedding维度是d_model512head数h8那么每个head的子空间维度d_k d_v d_model // h 64。此时Q、K、V的投影矩阵不再是(d_model, d_model)而是(d_model, d_k)、(d_model, d_k)、(d_model, d_v)。当你执行q self.w_q(x)时x的shape是(batch, seq_len, d_model)w_q的shape是(d_model, d_k)结果q的shape就是(batch, seq_len, d_k)。但注意这只是单个head的输出。为了支持h个head并行计算我们需要把q reshape为(batch, seq_len, h, d_k)再transpose为(batch, h, seq_len, d_k)——这个transpose操作才是Multi-Head能并行计算的物理基础。如果你漏掉这一步后续的q k.transpose(-2, -1)就会因维度不匹配而报错。2.3 前提三Positional Encoding不是“加个向量”而是“注入不可学习的周期性先验”几乎所有教程都会告诉你“把位置编码加到词向量上”。但没人告诉你这个“加法”操作是整个Transformer架构中最精妙的设计之一它用纯数学函数sin/cos构造出一个完全不可学习、但具有强归纳偏置的位置表示其周期性保证了模型能泛化到远超训练长度的序列。sin/cos函数的选择绝非偶然。原始论文中定义的PE(pos, 2i) sin(pos / 10000^(2i/d_model))PE(pos, 2i1) cos(pos / 10000^(2i/d_model))其中i是维度索引。这个公式确保了1不同维度对应不同波长的正弦波2波长形成一个从2π到10000·2π的几何级数3对于任意固定的偏移量kPE(posk)都可以被PE(pos)的线性组合所表示——这正是模型能学习到相对位置关系的数学基础。当你用torch.sin(position * div_term)生成编码时div_term的计算必须严格遵循1 / (10000 ** (2*i / d_model))任何近似如用1e4代替10000都会破坏这种优雅的周期性结构。注意位置编码的shape必须与词嵌入完全一致(1, seq_len, d_model)。这里的1代表batch维度的广播意味着同一个位置编码会被加到batch中所有样本的相同位置上。如果你不小心生成了(seq_len, d_model)在加法时会触发PyTorch的隐式广播导致难以察觉的bug。2.4 前提四Encoder-Decoder架构中的Masking本质是“在计算图中物理切断梯度流”在Encoder-Decoder结构中Decoder的Self-Attention需要使用causal mask也叫look-ahead mask确保每个位置只能看到它之前的token。很多人以为这只是“在softmax前把未来位置的logits设为-inf”但更深层的意义在于这个mask操作是在计算图中物理性地切断了未来token对当前token的梯度依赖路径。这意味着在反向传播时未来token的梯度不会回传到当前token的参数上从而强制模型学习单向的因果关系。实现时一个常见的坑是mask的shape。假设batch_size2seq_len5那么causal mask应该是(1, 1, 5, 5)其中1用于广播到batch和head维度。如果你生成了(5, 5)的mask在attn_scores attn_scores.masked_fill(mask 0, float(-inf))时PyTorch会尝试广播但一旦你的attn_scores是(2, 8, 5, 5)batch, head, seq, seq广播规则就可能出错。正确的做法是mask torch.tril(torch.ones(seq_len, seq_len)).view(1, 1, seq_len, seq_len)然后在计算前确保attn_scores的shape与mask兼容。这四个前提不是理论铺垫而是你后续每一行代码的“宪法”。当你在调试时遇到RuntimeError: The size of tensor a (64) must match the size of tensor b (512) at non-singleton dimension 2你应该立刻想到是不是在Multi-Head的reshape环节忘了把(batch, seq_len, d_model)转成(batch, seq_len, h, d_k)再transpose当你发现attention weights全是0.25序列长度为4时你应该检查是不是忘了除以sqrt(d_k)导致softmax饱和这些debug直觉全部建立在这四个反直觉前提的肌肉记忆之上。3. 逐行解剖从Positional Encoding到Multi-Head Attention的完整实现与深度验证现在让我们把白板上的推演变成可运行、可调试、可逐层inspect的代码。我会像带新人一样写出每一行并解释它存在的唯一理由。所有代码均基于PyTorch 2.0无需额外依赖可直接粘贴到.py文件中运行。3.1 第一块基石PositionalEncoding类——用数学函数编织时空坐标import torch import torch.nn as nn import math class PositionalEncoding(nn.Module): def __init__(self, d_model: int, dropout: float 0.1, max_len: int 5000): super().__init__() self.dropout nn.Dropout(pdropout) # 创建一个足够长的位置编码矩阵 (max_len, d_model) pe torch.zeros(max_len, d_model) # 创建位置索引向量 (max_len, 1) position torch.arange(0, max_len, dtypetorch.float).unsqueeze(1) # 计算div_term: (d_model//2,)对应sin/cos的分母 # 注意这里用的是 10000 ** (2*i / d_model)i从0到d_model//2-1 div_term torch.exp( torch.arange(0, d_model, 2, dtypetorch.float) * (-math.log(10000.0) / d_model) ) # 将sin应用到偶数维度 (0,2,4,...) pe[:, 0::2] torch.sin(position * div_term) # 将cos应用到奇数维度 (1,3,5,...) pe[:, 1::2] torch.cos(position * div_term) # 增加batch维度变为 (1, max_len, d_model)便于广播 pe pe.unsqueeze(0) # 将pe注册为buffer而非parameter因为它不需要学习 self.register_buffer(pe, pe) def forward(self, x: torch.Tensor) - torch.Tensor: x: (batch_size, seq_len, d_model) 返回: (batch_size, seq_len, d_model)位置编码已相加 # 确保x的seq_len不超过pe的最大长度 x x self.pe[:, :x.size(1)] return self.dropout(x)这段代码看似简单但藏着三个必须掌握的细节register_buffervsnn.Parameter位置编码是固定的、不可学习的所以必须用register_buffer。如果误用nn.Parameter(pe)PyTorch会把这个矩阵加入到model.parameters()中导致优化器试图更新它这不仅浪费计算还会破坏精心设计的周期性结构。你可以用list(model.named_buffers())来验证pe是否在buffer列表中。div_term的精确计算torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))这行是核心。它等价于1 / (10000 ** (2*i / d_model))但用exp(log())的方式计算是为了数值稳定性。如果你直接写1 / (10000 ** (2*i / d_model))当d_model很大时指数运算可能导致下溢或上溢。x self.pe[:, :x.size(1)]的广播机制self.pe的shape是(1, max_len, d_model)x是(batch, seq_len, d_model)。PyTorch的广播规则是从右向左对齐维度1可以广播到任意大小。所以self.pe[:, :x.size(1)]切片后是(1, seq_len, d_model)与x相加时1维度自动广播到batch大小。这是PyTorch高效实现的关键但也是初学者最容易困惑的地方。实操验证在你的脚本中添加以下测试代码# 测试PositionalEncoding pe PositionalEncoding(d_model12, dropout0.0) x torch.zeros(2, 5, 12) # batch2, seq_len5, d_model12 out pe(x) print(fInput shape: {x.shape}) print(fOutput shape: {out.shape}) print(fFirst tokens PE (first 4 dims): {out[0, 0, :4]}) print(fSecond tokens PE (first 4 dims): {out[0, 1, :4]})运行后你会看到out[0, 0, :4]和out[0, 1, :4]的值完全不同且呈现明显的正弦/余弦波形特征。这就是时空坐标的注入过程。3.2 第二块基石ScaledDotProductAttention——Attention机制的原子核def scaled_dot_product_attention( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor None, dropout_p: float 0.0 ) - torch.Tensor: 核心Attention计算函数 q, k, v: (batch, head, seq_len, d_k/d_v) mask: (batch, 1, seq_len, seq_len) 或 (1, 1, seq_len, seq_len) 返回: (batch, head, seq_len, d_v) # Step 1: 计算点积得分 (batch, head, seq_len, seq_len) # q k.transpose(-2, -1) - (batch, head, seq_len, d_k) (batch, head, d_k, seq_len) attn_scores torch.matmul(q, k.transpose(-2, -1)) # Step 2: 缩放 (scaled) # d_k 是每个head的维度必须是整数 d_k q.size(-1) attn_scores attn_scores / math.sqrt(d_k) # Step 3: 应用mask如果提供 if mask is not None: # mask为0的位置设置为负无穷使softmax后概率为0 attn_scores attn_scores.masked_fill(mask 0, float(-inf)) # Step 4: softmax得到注意力权重 (batch, head, seq_len, seq_len) # 沿着最后一个维度seq_len做softmax即每个位置对所有位置的权重和为1 attn_weights torch.softmax(attn_scores, dim-1) # Step 5: Dropout可选用于正则化 if dropout_p 0.0: dropout nn.Dropout(pdropout_p) attn_weights dropout(attn_weights) # Step 6: 加权求和 (batch, head, seq_len, seq_len) (batch, head, seq_len, d_v) # - (batch, head, seq_len, d_v) output torch.matmul(attn_weights, v) return output, attn_weights这个函数是整个Transformer的“心脏”它的每一行都值得深究k.transpose(-2, -1)这是张量操作的精髓。-2和-1代表倒数第二和倒数第一维度无论k的shape是(2, 8, 10, 64)还是(1, 1, 5, 64)这个操作都能正确地将最后两个维度交换。这是编写鲁棒代码的关键。attn_scores.masked_fill(mask 0, float(-inf))这里mask 0生成一个布尔张量masked_fill会将所有True位置的值替换为-inf。float(-inf)是PyTorch中表示负无穷的标准方式它能确保torch.softmax在计算时这些位置的输出为0。切记不能用-1e9之类的近似值因为在某些GPU上-1e9可能不会被完全视为负无穷导致微小的泄漏。torch.softmax(attn_scores, dim-1)dim-1指定了softmax沿着最后一个维度即seq_len进行这保证了每个位置的注意力权重之和为1。如果你错误地用了dim1那就会让每个head的权重和为1这是完全错误的。深度验证写一个单元测试强制验证mask的效果# 测试mask功能 q_test torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]]) # (1,1,2,2) k_test torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]]) # (1,1,2,2) v_test torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]]) # (1,1,2,2) # 创建一个causal mask: [[1,0],[1,1]] mask_test torch.tensor([[[[1, 0], [1, 1]]]], dtypetorch.bool) # (1,1,2,2) out, weights scaled_dot_product_attention(q_test, k_test, v_test, maskmask_test) print(Attention weights with causal mask:) print(weights.squeeze().numpy()) # 预期输出: [[1.0, 0.0], [0.5, 0.5]]第二行显示位置1能看到位置0和1但第一行位置0只能看到自己3.3 第三块基石MultiHeadAttention——将原子核组装成引擎class MultiHeadAttention(nn.Module): def __init__(self, d_model: int, num_heads: int, dropout: float 0.0): super().__init__() assert d_model % num_heads 0, d_model must be divisible by num_heads self.d_model d_model self.num_heads num_heads self.d_k d_model // num_heads # 每个head的维度 self.d_v d_model // num_heads # 定义Q, K, V的线性投影层 # w_q, w_k, w_v 的shape都是 (d_model, d_k) 或 (d_model, d_v) self.w_q nn.Linear(d_model, d_model) # 输出d_model后续再split self.w_k nn.Linear(d_model, d_model) self.w_v nn.Linear(d_model, d_model) # 输出投影层将h个head的输出拼接后映射回d_model self.w_o nn.Linear(d_model, d_model) self.dropout nn.Dropout(dropout) def forward( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor None ) - torch.Tensor: query, key, value: (batch, seq_len, d_model) mask: (batch, 1, seq_len, seq_len) or (1, 1, seq_len, seq_len) 返回: (batch, seq_len, d_model) batch_size query.size(0) # Step 1: 线性投影 (batch, seq_len, d_model) - (batch, seq_len, d_model) q self.w_q(query) # (batch, seq_len, d_model) k self.w_k(key) # (batch, seq_len, d_model) v self.w_v(value) # (batch, seq_len, d_model) # Step 2: Reshape and transpose to (batch, num_heads, seq_len, d_k/d_v) # 这是Multi-Head的核心操作将d_model维度拆分为(num_heads, d_k) q q.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) k k.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) v v.view(batch_size, -1, self.num_heads, self.d_v).transpose(1, 2) # 现在 q, k, v 的shape都是 (batch, num_heads, seq_len, d_k/d_v) # Step 3: 调用核心Attention函数 # 注意mask需要适配新的维度 (batch, 1, seq_len, seq_len) # 如果mask是 (batch, seq_len, seq_len)需要unsqueeze(1)变成 (batch, 1, seq_len, seq_len) if mask is not None and len(mask.shape) 3: mask mask.unsqueeze(1) # (batch, 1, seq_len, seq_len) x, self.attn_weights scaled_dot_product_attention( q, k, v, maskmask, dropout_pself.dropout.p ) # Step 4: 把num_heads维度合并回d_model # x: (batch, num_heads, seq_len, d_v) - (batch, seq_len, num_heads, d_v) x x.transpose(1, 2).contiguous() # contiguous() 是必须的因为transpose后内存不连续view会报错 x x.view(batch_size, -1, self.d_model) # (batch, seq_len, d_model) # Step 5: 最终线性投影 return self.w_o(x)这个类的实现集中体现了前面所有前提的落地q.view(...).transpose(1, 2)这是Multi-Head的物理实现。view将(batch, seq_len, d_model)重塑为(batch, seq_len, num_heads, d_k)transpose(1, 2)将其变为(batch, num_heads, seq_len, d_k)为并行计算做好准备。transpose(1, 2)中的1和2分别指代seq_len和num_heads维度的索引。contiguous()的必要性这是PyTorch中一个极其重要但常被忽略的细节。transpose操作会改变张量的内存布局使其变得“不连续”。而view操作要求张量在内存中是连续的。如果不加contiguous()view会抛出RuntimeError: view size is not compatible with input tensors size and stride。这个错误在调试时非常隐蔽因为它只在特定的维度变化下才出现。mask.unsqueeze(1)的智能适配我们的scaled_dot_product_attention函数期望mask的shape是(batch, 1, seq_len, seq_len)但用户传入的mask可能是(batch, seq_len, seq_len)。这个unsqueeze(1)操作就是为了让接口更友好避免用户在每次调用时都手动增加维度。实操心得在实际项目中我习惯在forward函数末尾添加一句self.attn_weights self.attn_weights.detach()。因为attn_weights是一个巨大的中间变量shape(batch, num_heads, seq_len, seq_len)如果它一直保留在计算图中会极大增加显存占用。detach()将其从计算图中剥离只保留数值用于可视化或分析这对训练大型模型至关重要。4. 组装与验证构建完整Encoder Block并用真实数据流进行端到端压力测试有了PositionalEncoding和MultiHeadAttention这两块基石现在我们可以组装出完整的Encoder Layer。但请注意一个标准的Encoder Block不仅仅是Attention它还包括Add Norm和Feed-Forward NetworkFFN。我们将严格按照原始论文的结构实现并在每一步都插入验证点。4.1 EncoderLayerAttention AddNorm FFN AddNorm的完整流水线class FeedForward(nn.Module): Feed-Forward Network原始论文中为两层线性变换 ReLU def __init__(self, d_model: int, d_ff: int, dropout: float 0.1): super().__init__() self.linear1 nn.Linear(d_model, d_ff) self.dropout nn.Dropout(dropout) self.linear2 nn.Linear(d_ff, d_model) # 注意原始论文使用的是ReLU但后来的实现如BERT常用GELU # 这里我们严格遵循原始论文用ReLU self.activation nn.ReLU() def forward(self, x: torch.Tensor) - torch.Tensor: # x: (batch, seq_len, d_model) x self.linear1(x) # (batch, seq_len, d_ff) x self.activation(x) # (batch, seq_len, d_ff) x self.dropout(x) # (batch, seq_len, d_ff) x self.linear2(x) # (batch, seq_len, d_model) return x class EncoderLayer(nn.Module): def __init__( self, d_model: int, num_heads: int, d_ff: int, dropout: float 0.1 ): super().__init__() self.self_attn MultiHeadAttention(d_model, num_heads, dropout) self.feed_forward FeedForward(d_model, d_ff, dropout) # 两个LayerNorm层 self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout1 nn.Dropout(dropout) self.dropout2 nn.Dropout(dropout) def forward(self, x: torch.Tensor, src_mask: torch.Tensor None) - torch.Tensor: x: (batch, seq_len, d_model) src_mask: (batch, 1, seq_len, seq_len) 或 None # Sublayer 1: Self-Attention # 先Norm再Attention再Add残差连接 # 这是Pre-LNPre-LayerNorm结构原始论文用的是Post-LN # 但Pre-LN在训练初期更稳定是现代实现的常见选择 norm_x self.norm1(x) attn_output self.self_attn(norm_x, norm_x, norm_x, masksrc_mask) x x self.dropout1(attn_output) # Sublayer 2: Feed-Forward norm_x self.norm2(x) ff_output self.feed_forward(norm_x) x x self.dropout2(ff_output) return x class Encoder(nn.Module): 堆叠N个EncoderLayer def __init__( self, layer: EncoderLayer, N: int ): super().__init__() self.layers nn.ModuleList([copy.deepcopy(layer) for _ in range(N)]) self.norm nn.LayerNorm(layer.d_model) def forward(self, x: torch.Tensor, mask: torch.Tensor None) - torch.Tensor: for layer in self.layers: x layer(x, mask) return self.norm(x)这个EncoderLayer的实现严格遵循了原始论文的“Add Norm”范式但做了一个关键的现代改进Pre-LNPre-LayerNorm。原始论文Post-LN的顺序是x Attn(LayerNorm(x))而Pre-LN是x Attn(LayerNorm(x))。这个改动看似微小但实测下来Pre-LN能让模型在训练初期收敛得更快、更稳定尤其是在层数较多时。这也是为什么Hugging Face的BertLayer默认使用Pre-LN。4.2 端到端压力测试用真实数据流驱动暴露所有潜在Bug光有代码还不够我们必须用真实的数据流来“压测”它。下面是一个完整的、可运行的端到端测试脚本它会模拟一个真实的训练步骤从随机输入开始经过Embedding、PositionalEncoding、Encoder最后计算一个简单的loss。这个过程会暴露出90%的实现错误。import torch import torch.nn as nn import torch.optim as optim import copy # 1. 构建一个最小的Transformer Encoder模型 class SimpleTransformerEncoder(nn.Module): def __init__(self, vocab_size1000, d_model512, num_heads8, d_ff2048, N2, dropout0.1): super().__init__() self.embedding nn.Embedding(vocab_size, d_model) self.pos_encoding PositionalEncoding(d_model, dropout) encoder_layer EncoderLayer(d_model, num_heads, d_ff, dropout) self.encoder Encoder(encoder_layer, N) # 添加一个简单的分类头用于生成loss self.classifier nn.Linear(d_model, 2) # 二分类 def forward(self, src: torch.Tensor, src_mask: torch.Tensor None) - torch.Tensor: # src: (batch, seq_len) x self.embedding(src) * math.sqrt(self.embedding.embedding_dim) # 缩放 x self.pos_encoding(x) x self.encoder(x, src_mask) # 取第一个token的输出作为句子表示类似[CLS] cls_token x[:, 0, :] return self.classifier(cls_token) # 2. 创建测试数据 batch_size 2 seq_len 10 vocab_size 1000 model SimpleTransformerEncoder(vocab_sizevocab_size, d_model128, num_heads4, d_ff512, N2) # 生成随机输入和标签 src torch.randint(0, vocab_size, (batch_size, seq_len)) # (2, 10) labels torch.randint(0, 2, (batch_size,)) # (2,) # 3. 创建mask一个简单的padding mask假设所有token都是有效的 # 在真实场景中这里会是根据padding token生成的mask src_mask torch.ones(batch_size, 1, seq_len, seq_len) # (2, 1, 10, 10) # 4. 前向传播 output model(src, src_mask) print(fModel output shape: {output.shape}) # 应该是 (2, 2) # 5. 计算loss并反向传播 criterion nn.CrossEntropyLoss() loss criterion(output, labels) loss.backward() # 6. 检查梯度 print(\nGradient check:) for name, param in model.named_parameters(): if param.grad is not None: print(f{name}: grad_norm {param.grad.norm().item():.4f}) # 7. 关键验证检查Attention权重是否合理 # 我们可以访问最后一层EncoderLayer的self_attn.attn_weights # 由于我们封装了多层需要稍微绕一下 last_encoder_layer model.encoder.layers[-1] if hasattr(last_encoder_layer.self_attn, attn_weights): attn_weights last_encoder_layer.self_attn.attn_weights print(f\nLast layer Attention weights shape: {attn_weights.shape}) print(fMean attention weight: {attn_weights.mean().item():.4f}) print(fMax attention weight: {attn_weights.max().item():.4f}) print(fMin attention weight: {attn_weights.min().item():.4f}) # 检查是否满足概率分布每行和应为1 row_sums attn_weights.sum(dim-1) print(fRow sums (should be ~1.0): {row_sums[0, 0].tolist()})运行这个脚本你会得到一系列关键输出。**请务必
