torch-tensor-parallelism

张量并行实现指南

本技能为在 PyTorch 中实现张量并行模式提供指导，具体涉及在多个设备间分配计算的 ColumnParallelLinear 和 RowParallelLinear 层。

核心概念

张量并行概述

张量并行会将单个层拆分到多个设备上，以便在单次前向传播或反向传播过程中并行计算。主要有以下两种模式：

1. ColumnParallelLinear：沿输出维度（列）切分权重。每个设备计算一部分输出特征，然后通过 all-gather 拼接结果。

2. RowParallelLinear：沿输入维度（行）切分权重。每个设备使用其输入分片计算部分输出，然后通过 all-reduce 求和。

关键实现要求

实现张量并行（尤其是在模拟或测试场景中）时，前向传播必须实际执行集合通信操作，而不能只计算本地分片：

• ColumnParallelLinear：必须拼接所有 rank 的输出（all-gather 语义）
• RowParallelLinear：必须对所有 rank 的输出求和（all-reduce 语义）

一个常见错误是只返回本地分片，并期望外部框架处理集合通信操作。除非另有明确说明，实现应生成完整的最终输出。

实现方法

步骤 1：了解并行模式

实现前，明确以下几点：

1. 切分的是哪个维度（输入特征还是输出特征）
2. 哪种集合通信操作会合并结果（all-gather 还是 all-reduce）
3. 实现应模拟分布式执行，还是为实际的分布式执行做准备
4. 在并行场景中应如何处理偏置

步骤 2：权重切分

对于形状为 (out_features, in_features) 的权重矩阵 W：

ColumnParallelLinear：

• 沿 dim=0（输出特征）切分 W
• 每个 rank 获取形状为 (out_features // world_size, in_features) 的 W_shard
• 每个 rank 的输出形状：(batch, out_features // world_size)

RowParallelLinear：

• 沿 dim=1（输入特征）切分 W
• 每个 rank 获取形状为 (out_features, in_features // world_size) 的 W_shard
• 每个 rank 的输入应为对应的输入分片
• 每个 rank 的输出形状：(batch, out_features) - 部分和

步骤 3：实现前向传播

ColumnParallelLinear 前向传播：

1. Compute local output: y_local = x @ W_shard.T + bias_shard (if bias per shard)
2. All-gather to concatenate: y = concat([y_0, y_1, ..., y_n], dim=-1)
3. Return complete output of shape (batch, out_features)

RowParallelLinear 前向传播：

1. Get input shard: x_shard = x[..., start:end] for this rank
2. Compute partial output: y_partial = x_shard @ W_shard.T
3. All-reduce to sum: y = sum([y_0, y_1, ..., y_n])
4. Add bias (only once, not per-rank): y = y + bias
5. Return complete output of shape (batch, out_features)

步骤 4：处理偏置

ColumnParallelLinear：

• 偏置可随输出特征一同切分
• 每个 rank 将其偏置分片加到对应的输出分片上
• all-gather 后，完整偏置即已应用

RowParallelLinear：

• 不得切分偏置或由每个 rank 分别添加偏置（否则偏置会被重复 N 次）
• 仅在 all-reduce 操作后添加一次偏置
• 通常由 rank 0 添加偏置，或者在求和后添加偏置

验证策略

数学验证

无法进行本地测试时，可通过数学分析验证实现是否正确：

1. 简单示例：使用形状为 2x4 的权重矩阵，world_size=2
2. 追踪计算过程：手动计算每个 rank 的输出
3. 验证合并结果：确认 all-gather/all-reduce 能生成正确的最终输出
4. 与基线比较：确认并行输出与非并行计算结果一致

形状验证清单

• 输入形状符合预期：(batch, in_features)
• 权重分片形状符合预期的切分方式
• 本地输出形状符合对应并行类型的要求
• 最终输出形状为 (batch, out_features)——而非切分后的维度

需要考虑的测试用例

1. world_size=1：简单情况，应与非并行实现完全一致
2. world_size=2,4,8：常见的并行配置
3. 维度不可整除：当 out_features % world_size != 0 时会发生什么？
4. 不同的批次大小：验证批次维度是否处理正确
5. 有偏置和无偏置：测试这两种配置

常见陷阱

陷阱 1：仅返回本地分片

症状：输出张量大小为 (out_features / world_size)，而不是 (out_features)

原因：实现计算了本地分片，但没有执行 all-gather

解决方法：实现集合操作，将所有 rank 的结果合并

陷阱 2：RowParallelLinear 中偏置处理不正确

症状：输出值比预期大 N 倍（其中 N 是 world_size）

原因：每个 rank 都加上了完整偏置，然后将这些值求和

解决方法：仅在 all-reduce 之后添加一次偏置，而不是在每个 rank 上都添加

陷阱 3：误解“模拟”要求

症状：实现适用于 world_size=1，但在更大的 world_size 下失败

原因：假设外部框架会处理集合操作

解决方法：仔细阅读要求——“如同使用 all_gather”意味着要实现该操作

陷阱 4：文件写入被截断

症状：实现存在语法错误或缺少代码

原因：文件写入操作被截断

解决方法：写入后务必读回完整文件，以验证其完整性

陷阱 5：分片维度错误

症状：矩阵乘法期间出现形状不匹配错误

原因：沿错误的维度分片（混淆了行和列）

解决方法：ColumnParallel 沿输出特征分片（权重的 dim=0），RowParallel 沿输入特征分片（权重的 dim=1）

实现前检查清单

编写代码前，请确认理解以下内容：

• 需要哪种集合操作（all-gather 还是 all-reduce）
• 最终输出形状应是什么
• 模拟是否应实际执行集合操作，还是将其延后处理
• 此并行类型应如何处理偏置
• 边界情况如何处理（world_size=1、维度不可整除）

实现后检查清单

编写代码后：

• 读回完整的实现文件，确认没有被截断
• 验证所有 world_size 下的输出形状均符合预期维度
• 手动推演一个简单示例，验证其正确性
• 测试简单情况（world_size=1），确认结果与非并行基线一致
• 至少测试一种非简单情况（world_size=2 或 4）