企业官网建设流程全解析

最近和一场模拟面试的经历撞了个正着：候选人很流畅地写出了多头注意力模块，手推公式也没问题，但当我问起“为什么源码里要把 Q、K、V 一起经过一个大线性层再拆开”“permute 之后张量连续吗”“训练和推理的 mask 到底差在哪”时，对方明显卡住了。这不是个例，而是很多刚开始接触 Transformer 源码的同学共同的痛点：嘴上会讲 attention，手上却不知道每一行程式码在干什么。这篇记录就是从那场模拟面试延伸出来的——把 Attention 计算、维度变换、KV Cache、分布式训练这几个高频追问点，掰开揉碎讲清楚，顺便给出我这些年带人和面试时沉淀的一条源码阅读路线。

1. 为什么学了Transformer的人，一被追问源码就露馅

1.1 一次模拟面试的真实开场

这场模拟面试完全没有用真实项目，就考察一个最基础的多头注意力模块。候选人 A 同学在共享屏幕上敲出了类似下面这段代码，速度很快：

class MultiHeadAttention(nn.Module): def __init__(self, d_model, n_heads): super().__init__() self.n_heads = n_heads self.d_model = d_model self.head_dim = d_model // n_heads self.qkv = nn.Linear(d_model, 3 * d_model) self.proj = nn.Linear(d_model, d_model) def forward(self, x, mask=None): B, T, D = x.shape qkv = self.qkv(x).reshape(B, T, 3, self.n_heads, self.head_dim).permute(2, 0, 3, 1, 4) q, k, v = qkv[0], qkv[1], qkv[2] attn = q @ k.transpose(-2, -1) * (self.head_dim ** -0.5) if mask is not None: attn = attn.masked_fill(mask == 0, float('-inf')) attn = attn.softmax(dim=-1) out = attn @ v out = out.transpose(1, 2).reshape(B, T, D) return self.proj(out)

这段代码本身没什么大问题，甚至已经在用* (head_dim ** -0.5)替代torch.div来避免额外传参。但面试不是让你默写代码，面试官真正关心的是“中间每一步的形状变化是什么”“这样写为什么能工作”“换成view会怎么样”“mask 里填-inf之后 softmax 会不会意外把 NaN 算出来”。

当我顺着代码逐行追问之后，A 同学承认，自己平时只调用封装好的模块，很少真的打开源码看 shape 是怎么一路变过来的。这是很多准备模型岗实习的同学的通病：背了结构图，但没掌握源码的计算图。

1.2 面试官真正想考察的底层能力模型

后来我把这场模拟面试中反复追问的角度拆解了一下，发现面试官考察的其实不是“记性”，而是三层能力。

第一层是张量形状变换能力。你能不能在心里跟着reshape -> permute -> transpose -> view走完一次 forward？能不能说出来什么时候需要.contiguous()，什么时候不需要？这决定了你调试模型时能不能快速定位维度不匹配的问题。

第二层是数值稳定性与计算图意识。为什么缩放因子要除以根号d_k？为什么 softmax 前要 mask 成-inf？为什么训练设model.train()的时候 dropout 会改输出分布？这些细节直接关系到你训练一个大模型时会不会出现 loss 震荡、推理时会不会每步结果都不一样。

第三层是工程实现与加速视角。同样的 attention 计算，在单卡、多卡、推理、训练这四种场景下，源码组织方式完全不同。KV Cache 在哪里生效？张量并行时哪些矩阵需要切，哪些不需要？FlashAttention 到底省的是哪部分显存？面试官不会问得太虚，但如果你能主动讲到这些，和只会背公式的候选人马上就拉开了差距。

接下来我按那天追问的顺序，把每一条都展开成“源码级”的解析。

2. Attention核心计算的源码级拆解：Q、K、V到底经历了什么

2.1 从公式到实现：缩放点积注意力的每个算子的作用

先看原始公式：Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V。这里最容易被追问的就是sqrt(d_k)到底解决什么问题。

假设q和k的每个分量都是均值 0、方差 1 的随机变量，它们的点积期望是 0，但方差是d_k。如果不做缩放，当d_k比较大时，点积结果就可能非常大，比如落在几百到上千的区间。之后接 softmax，这些大的 logits 会直接把概率分布推向极端：最大的 logit 对应的概率接近 1，其余接近 0。你可以想象一张灰度图几乎只剩纯黑和纯白，所有中间细节都没了。这个状态对应的梯度会非常小，因为 softmax 在饱和区的导数趋近于 0，模型几乎学不动。

所以源码里通常写成attn = q @ k.transpose(-2, -1) * (self.head_dim ** -0.5)。这里的head_dim ** -0.5就是1/sqrt(head_dim)，不是随手乘的，而是在保持点积方差为 1 的尺度。理解这一点后，面试官再问你“为什么不用d_model而用head_dim”，你就能答上来：因为每个 head 内部点积的维度是head_dim = d_model // n_heads，缩放只作用于这个维度。

接着是 softmax。源码里往往写成attn = attn.softmax(dim=-1)，默认不是对最后一维做归一化吗？对，它的作用是把注意力权重变成总和为 1 的分布，并且对负无穷值会自然得到 0 的概率。这里需要注意：如果 mask 填充的是-inf，softmax 计算时exp(-inf)是 0，不会报错；但如果你不小心写成-1e9这样很小的负数，exp也可能得到 0，不过数值上可能存在上溢风险，因此正规实现都用-inf。

最后out = attn @ v就是把加权后的 value 聚合起来。很多初学者会忽略这一步的 shape：attn的形状是[B, heads, T, T]，v的形状是[B, heads, T, head_dim]，矩阵乘出来是[B, heads, T, head_dim]。

2.2 reshape、transpose、view的三重陷阱与维度变化细节

源码里最让新手头疼的其实是reshape和permute的组合。我们拆开看。

输入x的形状是[B, T, D]，其中D = d_model。经过self.qkv(x)，得到[B, T, 3*D]。然后执行：

qkv.reshape(B, T, 3, self.n_heads, self.head_dim).permute(2, 0, 3, 1, 4)

为什么先 reshape 成 5 维，再 permute？因为qkv最后那个3*D维里，顺序是“先放 Q 的全部分量，再放 K 的全部分量，最后放 V 的全部分量”。如果直接 reshape 成[B, T, 3, n_heads, head_dim]，那么这个三维中的第 3 维（索引 2）实际上表示 Q/K/V 的编号。随后permute(2, 0, 3, 1, 4)把这一维挪到最前面，得到形状[3, B, n_heads, T, head_dim]。

上面代码里q, k, v = qkv[0], qkv[1], qkv[2]，其实就在第一维上切“第 0 份、第 1 份、第 2 份”。这里有个很隐蔽的点：permute之后张量在内存里不是连续的，因为原来的存储顺序是B, T, (Q/K/V), heads, head_dim，现在逻辑顺序变成了Q/K/V, B, heads, T, head_dim，但底层的物理排列没变。若你在permute之后调用view而不是reshape，大概率会报错。正确做法是写q = qkv[0].contiguous()之后再用view，或者干脆继续用reshape，让它内部帮你处理内存拷贝。

这部分太适合被连环追问了。面试官通常会这么问：

• “reshape和view的区别？”
• “permute之后为什么需要contiguous()？”
• “如果不加contiguous()，后面的矩阵乘能正常跑吗？”

第一个问题：view要求原张量满足内存连续，reshape不要求，不连续时它会自动复制到连续内存再变换。第二个问题：permute只改变逻辑视图，不改变物理布局，只要后续算子不是对最后一维做简单索引，很多情况下也能跑，但会多出非连续访存的开销。第三个问题：矩阵乘法q @ k.transpose(...)本身会自动处理步长信息，不一定报错，但如果你再派生出其他操作，就可能因为张量不连续触发行列访问顺序混乱，性能骤降、内存占用却激增。

2.3 Mask的三种形态：padding mask、causal mask、组合mask

在源码中，mask 处理是 attention 正确性的关键。不要把 mask 理解成一张简单的布尔面具，它要跟实际形状匹配。

第一种是padding mask：用于忽略句子中补齐部分的 token。形状通常是[B, T]或者[B, 1, 1, T]，表示哪些位置是有效 token。在注意力分数矩阵里，对 key 维度做 mask，也就是把[B, heads, T_q, T_k]中对应 key 位置的分数填成-inf。

第二种是causal mask：用于自回归任务，保证第i个 token 只能看到0~i的位置，不能看到未来。双向 attention 的 Transformer 如果用在生成任务上，不加 causal mask 就等于作弊。源码中常用torch.triu生成一个上三角矩阵，把右上角填成-inf：

causal_mask = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)

然后在注意力分数上执行attn = attn.masked_fill(causal_mask, float('-inf'))。

第三种是组合 mask：padding mask 和 causal mask 同时存在时，要先做逻辑或或者逐个填-inf。实际实现里可能会把两个 mask 都 expand 成相同的[B, 1, T, T]形状，再融合成一个 mask。常见错误是把 padding mask 直接 broadcast 到[B, T, T]时维度对不上，结果 mask 被作用于 query 维度而不是 key 维度，模型看到了一堆错误位置。

这里还有一个面试官爱问的小点：masked_fill是不是 in-place 的？严格说，masked_fill在 PyTorch 文档里“不是原位操作”，它会返回一个新的张量。很多人会把attn = attn.masked_fill(...)里的重新赋值省略，写成attn.masked_fill(...)，导致 mask 根本没生效。这属于写法上的隐性 bug，但面试时能主动提出来，说明你真的抠过源码。

3. 那些“一读就懂、一问就忘”的Transformer源码细节

3.1 多头注意力为什么要按head拆，而不是直接算一整块矩阵

原始论文给出的 motivation 是让模型在不同的表示子空间里关注不同信息。但用代码的角度看，多头的本质是什么？是把d_model维的向量切分成n_heads份，每一份维度是head_dim，然后每一份单独做 attention，最后拼回去。

如果我们不做多头，直接让Q, K, V的维度都是d_model，注意力计算就是一个d_model × d_model的矩阵乘，相当于用一个高维参数矩阵去拟合所有的注意力模式。这个做法的缺点是：所有信息混在同一组矩阵里，很难让模型学到“有的头关注位置关系、有的头关注语义关系、有的头关注局部句法”这类分离表示。多头相当于给模型提供了多个相关性度量通道，每个头只在自己的低维子空间里建模，避免单个参数矩阵把所有模式混成一团。

从参数数量上看，多头并不会增加多少参数。一个多头注意力模块里，qkv线性层的输出还是3*d_model，proj线性层的输出还是d_model。只不过在内部把矩阵分块了。很多同学以为多头会带来参数量爆炸，其实没有。

面试现场我还追问了 A 同学一句：“如果d_model不能整除n_heads，代码会发生什么？”答案是直接报错，因为d_model // n_heads会向下取整，剩下的维数要么被丢掉，要么导致维度对不上。所以实际实现中d_model必须能被n_heads整除。这在设计模型时就要提前想好。

3.2 LayerNorm与残差连接的顺序：pre-norm vs post-norm的源码差异

Transformer 论文里的结构是x + Sublayer(LayerNorm(x))，也就是先归一化，进子层，再加残差。但今天大量开源实现的 Transformer Block 是另一种顺序：

x = x + attn(LayerNorm(x)) x = x + ffn(LayerNorm(x))

这就是 pre-norm 结构。区别在哪呢？

我们先看 post-norm（论文原始版本）的逻辑：先过attention或ffn，再归一化，再加到输入上。残差分支上直接输入的是未经归一化的表征，优化时梯度要穿过一个非线性层再加回来，路径比较曲折，容易造成训练初期不稳定。

pre-norm 把归一化放在子层之前，主分支上经过残差加和的是已经归一化后的信息，相当于把“尺度归一”这件事提前做掉，让每一层输入分布更稳定。另一个好处是，即便网络很深，梯度也能通过残差短接线直接回传，和归一化模块相对独立。因此现在主流大模型几乎全是 pre-norm。

面试时还可以补充一个细微差别：pre-norm 在输出层的均值和方差统计可能与 post-norm 不同，所以有些实现最终还会在最外面接一个 final LayerNorm。这也是为什么你会在源码里看到 Transformer Block 的输出过一个ln_f或final_norm。

3.3 推理加速背后的KV Cache：源码中缓存的是什么

每次生成一个 token 时，如果都把之前所有 token 的 Q/K/V 重新算一遍，计算量会随序列长度线性增长，非常浪费。KV Cache 的思路是：在自回归解码中，每个新 token 的 Q 需要计算，但它对应的 K、V 会参与后续所有 token 的注意力计算，所以只要把已经生成的 token 的 K、V 保存下来，下一轮直接拼接新 token 的 K、V，再用新 query 和拼接后的 cache 做 attention。

具体到源码，大致是：

if past_key_value is not None: k = torch.cat([past_key_value[0], k], dim=-2) v = torch.cat([past_key_value[1], v], dim=-2)

这里在倒数第二维（序列维度）上拼接，新的past_key_value再传给后续层。面试官可能会追问“为什么只缓存 K 和 V，不缓存 Q？”关键点在于 Q 只用于当前 token 的位置，之后的 token 不会再去看“当前 token 的 query”，但后续 token 需要和“当前 token 的 key”做匹配，也要用“当前 token 的 value”做加权，所以 K 和 V 是必须长期保留的。Q 每步都不同，用完就没有保留意义了。

如果继续追问“为什么 cache 增大了显存反而生成变快”，可以从计算量角度答：没有 cache 时，第 N 步要重新计算前 N 个 token 的 K/V；有 cache 后每步只计算新的一个 token 的 K/V，操作量从 O(N^2) 降到 O(N)，省下的时间远大于拼接操作的开销。当然代价是显存，这也是长文本生成时显存容易爆炸的原因之一。

4. 从单卡到多卡：Transformer源码在分布式场景下的隐藏逻辑

4.1 数据并行时，梯度同步的瓶颈为什么集中在LayerNorm和Embedding

面试官把一个单卡 Transformer 放到多卡场景，问的第一句往往是“数据并行怎么做？”最直接的答案是：把一个小批量切成多份，每张卡一个副本，前反向独立计算，算完梯度后用 all-reduce 做平均同步，再更新各自副本上的参数。

那为什么说瓶颈集中在 Embedding 和 LayerNorm？

先看 Embedding 表。假设字典大小是 100k，隐藏维度是 4096，Embedding 就有 4 亿个参数，占了模型参数很大一块。数据并行下每张卡都要维护一份完整的 Embedding，all-reduce 梯度的通信量和参数数量成正比，因此 Embedding 的通信开销非常大。LayerNorm 的参数虽然少，但它也参与反向传播，会生成梯度，且在每张卡上都会计算和同步，中间包括均值和方差统计，同样逃不掉通信。

面试中如果答到这里，还能主动提一句“所以很多人会做 Embedding 的分片，比如让每张卡只管部分词表，再通过 all-gather 拼起来”，马上会加分。部分并行训练框架里甚至会单独为 Embedding 做词表并行，就是为了绕过这个通信瓶颈。

4.2 张量并行中，Attention的列切分与行切分如何避免大量通信

张量并行解决的是“单张卡放不下模型”的问题。思路是把一个矩阵乘法切到多张卡上算，然后在合适的位置做通信。

对 Transformer Block 里的 MLP 来说，第一层线性层一般做“列切分”，即把权重按列切成多块，每个并行 rank 各算输出的一部分；第二层线性层做“行切分”，把权重按行切成多块，每个 rank 在各自分块上做局部计算，最后通过 all-reduce 把各部分结果累加起来。这样设计的好处是：只在两层之间做一次 all-reduce，中间不需要全通信，避免每层都同步。

Attention 部分更特殊。多头的天然并行性让它很适合按 head 维度切分：把n_heads个 head 平均分给多个 rank，每个 rank 算自己负责的那些 head，算完输出后把各 rank 的[B, local_heads, T, head_dim]拼回[B, n_heads, T, head_dim]，再通过一个通信操作完成拼接。这套方案在 Megatron-LM 之类的并行范式里被称为多头并行。

但要注意：拼完之后还有一个proj线性层，它会对d_model维做线性变换。为了减少通信，通常会把这个proj的输入分布到各 rank，再做一次 all-reduce 或者把拼接和 proj 融合起来。所以面试时不要只记住“按 head 切分”，要能说清楚通信点到底在哪几次矩阵运算之间。

4.3 序列并行与FlashAttention的融合技巧：面试中的加分回答

序列并行是分布式场景里比较进阶的题目。它把序列长度T切到多张卡上，每张卡只处理一个序列片段。这样做能同时降低 per-card 的激活显存和计算量，但有一个天然矛盾：注意力分数矩阵的 softmax 归一化需要跨整个序列的统计量，如果序列被切开，每个片段算出的 softmax sum 并不全局。

FlashAttention 的思路恰好能帮忙。它通过 online softmax，把归一化统计量拆成可累加的部分和，在扫描分块时动态更新。序列并行场景下，每个 rank 可以独立扫描自己的序列块，再把局部的统计量通过通信合并。面试时如果能说出“online softmax 维护当前最大值和当前指数和，合并分块时用新最大值重新缩放旧的指数和”这个细节，几乎等于给了满分。

这里还有一层融合：注意力计算里的 dropout、mask、矩阵乘都可以在切分块内完成，尽力避免把完整的T × T注意力矩阵物化出来。FlashAttention 本身就是不物化大矩阵、利用分块减少 HBM 访问的典型，这和序列并行天然契合。回答时可以把两者的共同点概括为“把大的矩阵运算拆成小块，用少量额外计算换大量显存和访存”。

5. 高频连环追问拆解：像面试官一样追问，直到源码细节

5.1 第一追问链：为什么除以根号d、为什么软max、为什么多头

这条链我几乎每次模拟面试都会问一遍，每个问题都可以往下钻两层。

先问“为什么除以根号 d”：不缩放会怎样？计算方差的公式是什么？如果 d 是 512，缩放和不缩放的注意力分数量级差多少？要能答出“点积方差正比于维度”这一点。

再问“为什么用 softmax”：能不能换成 sigmoid？softmax 保证概率和为 1，可以处理变长序列；sigmoid 不能。当然也存在一些变体用 ReLU attention，但那是另一套设计。面试里答“softmax 是为了归一化”还不够，要补充“它让较大的 logits 在指数运算下受到抑制，梯度不会全部消失”。

最后问“为什么多头”：除了子空间表征之外，还要提并行计算的方便性。多个 head 可以在同一个 batch 内矩阵计算，把[B, heads, T, head_dim]的四维张量合并成三维矩阵乘，底层 GEMM 效率更高。也就是说，多头不只是建模能力的需求，也是工程计算效率的需求。这个角度很多人没提，提出来会让人眼前一亮。

5.2 第二追问链：训练时用dropout，为什么推理时通常不开启

这条链围绕 dropout 展开。

源码里nn.Dropout(p=0.1)的行为受self.training控制。训练时，每个元素有概率p被置 0，同时剩余元素除以1-p，目的是保持期望不变。推理时如果还开着 dropout，模型输出会随机波动，生成结果不稳定，而且等于在预测时引入了无意义的噪声。所以推理前要调用model.eval()，源码里模块内部的dropout.training标志也会被设成 False。

面试官通常还会追问一句：“如果训练时忘记改成 eval，模型会怎样？”A 同学当时答“结果可能每次不一样”，但更关键的是“性能会下降，因为 dropout 会破坏已经习得的表征，尤其在生成任务里可能出现明显的随机重复或突变”。能抓住这两层的候选人，证明他不是只会调 API。

5.3 第三追问链：显存不够怎么办？从重计算到混合精度的模型策略

当模型大到单卡放不下，面试官大概率会问“训练时怎么省显存”。大多数人都能想到混合精度，但能说出具体手段的人不多。

第一个方向是重计算，也叫 activation checkpointing / gradient checkpointing。训练时我们不保存每一层的所有激活，只保存每层的输入或选定的 checkpoint，反向传播用到中间激活时重新算一遍。这样做的效果是前向激活显存从 O(层数 × 每层激活大小) 降到 O(层数 × checkpoint 数量)，典型节省比例能到 60% 以上，代价是约 30%~50% 的额外计算时间。

第二个方向是混合精度。用 FP16 或 BF16 存权重、梯度和激活能大幅减少显存，但 FP16 容易造成梯度下溢，所以需要 loss scaling；BF16 的动态范围更大，但精度更低。现在很多框架默认就是 BF16 + 重计算 + 序列并行三件套。

第三个方向是优化器状态分片，比如 ZeRO 的思想：把优化器状态平均分到多卡，每张卡只维护自己分到的部分。这样在 64 卡上，优化器状态每卡只占原来的 1/64。提到这个，面试官如果深问，就能引出“参数、梯度、优化器状态三层分片”的顺序。

这条链回答完毕，基本就覆盖了大模型训练中最常被问的显存优化全家桶。

6. 三个月吃透Transformer源码的实践路线

6.1 用日志和断点把张量维度“看”在手里

我建议所有准备模型源码面试的同学，不要直接对着源码发呆，而是写一个“维度探测器”。比如在.forward里加几个日志打印：

def log_tensor(name, t): print(f"{name}: shape={tuple(t.shape)}, dtype={t.dtype}, " f"mean={t.mean():.4f}, std={t.std():.4f}")

然后在每个子层输入输出调用log_tensor。这样你就建立了“抽象维度到实际数值”的感知。比如在 attention 内部，你能亲眼看到q在进入@ k.transpose(-2, -1)之前的 shape 是[B, heads, T, head_dim]，也能感觉到为什么需要转置。

善用 PyTorch 的register_forward_hook也能省事，但自己 print 更直观。调试时随便跑一个小序列，序列长度 8，batch 2，隐藏维度 8，打印出来的维度变化会非常清晰。把每一层的[B,T,D]变化记录在一张纸上，基本上不需要死记源码了。

6.2 必须手写一遍的模块清单和验证方法

我给 A 同学的建议是，在三个月内按顺序手写以下几个模块，并做一致性验证：

每个模块验证时，用同一个输入张量，固定所有随机种子，对比手写实现和标准实现输出的最大误差max(|out_hand - out_lib|) < 1e-5。如果误差偏大，先检查自己是否漏了head_dim ** -0.5，再检查维度 permute 是否多转了一次。这套排查流程本身就是很宝贵的面试素材。

6.3 最后的避坑建议：别背源码，要理解设计意图

吃透源码不等于把每行代码背下来。面试官问你“为什么这里用 reshape 而不用 view”，本质上不是考察 API 用法，而是考察你是否知道 reshape 和 view 背后的内存连续性问题。同理，“为什么推理时要关 dropout”，也不是考你 API 开关，而是考你对训练和推理分布差异的感知。

我自己的体会是，读 Transformer 源码最顺的路径是：先手推一遍 attention 的数学公式，再对照一段最小的多头注意力代码，最后再往工程方向扩展（mask、KV Cache、分布式、 FlashAttention）。每一步都问自己三个问题——输入形状是什么、输出形状是什么、如果我把某个算子换成替代实现会怎样。带着这三个问题去读，三个月后你不仅能应对模拟面试，还能在真实模型调试时快速定位问题。

最后再分享一个实用小技巧：很多同学把源码读完了还是记不住维度，那就自己画一张“张量流向图”。不需要软件，纸上画也行，标清楚每层开始时的[B,T,D]，以及进了 attention 后变成[B,heads,T,head_dim]，出来又变回[B,T,D]。这张图比任何笔记都管用，也是你面试时可以随手在草稿纸上复现的重要武器。