张量并行性的工作原理 - Amazon SageMaker
库如何调整张量并行性以适应 PyTorch nn.Linear 模块
Amazon Web Services 文档中描述的 Amazon Web Services 服务或功能可能因区域而异。要查看适用于中国区域的差异，请参阅 中国的 Amazon Web Services 服务入门 (PDF)。

张量并行性的工作原理

张量并行性在级别 nn.Modules 上实现；它在张量并行秩之间对模型中的特定模块分区。这是对管道并行性中使用的模块集的现有分区的补充。

通过张量并行性对模块进行分区时，其向前和向后传播是分布式的。库处理设备间的必要通信，以实施这些模块的分布式执行。这些模块在多个数据并行秩之间进行分区。与传统的工作负载分布相反，在使用库的张量并行性时，每个数据并行秩并没有完整的模型副本。相反，在所有未分布而保持完整的模块之外，每个数据并行秩可能只有分布式模块的一个分区。

示例：考虑跨数据并行秩的张量并行性，其中数据并行度为 4，张量并行度为 2。假设在对模块集进行分区后，您有一个存有以下模块树的数据并行组。

A
├── B
|   ├── E
|   ├── F
├── C
└── D
    ├── G
    └── H

假设模块 B、G 和 H 支持张量并行性。此模型的张量并行分区的一种可能结果可能是：

dp_rank 0 (tensor parallel rank 0): A, B:0, C, D, G:0, H
dp_rank 1 (tensor parallel rank 1): A, B:1, C, D, G:1, H
dp_rank 2 (tensor parallel rank 0): A, B:0, C, D, G:0, H
dp_rank 3 (tensor parallel rank 1): A, B:1, C, D, G:1, H

每行代表存储在该 dp_rank 中的一组模块，标记 X:y 表示模块 X 的第 y 个部分。请注意以下几点：

  1. 分区跨数据并行秩的子集进行，我们调用 TP_GROUP，而不是整个 DP_GROUP，因此精确的模型分区会跨 dp_rank 0 和 dp_rank 2 复制，类似地跨 dp_rank 1 和 dp_rank 3 复制。

  2. 模块 E 和 F 不再是模型的一部分，因为它们的父模块 B 已分区，并且通常属于 E 和 F 的任何执行在（已分区）B 模块内进行。

  3. 尽管 H 支持张量并行性，但在此示例中，它没有进行分区，这强调了是否对模块进行分区取决于用户输入。模块支持张量并行性并不一定意味着它需要分区。

库如何调整张量并行性以适应 PyTorch nn.Linear 模块

对数据并行秩执行张量并行性时，对于所分区的模块，参数、梯度和优化器状态的子集在张量并行设备上分区。对于其余模块，张量并行设备以常规的数据并行方式操作。为了执行分区模块，设备首先在同一个张量并行度组中，跨对等设备收集所有数据样本的必要部分。接下来，设备对所有这些数据样本运行模块的本地部分，然后进行另一轮同步，此过程既合并了每个数据样本的输出部分，又将组合的数据样本返回到最初生成数据样本的 GPU。下图显示了在已分区 nn.Linear 模块上执行此过程的示例。

第一张图显示了一个带有 nn.Linear 模块的小模型，带有在两个张量并行性秩上的数据并行性。nn.Linear 模块复制到两个并行秩中。

第二张图显示了在拆分 nn.Linear 模块时应用于较大模型的张量并行性。每个 tp_rank 存有一半的线性模块，以及剩余的全部操作。当线性模块运行时，每个 tp_rank 都会收集所有数据样本中相关的一半，并传递到自己的一半 nn.Linear 模块中。结果必须是减少分散的（使用求和作为归约运算），这样每个秩都有自己数据样本的最终线性输出。模型的其余部分以典型的数据并行方式运行。