LLM Lab 实验手记 04：张量并行入门——Column/Row Parallel 与 Megatron FFN

代码文件：llm_lab/03_tensor_parallel/tensor_parallel_demo.py、row_parallel_demo.py、megatron_ffn_demo.py

张量并行（Tensor Parallelism, TP）解决的是"单个权重矩阵放不下"的问题：把一个大矩阵切开分到多卡，每卡只存一部分，计算时用集合通信把结果拼回来。本篇用三个最小实验，把 TP 的两种切法（Column / Row）和 Megatron 的经典 FFN 拆分全部讲透。

1. 实验目标

1. 实现 Column Parallel（沿列切分权重）：每卡各算一部分输出，AllGather 拼回完整输出；
2. 实现 Row Parallel（沿行切分权重）：每卡算部分和，AllReduce 汇总；
3. 复刻 Megatron 的 FFN 两层（Column → GELU → Row），并讲清为什么这样配对。

2. 代码逐段拆解

实验 A：Column Parallel（tensor_parallel_demo.py）

world = dist.get_world_size()

hidden = 4096
chunk = hidden // world              # 每卡分到的列数

weight = torch.randn(hidden, chunk, device="cuda")   # ① 列切：hidden x chunk

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

local_out = x @ weight               # ② 每卡算出一部分输出

gather_list = [torch.zeros_like(local_out) for _ in range(world)]
dist.all_gather(gather_list, local_out)              # ③ AllGather 拼列

out = torch.cat(gather_list, dim=1)  # ④ 完整输出 [32, hidden*world]

• ① 权重按列切成 world 份，每卡持有 [hidden, chunk]；
• ② 每卡只算自己的输出列，结果形状 [32, chunk]；
• ③④ all_gather 把各卡的输出列收集齐，沿列拼成 [32, hidden]——和完整权重算出的结果逐元素一致。

一句话：Column Parallel = 输出拼列（AllGather）。

实验 B：Row Parallel（row_parallel_demo.py）

chunk = hidden // world

weight = torch.randn(chunk, hidden, device="cuda")   # ① 行切：chunk x hidden

x = torch.randn(32, chunk, device="cuda")            # ② 输入也按行分

local_out = x @ weight               # ③ 每卡算部分和

dist.all_reduce(local_out)           # ④ AllReduce 汇总

• ① 权重按行切成 world 份，每卡 [chunk, hidden]；
• ② 输入 x 也必须按同样方式切好（各卡持有不同的行片段）；
• ③④ 每卡算出的 local_out 是完整结果的部分和，all_reduce 求和得到完整输出——且所有卡拿到同一份。

一句话：Row Parallel = 输出求和（AllReduce）。

实验 C：Megatron FFN（megatron_ffn_demo.py）

chunk = ffn_hidden // world

w1 = torch.randn(hidden, chunk, device="cuda")   # Column：hidden x chunk
w2 = torch.randn(chunk, hidden, device="cuda")   # Row：chunk x hidden

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

local_mid = x @ w1                 # 第一层：Column Parallel
local_mid = F.gelu(local_mid)      # 激活函数，逐元素，无需通信

local_out = local_mid @ w2         # 第二层：Row Parallel
dist.all_reduce(local_out)         # 汇总

关键观察：Column 后接 Row 时，中间的 AllGather 被省掉了。

3. 原理讲解：为什么 Column → Row 不需要中间通信

这是 Megatron 设计最精妙的一处。先看朴素拆法：

x ──(Column w1)──> mid（需 AllGather 拼全）──(激活)──> mid ──(Row w2)──> out（需 AllReduce）

Column 层的输出若在逻辑上"拼全"，就要 AllGather；但如果下一层正好是 Row Parallel，Row 层要求输入按行分片——Column 层的分片输出正好就是 Row 层的分片输入！于是中间的 AllGather 彻底省略：

x ──(Column w1)──> mid_本地（天然分片）──(GELU 逐元素)──> mid_本地 ──(Row w2)──> part ──(AllReduce)──> out

两层共享一次通信：整个 FFN 块只需要一次 AllReduce（在 Row 层输出端），而不是两次。扩展到整个 Transformer：每个 Block 的 FFN 一次通信、Attention 的 QKV（Column）与输出投影（Row）也配对省略中间通信——全模型通信点大幅减少，这正是 Megatron 张量并行能跑得动的关键。

（GELU 逐元素、不需要通信，也是能省中间同步的原因之一。）

4. 实验结果

三个实验以 4 卡运行，预期输出（rank 0 打印形状）：

# 实验 A：Column Parallel
Output shape: torch.Size([32, 4096])

# 实验 B：Row Parallel
Output shape: torch.Size([32, 4096])

# 实验 C：Megatron FFN
FFN output shape: torch.Size([32, 4096])

如何判断实验成功：将单卡完整权重计算的输出作为基准，分布式输出应与之一致（可加一行 torch.allclose 断言）。形状正确只是第一步，数值一致才是通信与切分正确的证明。

5. 工程启示

1. 两种切法的通信成本不同：Column 用 AllGather（通信量 = 输出大小），Row 用 AllReduce（通信量 ≈ 2× 输出）。所以通信敏感场景优先 Column 切（反向还不需要通信），Row 只用在必须配对的位置（FFN 第二层、Attention 输出投影）。

2. "切分方向 × 邻接层"决定通信点数量：设计并行方案时，逐层画"切分状态传播图"——输入处于什么分片状态、下一层需要什么状态、中间需不需要集合操作。Megatron 的精髓就是让分片状态自然衔接，把通信点压到最少。

3. TP 的适用边界：TP 通信最密集、且每步都通信，通常只在节点内（NVLink 高速互联）使用；跨节点 TP 会让节点间流量爆炸。这是 3D 并行里 TP 度一般取 4–8 的硬件原因。

4. world_size 必须能整除维度：本实验 chunk = hidden // world 假设整除。真实代码要显式检查，否则切片错位会造成静默算错——比报错更难排查。

下一步：FFN 会拆了，Attention 怎么拆？第 5 篇按 Head 切分 QKV：LLM Lab 实验手记 05：注意力张量并行。