张量并行性 - torch.distributed.tensor.parallel

张量并行（TP）是建立在分布式张量（DTensor）之上的，并提供了几种并行方式：行方向并行、列方向并行和成对并行。

警告

张量并行 API 是实验性的，可能会发生变化。

使用张量并行化的并行化您的 nn.Module 的入口是：

torch.distributed.tensor.parallel.parallelize_module(module, device_mesh, parallelize_plan, tp_mesh_dim=0)

在PyTorch中应用张量并行（TP）的API。我们根据并行计划（parallelize_plan）对模块或子模块进行并行化。并行计划包含 ParallelStyle，这表示用户希望如何对模块或子模块进行并行化。

用户还可以根据模块的完整限定名称（FQN）指定不同的并行风格。 该 API 本机支持二维并行，通过接受一个 n 维 device_mesh，并且用户只需指定执行张量并行的维度即可。

Parameters:

• 模块 (nn.Module) – 需要并行化的模块。

• device_mesh (DeviceMesh) – 描述用于 DTensor 的设备网格拓扑结构的对象。

• parallelize_plan (Union[ParallelStyle, Dict[str, ParallelStyle]]) – 用于并行化模块的计划。它可以是 一个ParallelStyle对象，其中包含我们如何为张量并行准备输入/输出，或者它也可以是一个字典，键为模块的完全限定名称（FQN），值为其对应的ParallelStyle对象。

• tp_mesh_dim (int) – 我们在维度 device_mesh 上执行张量并行的维度。

Returns:

一个 nn.Module 对象并行化。

Return type:

模块

Example::

>>> from torch.distributed._tensor.parallel import parallelize_module, PairwiseParallel
>>>
>>> # Define the module.
>>> m = Model(...)
>>> m = parallelize_module(m, PairwiseParallel())
>>>

警告

PairwiseParallel 目前带有约束条件。如果你需要更细粒度的控制，需要传入一个包含模块全限定名和并行风格的字典代替。

张量并行支持以下并行风格：

class torch.distributed.tensor.parallel.style.RowwiseParallel

将模块的行进行划分。 我们假设输入为分片的 DTensor，输出为复制的 DTensor。

class torch.distributed.tensor.parallel.style.ColwiseParallel

将张量或模块的列进行分区。 我们假设输入是一个复制的 DTensor，输出是一个分片的 DTensor。

class torch.distributed.tensor.parallel.style.PairwiseParallel

PairwiseParallel 将列方向和行方向的样式作为固定对进行拼接，就像 Megatron-LM(https://arxiv.org/abs/1909.08053) 所做的那样。 我们假设输入和输出都需要是 DTensor 的副本。

警告

PairwiseParallel 目前仅支持 nn.Multihead Attention, nn.Transformer 或偶数层的 MLP。

由于张量并行性是建立在 DTensor 之上的，我们需要使用 DTensor 指定模块的输入和输出位置，以便它能够与前后的模块预期地进行交互。以下是一些用于输入/输出准备的函数：

torch.distributed.tensor.parallel.style.make_input_replicate_1d(input, device_mesh=None)

在1-D设备网格上复制输入张量。此函数将在ParallelStyle中使用。

Parameters:

• 输入 (Union[torch.Tensor, DTensor]) – 此输入张量将在1-D DeviceMesh 上被复制。

• device_mesh (DeviceMesh, 可选) – 用于放置 input 的 1-D 设备网格。 如果未传入 DeviceMesh 并且 input 是一个 DTensor， 将使用 input.device_mesh。 如果 DeviceMesh 不是 1-D，将抛出异常。 默认值: None

Returns:

A DTensor 复制到 device_mesh。

Return type:

DTensor

torch.distributed.tensor.parallel.style.make_input_shard_1d(input, device_mesh=None, dim=0)

在1-D设备网格上将输入张量分片到dim。此函数将在ParallelStyle中使用。

Parameters:

• 输入 (Union[torch.Tensor, DTensor]) – 单个张量将在维度 dim 上沿 1-D DeviceMesh 进行分片。

• device_mesh (DeviceMesh, 可选) – 用于将 input 进行分片的 1-D 设备网格。 如果未传递 DeviceMesh 并且 input 是一个 DTensor， 将使用 input.device_mesh。 如果 DeviceMesh 不是 1-D，将抛出异常。 默认值: None

• dim (int, 可选) – input 张量的分片维度。 默认值: 0

Returns:

A DTensor 在维度 dim 上分片为 device_mesh。

Return type:

DTensor

torch.distributed.tensor.parallel.style.make_input_shard_1d_last_dim(input, device_mesh=None)

Wrapper func of make_input_shard_1d with dim = -1.

Parameters:

• 输入 (Union[torch.Tensor, DTensor]) – 这个单一张量将在维度 dim 上沿着 1-D DeviceMesh 进行分片。

• device_mesh (DeviceMesh, 可选) – 用于将 input 进行分片的 1-D 设备网格。 如果未传递 DeviceMesh 并且 input 是一个 DTensor， 将使用 input.device_mesh。 如果 DeviceMesh 不是 1-D，将抛出异常。 默认值: None

Returns:

A DTensor 在维度 dim 上分片为 device_mesh。

Return type:

DTensor

torch.distributed.tensor.parallel.style.make_output_replicate_1d(output, device_mesh=None)

将输出 DTensor 转换为复制的 DTensor。这将在并行风格中使用。

Parameters:

• 输出 (DTensor) – 要转换的模块的输出。

• device_mesh (DeviceMesh, 可选) – 用于复制输出所需的对象，它必须是一个1D device_mesh 如果传入的不是1D device_mesh，我们将抛出异常。 如果没有传入 device_mesh，我们将复用输出中的那个。 默认值: None

Returns:

一个 DTensor 对象被复制。

Return type:

DTensor

torch.distributed.tensor.parallel.style.make_output_tensor(output, device_mesh=None)

首先将输出的 DTensor 转换为复制的 DTensor，然后再将其转换为 Tensor。

Parameters:

• 输出 (DTensor) – 要转换的模块的输出。

• device_mesh (DeviceMesh, 可选) – 用于复制输出所需的对象，它必须是一个一维的 device_mesh，如果传入非一维的 device_mesh，我们将抛出异常。如果没有传入 device_mesh， 我们将复用输出中的那个。默认值： None

Returns:

一个从输出DTensor转换而来的 torch.Tensor 对象。

Return type:

张量

torch.distributed.tensor.parallel.style.make_output_shard_1d(output, device_mesh=None, dim=0)

将输出 DTensor 转换为分片的 DTensor。这将在 ParallelStyle 中使用。

Parameters:

• 输出 (DTensor) – 要转换的模块的输出。

• device_mesh (DeviceMesh, 可选) – 用于分片输出所需的对象，它必须是一个 1D device_mesh， 如果传入了非 1D device_mesh，我们将抛出异常。 如果没有传入 device_mesh，我们将复用输出中的那个。 默认值: None

• dim (int) – 输出的分片维度。默认值：0

Returns:

A DTensor 对象在给定维度上进行分片。

Return type:

DTensor

目前，有一些限制使得 nn.MultiheadAttention 模块难以开箱即用的方式支持张量并行，所以我们为张量并行用户构建了这个 multihead_attention 模块。此外，在 parallelize_module 中，当我们指定 PairwiseParallel 时，我们会自动将 nn.MultiheadAttention 替换为这个自定义模块。

class torch.distributed.tensor.parallel.multihead_attention_tp.TensorParallelMultiheadAttention(embed_dim, num_heads, dropout=0.0, bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None, batch_first=False, device=None, dtype=None, tp_size=1, self_attention=True)

来自Transformer模型的多头注意力块。 由于我们需要对注意力层进行一些自定义， 我们编写了一个自定义但数学上等效的 注意力模块，如torch.nn中所定义。

请注意： 我们现在仅支持自注意力的情况，且输入参数有限，并假设输入张量的维度为三。虽然我们实现了多头注意力的逻辑，但尚未进行完整测试。

我们还启用了2D并行性以与FullyShardedDataParallel集成。 用户只需显式调用以下API:

torch.distributed.tensor.parallel.fsdp.enable_2d_with_fsdp()

该 API 注册了 Tensor Parallelism（TP）与 FullyShardedDataParallel（FSDP）协同工作所需的扩展。我们首先根据 parallelize_plan 在一个模块或子模块内对参数进行并行化，然后让 FSDP 对分布式参数的本地张量进行重新分片，该张量本质上是一个 DTensor。

Returns:

一个 bool 表示扩展注册是否成功。

Return type:

布尔