LLM Lab 实验手记 05：注意力张量并行——按 Head 切分 QKV 的 Megatron 写法

代码文件：llm_lab/03_tensor_parallel/tp_attention_demo.py

上一篇讲完 FFN，本篇处理 Transformer 的另一半——Self-Attention。Attention 有个天然优势：多头彼此独立，这是它比 FFN 更适合并行的结构红利。

1. 实验目标

1. 实现多头注意力的张量并行：按 Head 数切分 QKV 权重；
2. 在本地完成 Q/K/V 投影、注意力计算，AllGather 拼回完整输出；
3. 理解 Megatron Self-Attention 的"三列权重 + 一拼"结构。

2. 代码逐段拆解

完整代码（tp_attention_demo.py，63 行）：

world = dist.get_world_size()

batch = 4
seq = 128
hidden = 4096

heads = 32
local_heads = heads // world      # ① 每卡分到的 head 数
head_dim = hidden // heads        # 每 head 维度

local_qkv_dim = local_heads * head_dim

wq = torch.randn(hidden, local_qkv_dim, device="cuda")   # ② 三列权重
wk = torch.randn(hidden, local_qkv_dim, device="cuda")
wv = torch.randn(hidden, local_qkv_dim, device="cuda")

x = torch.randn(batch, seq, hidden, device="cuda")

q = x @ wq                        # ③ 本地 Q/K/V 投影
k = x @ wk
v = x @ wv

q = q.view(batch, seq, local_heads, head_dim)
k = k.view(batch, seq, local_heads, head_dim)
v = v.view(batch, seq, local_heads, head_dim)

scores = torch.matmul(q, k.transpose(-1, -2)) / (head_dim ** 0.5)   # ④ 本地注意力
attn = F.softmax(scores, dim=-1)

out = torch.matmul(attn, v)

out = out.reshape(batch, seq, local_qkv_dim)

gather_list = [torch.zeros_like(out) for _ in range(world)]
dist.all_gather(gather_list, out)         # ⑤ AllGather 拼回全部 head

full_out = torch.cat(gather_list, dim=-1) # 完整输出 [batch, seq, hidden]

① local_heads = heads // world

Attention 并行的关键决策：不切 head 内部的维度，而是把 32 个 head 平分给各卡。每卡只负责 local_heads 个 head 的完整计算。相比切 head_dim，切 head 让每卡的计算自包含（Q/K/V 都在本地），注意力内部零通信。

② 三列权重（Column Parallel 风格）

Q/K/V 三个投影权重都是 [hidden, local_qkv_dim]——沿列切分，不需要 AllGather 拼回（因为下一层"拼接 head"正好是 Row 层想要的分片状态）。这就是 Megatron 的 Self-Attention：三个 Column 权重，无中间通信。

③④ 本地投影与本地注意力

投影后 reshape 成 4D（batch, seq, local_heads, head_dim），接下来 q @ kᵀ / √d、softmax、attn @ v 全部在本地完成——注意力分数、softmax、加权求和都不需要跨卡。

⑤ 唯一的通信点：输出端 AllGather

计算完本地 head 的输出后，all_gather 收集各卡的 head 片段，沿最后一维拼接成 [batch, seq, hidden]。整个 Attention 块只有这一个集合操作。

3. 原理讲解：注意力并行的"结构红利"

为什么 Attention 比 FFN 更好并行？三个原因：

1. Head 天然独立：32 个 head 互不相干，切到 4 卡就是每卡 8 个 head 的独立注意力，并行度直接等于 head 数（本实验 32 head → 最多 TP=32）。

2. Q/K/V 投影可免通信：Column 切分下，各卡只算自己的 head，不需要看别的 head 的 Q/K——注意力计算完全本地化。

3. 通信只在收尾：唯一一次 AllGather 发生在输出拼接，且与后续层（Row 风格的输出投影）配对时还能再省一次（同第 4 篇的 Column→Row 衔接）。

4. 实验结果

以 4 卡运行（32 head → 每卡 8 个），预期输出：

Attention output shape: torch.Size([4, 128, 4096])

如何判断实验成功：与单卡完整注意力对比 allclose。需要特别核对每卡的 head 顺序——拼接后的 head 顺序必须与切分约定一致（head 0–7 在卡 0，8–15 在卡 1……），否则输出"数值全错但形状全对"，是最隐蔽的一类 bug。

5. 工程启示

1. head 数与 TP 度的整除约束：heads // world 必须整除。32 head 可支持 TP=2/4/8/16/32，但不能 TP=6。设计集群并行配置时，模型 head 数是先验约束。

2. 通信量分析：一次 AllGather（输出大小 = batch × seq × hidden）——比 FFN 的一次 AllReduce（2× 输出）还少。Attention 是张量并行里通信最省的部分，这也是为什么 TP 下 Attention 通常不是瓶颈。

3. 拼接顺序约定要写进注释：跨卡 head 的顺序、local_heads 的编号规则，是 TP 代码最容易"静默算错"的地方。生产代码里建议在切分后加 torch.allclose 断言（对照单卡），成本极低。

4. 与 FlashAttention 的组合（第 7 篇预告）：生产实现中 local_heads 的注意力计算会替换为 flash_attn_func——TP 的分片逻辑不变，只换本地计算内核，这正是并行框架的模块化威力。

下一步：Attention 和 FFN 都会切了，现在把它们拼成一个完整的 Megatron 风格 Transformer Block——第 6 篇：LLM Lab 实验手记 06：从零拼一个 Megatron 风格 Transformer Block。