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

张量并行（TP）建立在 PyTorch DistributedTensor（DTensor）之上，并提供不同的并行风格：按列并行、按行并行和序列并行。

警告

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

使用Tensor Parallelism并行化您的nn.Module的入口点是：

torch.distributed.tensor.parallel.parallelize_module(module, device_mesh, parallelize_plan)[源代码]¶

在 PyTorch 中应用张量并行，通过根据用户指定的计划并行化模块或子模块。

我们基于一个 parallelize_plan 来并行化模块或子模块。parallelize_plan 包含 ParallelStyle，它指示用户希望如何并行化模块或子模块。

用户还可以为每个模块的完全限定名称（FQN）指定不同的并行样式。

请注意，parallelize_module 仅接受一维的 DeviceMesh，如果你有一个二维或N维的 DeviceMesh，请先将其切片为一维的子 DeviceMesh，然后再传递给此 API（即 device_mesh["tp"]）

Parameters

模块 (nn.Module) – 要并行化的模块。
device_mesh (DeviceMesh) – 描述DTensor设备网格拓扑的对象。
parallelize_plan (Union[ParallelStyle, Dict[str, ParallelStyle]]) – 用于并行化模块的计划。它可以是一个包含如何准备输入/输出以进行张量并行化的ParallelStyle对象，也可以是一个模块FQN及其对应ParallelStyle对象的字典。

Returns

一个并行化的 nn.Module 对象。

Return type

模块

Example::

>>> from torch.distributed.tensor.parallel import parallelize_module, ColwiseParallel
>>> from torch.distributed.device_mesh import init_device_mesh
>>>
>>> # 定义模块。
>>> m = Model(...)
>>> tp_mesh = init_device_mesh("cuda", (8,))
>>> m = parallelize_module(m, tp_mesh, {"w1": ColwiseParallel(), "w2": RowwiseParallel()})
>>>

注意

对于像Attention、MLP层这样的复杂模块架构，我们建议将不同的ParallelStyles组合在一起（即ColwiseParallel和RowwiseParallel），并将其作为parallelize_plan传递，以实现所需的分布式计算。

张量并行支持以下并行方式：

class torch.distributed.tensor.parallel.ColwiseParallel(*, input_layouts=None, output_layouts=None, use_local_output=True)[源代码]¶

以列方式对兼容的 nn.Module 进行分区。目前支持 nn.Linear 和 nn.Embedding。用户可以将其与 RowwiseParallel 组合在一起，以实现更复杂模块的分片。（即 MLP、Attention）

Keyword Arguments

input_layouts (布局, 可选) – nn.Module 输入张量的 DTensor 布局，用于注释输入张量以成为 DTensor。如果未指定，我们假设输入张量是复制的。
output_layouts (布局, 可选) – nn.Module 输出的 DTensor 布局，用于确保 nn.Module 的输出符合用户期望的布局。如果未指定，输出张量将在最后一个维度上进行分片。
use_local_output (布尔值, 可选) – 是否使用本地的 torch.Tensor 而不是 DTensor 作为模块输出，默认值：True。

Returns

一个表示 nn.Module 列式分片的 ParallelStyle 对象。

Example::

>>> from torch.distributed.tensor.parallel import parallelize_module, ColwiseParallel
>>> from torch.distributed.device_mesh import init_device_mesh
>>> ...
>>> m = Model(...)  # m 是一个包含 "w1" nn.Linear 子模块的 nn.Module
>>> tp_mesh = init_device_mesh("cuda", (8,))
>>>
>>> # 默认情况下，"w1" Linear 的输入将被转换为 Replicated DTensor
>>> # 并且 "w1" 的输出将返回在最后一个维度上分片的 :class:`torch.Tensor`。
>>>
>>> sharded_mod = parallelize_module(m, tp_mesh, {"w1": ColwiseParallel()})
>>> ...

注意

默认情况下，如果未指定output_layouts，ColwiseParallel的输出会在最后一个维度上进行分片。如果存在需要特定张量形状的操作符（例如在配对的RowwiseParallel之前），请记住，如果输出被分片，操作符可能需要调整为分片的大小。

class torch.distributed.tensor.parallel.RowwiseParallel(*, input_layouts=None, output_layouts=None, use_local_output=True)[源代码]¶

以行方式对兼容的 nn.Module 进行分区。目前支持 nn.Linear 和 nn.Embedding。用户可以将其与 ColwiseParallel 组合以实现更复杂模块的分片。（即 MLP、Attention）

Keyword Arguments

input_layouts (布局, 可选) – nn.Module 输入张量的 DTensor 布局，用于注释输入张量以成为 DTensor。如果未指定，我们假设输入张量在最后一个维度上分片。
output_layouts (布局, 可选) – nn.Module 输出的 DTensor 布局，用于确保 nn.Module 的输出符合用户期望的布局。如果未指定，输出张量将被复制。
use_local_output (布尔值, 可选) – 是否使用本地的 torch.Tensor 而不是 DTensor 作为模块输出，默认值：True。

Returns

一个表示 nn.Module 行分片的 ParallelStyle 对象。

Example::

>>> from torch.distributed.tensor.parallel import parallelize_module, RowwiseParallel
>>> from torch.distributed.device_mesh import init_device_mesh
>>> ...
>>> m = Model(...)  # m 是一个包含 "w2" nn.Linear 子模块的 nn.Module
>>> tp_mesh = init_device_mesh("cuda", (8,))
>>>
>>> # 默认情况下，"w2" Linear 的输入将被转换为在最后一个维度上分片的 DTensor
>>> # 并且 "w2" 的输出将返回一个复制的 :class:`torch.Tensor`。
>>>
>>> sharded_mod = parallelize_module(m, tp_mesh, {"w2": RowwiseParallel()}),
>>> ...

class torch.distributed.tensor.parallel.SequenceParallel(*, sequence_dim=1, use_local_output=False)[源代码]¶

SequenceParallel 复制一个兼容的 nn.Module 参数，并在序列维度上对输入进行分片计算。目前支持 nn.LayerNorm、nn.Dropout 以及 RMSNorm python 实现

这种风格实现了论文中描述的操作减少大型Transformer模型中的激活重计算

输入和输出的 nn.Module 将在序列维度上进行分片。

Keyword Arguments

sequence_dim (int, 可选) – 输入张量的序列维度，用于注释输入张量以成为一个在序列维度上分片的DTensor，默认值：1。
use_local_output (布尔值, 可选) – 是否使用本地的 torch.Tensor 而不是 DTensor 作为模块输出，默认值：False。

Returns

一个表示 nn.Module 序列并行的 ParallelStyle 对象。

Example::

>>> from torch.distributed.tensor.parallel import parallelize_module, SequenceParallel
>>> from torch.distributed.device_mesh import init_device_mesh
>>> ...
>>> m = Model(...)  # m 是一个包含 "norm" nn.LayerNorm 子模块的 nn.Module
>>> tp_mesh = init_device_mesh("cuda", (8,))
>>>
>>> # 默认情况下，"norm" 的输入将被转换为在序列维度上分片的 DTensor
>>> # 并且 "norm" 的输出将返回一个在序列维度上分片的 :class:`DTensor`。
>>>
>>> sharded_mod = parallelize_module(m, tp_mesh, {"norm": SequenceParallel()}),
>>> ...

注意

SequenceParallel 风格假设如果 nn.Module 中有权重（例如 nn.LayerNorm 或 RMSNorm，并且它们默认具有 ones 初始化）。如果你对这些模块的权重有自定义初始化，你需要在并行化之前/之后广播权重，以确保它们被复制。

要简单地使用DTensor布局配置nn.Module的输入和输出，并执行必要的布局重分布，而不将模块参数分发到DTensors，可以在调用parallelize_module时使用以下ParallelStyle：

class torch.distributed.tensor.parallel.PrepareModuleInput(*, input_layouts, desired_input_layouts, use_local_output=False)[源代码]¶

配置 nn.Module 的输入，以根据 input_layouts 在运行时将 nn.Module 的输入张量转换为 DTensor，并根据 desired_input_layouts 执行布局重分布。

Keyword Arguments

input_layouts (Union[Placement, Tuple[Placement]]) – nn.Module 输入张量的 DTensor 布局，用于将输入张量转换为 DTensor。如果某些输入不是 torch.Tensor 或不需要转换为 DTensor，则需要指定 None 作为占位符。
desired_input_layouts (Union[Placement, Tuple[Placement]]) – 期望的 nn.Module 输入张量的 DTensor 布局，用于确保 nn.Module 的输入具有期望的 DTensor 布局。此参数需要与 input_layouts 的长度相同。
use_local_output (布尔值, 可选) – 是否使用本地的 torch.Tensor 而不是 DTensor 作为模块输入，默认值：False。

Returns

一个 ParallelStyle 对象，用于准备 nn.Module 输入的分片布局。

Example::

>>> from torch.distributed.tensor.parallel import parallelize_module, PrepareModuleInput
>>> from torch.distributed.device_mesh import init_device_mesh
>>> ...
>>> block = TransformerBlock(...)  # block 是一个包含 "attn" Attention 子模块的 nn.Module
>>> tp_mesh = init_device_mesh("cuda", (8,))
>>>
>>> # 根据下面指定的样式，attn 的第一个输入将被注释为 Sharded DTensor
>>> # 然后重新分配为 Replicated DTensor。
>>> parallelize_module(
>>>     block, # 这可以是子模块或模块
>>>     tp_mesh,
>>>     parallelize_plan={
>>>         "attn": PrepareModuleInput(
>>>             input_layouts=(Shard(0), None, None, ...),
>>>             desired_input_layouts=(Replicate(), None, None, ...)
>>>         ),
>>>     }
>>> )

class torch.distributed.tensor.parallel.PrepareModuleOutput(*, output_layouts, desired_output_layouts, use_local_output=True)[源代码]¶

配置 nn.Module 的输出，以便在运行时根据 output_layouts 将 nn.Module 的输出张量转换为 DTensor，并根据 desired_output_layouts 执行布局重分布。

Keyword Arguments

output_layouts (Union[Placement, Tuple[Placement]]) – nn.Module 输出张量的 DTensor 布局，用于将输出张量转换为 DTensor（如果它们是 torch.Tensor）。如果某些输出不是 torch.Tensor 或不需要转换为 DTensor，则需要指定 None 作为占位符。
desired_output_layouts (Union[Placement, Tuple[Placement]]) – 期望的 nn.Module 输出张量的 DTensor 布局，用于确保 nn.Module 的输出具有期望的 DTensor 布局。
use_local_output (布尔值, 可选) – 是否使用本地的 torch.Tensor 而不是 DTensor 作为模块的输出，默认值：True。

Returns

一个ParallelStyle对象，用于准备nn.Module输出的分片布局。

Example::

>>> from torch.distributed.tensor.parallel import parallelize_module, PrepareModuleOutput
>>> from torch.distributed.device_mesh import init_device_mesh
>>> ...
>>> block = TransformerBlock(...)  # block 是一个包含 "attn" Attention 子模块的 nn.Module
>>> tp_mesh = init_device_mesh("cuda", (8,))
>>>
>>> # 根据下面指定的样式，TransformerBlock 的输出将被转换为 Replicated DTensor
>>> # 然后重新分配到 Sharded DTensor。
>>> parallelize_module(
>>>     block, # 这可以是子模块或模块
>>>     tp_mesh,
>>>     parallelize_plan = PrepareModuleOutput(
>>>         output_layouts=Replicate(),
>>>         desired_output_layouts=Shard(0)
>>>     )
>>> )

注意

当使用 Shard(dim) 作为上述 ParallelStyle 的输入/输出布局时，我们假设输入/输出激活张量在 DeviceMesh 上沿张量维度 dim 均匀分片。例如，由于 RowwiseParallel 接受在最后一个维度上分片的输入，它假设输入张量已经在最后一个维度上均匀分片。对于不均匀分片的激活张量的情况，可以直接将 DTensor 传递给分区模块，并使用 use_local_output=False 在每个 ParallelStyle 之后返回 DTensor，其中 DTensor 可以跟踪不均匀的分片信息。

对于像Transformer这样的模型，我们建议用户在parallelize_plan中同时使用ColwiseParallel 和RowwiseParallel，以实现整个模型的所需分片（即Attention和MLP）。

并行化的交叉熵损失计算（损失并行化），通过以下上下文管理器实现支持：

torch.distributed.tensor.parallel.loss_parallel()[源代码]¶

一个上下文管理器，启用损失并行化，当输入在类别维度上分片时，可以执行高效的并行化损失计算。目前仅支持交叉熵损失。

在这个上下文管理器中，可以像往常一样使用 cross_entropy() 或 CrossEntropyLoss，并对输入参数做以下假设。相应的 backward() 调用（如果有）也需要在这个上下文管理器下进行。

Parameters

输入 (DTensor) – 输入的logits。假设在类别维度上进行分片。
目标 (Union[torch.Tensor, DTensor]) – 必须是真实类别索引（当前不支持类别概率）。假设在 DeviceMesh 中被复制。
权重（Union[torch.Tensor, DTensor]，可选）– 如果给出，假设在DeviceMesh中被复制。
label_smoothing – 目前不支持。

Returns

一个复制的 DTensor。

示例

这里手动创建了一个分片的DTensor，以展示其用法。在实际操作中，它通常是TP模块的输出。

>>> from torch.distributed.tensor.parallel import loss_parallel
>>> from torch.distributed.device_mesh import init_device_mesh
>>> ...
>>> device_mesh = init_device_mesh("cuda", (8,))
>>> input = torch.randn(4, 16, device="cuda", requires_grad=True)
>>> dist_input = distribute_tensor(input, device_mesh, placements=[Shard(1)])
>>> target = torch.randint(16, (4,), device="cuda")
>>> with loss_parallel():
>>>     loss = F.cross_entropy(dist_input, target, reduction="mean")
>>>     loss.backward()
>>> ...

警告

loss_parallel API 是实验性的，可能会发生变化。