张量模型并行

数据并行在不同设备间复制模型并处理不同的数据批次时，当模型本身过大，无法载入单个加速器（GPU）的内存时，这种方法就无能为力了。针对此类情况，我们需要模型并行，它将一个模型划分到多个设备上。张量模型并行（TMP）是一种模型并行方式，它将单个层或层内的特定操作拆分到多个设备上。

这种方法对于那些在特定层中参数 (parameter)量巨大的模型尤其适用，例如现代Transformer架构中常见的大型嵌入 (embedding)表或前馈网络（FFN）层。

基本原理：张量拆分

张量模型并行的基本原理是有序地将层的权重 (weight)张量（有时也包括激活张量）拆分到多个GPU上。然后，计算在每个GPU上部分地执行，接着通过通信步骤来同步或合并结果，最终得到与原始未拆分层相同的输出。

我们以一个标准线性层为例，其操作定义为 $Y = XA + b$，其中 $X$ 是输入激活，$A$ 是权重矩阵，$b$ 是偏置 (bias)，$Y$ 是输出激活。张量模型并行提供不同的策略来并行化此操作。

列并行

在列并行中，权重 (weight)矩阵 $A$ 按列拆分到 $N$ 个GPU上。设 $A = [A_1, A_2, ..., A_N]$，其中每个 $A_i$ 位于不同的GPU上。

1. 输入广播：输入 $X$ 广播（或已可用）到模型并行组中的所有GPU。
2. 并行计算：每个GPU $i$ 计算一个部分结果 $Y_i = X A_i$。
3. 收集/拼接：部分结果 $Y_i$ 沿列维度收集并拼接，形成最终输出 $Y = [Y_1, Y_2, ..., Y_N]$。偏置 (bias) $b$ 也可以按列拆分，$b = [b_1, b_2, ..., b_N]$，并在收集步骤前在每个GPU上独立添加。

列并行应用于线性层。输入 $X$ 是共享的，权重 $A$ 和偏置 $b$ 按列拆分，部分输出被拼接。

这种方法在并行计算后需要通信（收集或全收集操作），用于组装完整的输出张量 $Y$。

行并行

或者，我们可以按行拆分权重 (weight)矩阵 $A$：$A = \begin{bmatrix} A_1 \\ A_2 \\ \vdots \\ A_N \end{bmatrix}$。

1. 输入拆分/分散：输入 $X$ 沿其最后一个维度（特征维度）在GPU间拆分：$X = [X_1, X_2, ..., X_N]$。注意，这与列并行不同，在列并行中，每个GPU都需要完整的 $X$。
2. 并行计算：每个GPU $i$ 使用其部分的输入和权重计算一个部分结果：$X_i A_i$。
3. 全归约：部分结果使用全归约操作在所有GPU上求和，生成最终输出 $Y = \sum_i X_i A_i$。偏置 (bias) $b$ 通常只在一个GPU上（例如，rank 0）在全归约求和之后添加。

行并行应用于线性层。输入 $X$ 被拆分，权重 $A$ 按行拆分，计算部分结果，然后通过全归约求和。偏置 $b$ 在归约后添加。

行并行在矩阵乘法之后需要通信（一个全归约操作）。一个重要好处是，输出激活 $Y$ 在所有参与的GPU上都得到复制，这可能是后续层（例如层归一化 (normalization)或另一个行并行线性层）所需的输入格式。

Transformer中列并行与行并行的结合

Transformer模型通常结合使用这些技术。例如，在一个由两个线性层组成的标准前馈网络（FFN）块中：

1. 第一个线性层可能使用列并行。这将较大的中间维度拆分到多个GPU上。输出激活是分布的（分片的）。
2. 第二个线性层可能使用行并行。它将分片激活作为输入（如果拆分维度匹配则无需通信），并生成通过全归约求和的输出，使FFN块的最终输出在GPU之间复制，为Transformer的下一部分（如残差连接或层归一化 (normalization)）做好准备。

这种策略性组合有助于最大限度地减少通信开销，通过在FFN块内的两个线性层之间保持激活分片。

嵌入 (embedding)并行化

对于词汇量非常大的模型，嵌入表可能成为一个显著的内存瓶颈。在这里，张量模型并行可以通过将嵌入表按行（沿词汇维度）在GPU之间拆分来应用。当查找输入token ID的嵌入时：

1. 每个GPU接收完整的输入ID序列。
2. 每个GPU仅查找其所持有的词汇部分对应的嵌入。如果某个ID对应于另一个GPU持有的行，它会产生一个零向量 (vector)。
3. 全归约操作会汇总所有GPU上的部分嵌入向量。由于对于任何给定的token ID，只有一个GPU生成了实际的嵌入向量（其他GPU生成了零），因此总和有效地收集了正确的嵌入。

实现细节与通信

手动实现张量模型并行需要仔细处理张量分片、计算和同步，使用 torch.distributed 包中的原语：

• torch.distributed.broadcast：将张量从一个进程发送到所有其他进程。
• torch.distributed.scatter：将张量块分散到各个进程。
• torch.distributed.gather：将张量从所有进程收集到一个进程。
• torch.distributed.all_gather：将张量从所有进程收集到所有进程。
• torch.distributed.reduce_scatter：在进程间执行操作（如求和）并分散结果。
• torch.distributed.all_reduce：在进程间执行操作（如求和）并使结果在所有进程上可用。

通常会创建专用函数或包装器来封装这些操作，以适用于特定层类型（例如 ColumnParallelLinear、RowParallelLinear）。像NVIDIA的Megatron-LM这样的库率先使用了许多这些技术，并且其中一部分功能正在通过诸如 torch.distributed.tensor.parallel 等模块集成到PyTorch核心中，这些模块提供了更高级的API来简化这些实现。

考虑一个使用辅助函数的列并行线性层示例：

import torch
import torch.nn as nn
import torch.distributed as dist

# 假设这些辅助函数用于管理并行组和通信
from .parallel_utils import (
    get_tensor_model_parallel_group,
    get_tensor_model_parallel_rank,
    get_tensor_model_parallel_world_size,
    copy_to_tensor_model_parallel_region, # 处理输入广播/拆分
    gather_from_tensor_model_parallel_region # 处理输出收集/归约
)

class ColumnParallelLinear(nn.Module):
    def __init__(self, input_size, output_size, bias=True, **kwargs):
        super().__init__()
        world_size = get_tensor_model_parallel_world_size()
        # 确保 output_size 可以被 world_size 整除
        assert output_size % world_size == 0
        self.output_size_per_partition = output_size // world_size
        self.input_size = input_size

        # 权重矩阵沿输出维度（列）拆分
        self.weight = nn.Parameter(torch.empty(
            self.output_size_per_partition, self.input_size, **kwargs
        ))
        # 初始化权重...（例如，使用 init.kaiming_uniform_）

        if bias:
            # 偏置也沿输出维度拆分
            self.bias = nn.Parameter(torch.empty(
                self.output_size_per_partition, **kwargs
            ))
            # 初始化偏置...（例如，使用 init.zeros_）
        else:
            self.register_parameter('bias', None)

    def forward(self, input_):
        # 如果前一层输出已复制（例如 LayerNorm），
        # 输入可能需要广播或已可用。
        # 此函数处理必要的通信。
        parallel_input = copy_to_tensor_model_parallel_region(input_)

        # 执行局部矩阵乘法
        output_parallel = nn.functional.linear(parallel_input, self.weight, self.bias)

        # 从张量并行组中的所有GPU收集结果
        # 沿列维度拼接。
        output_ = gather_from_tensor_model_parallel_region(output_parallel)

        return output_

列并行线性层的实现草图。请注意权重 (weight)/偏置 (bias)的显式分片以及通信包装器的使用。

权衡与考虑

• 复杂性：手动实现张量模型并行需要对模型架构代码进行大量修改，这与通常作为包装器运行的DDP不同。调试并行逻辑可能具有挑战性。
• 通信开销：张量模型并行在单个层的前向和反向传播 (backpropagation)中引入了通信。全收集或全归约等操作可能成为瓶颈，特别是当GPU之间的互连带宽（例如NVLink与PCIe）受限时。
• 粒度：对于参数 (parameter)量大或激活值相对于执行的计算量大的层，它最有效。
• 组合性：张量模型并行通常与数据并行（DDP）和流水线并行结合使用，以训练真正大型的模型，形成混合并行策略。

总而言之，张量模型并行是一种不可或缺的技术，当单个模型层超出单个设备的内存限制时。通过在多个设备上拆分层内的权重 (weight)和计算，它使得训练比以往可能更大的模型成为可能，尽管代价是增加了实现复杂性和通信开销。它是当前扩展大型神经网络 (neural network)的基础组成部分。