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

张量并行（TP）基于PyTorch分布式张量（DTensor）构建，并提供了不同的并行风格：列式并行、行式并行和序列并行。

警告

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

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

torch.distributed.tensor.parallel.parallelize_module(module, device_mesh, parallelize_plan)[source]¶

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

我们根据并行计划对模块或子模块进行并行化。该并行计划包含 ParallelStyle，这表示用户希望如何对该模块或子模块进行并行化。

用户还可以为每个模块的完整限定名（FQN）指定不同的并行风格。

注意，parallelize_module 只接受一维的 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
>>>
>>> # Define the module.
>>> m = Model(...)
>>> tp_mesh = init_device_mesh("cuda", (8,))
>>> m = parallelize_module(m, tp_mesh, {"w1": ColwiseParallel(), "w2": RowwiseParallel()})
>>>

注意

对于像注意力机制和多层感知机这样的复杂模块架构，我们建议组合不同的并行样式（即 ColwiseParallel 和 RowwiseParallel），并将它们作为并行计划传递，以实现所需的分片计算。

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

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

按列方式拆分一个兼容的 nn.Module。目前支持 nn.Linear 和 nn.Embedding。用户可以将其与 RowwiseParallel 组合使用，以实现更复杂模块（例如 MLP、Attention）的拆分。

Keyword Arguments

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

Returns

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

Example::

>>> from torch.distributed.tensor.parallel import parallelize_module, ColwiseParallel
>>> from torch.distributed.device_mesh import init_device_mesh
>>> ...
>>> m = Model(...)  # m is a nn.Module that contains a "w1" nn.Linear submodule
>>> tp_mesh = init_device_mesh("cuda", (8,))
>>>
>>> # By default, the input of the "w1" Linear will be converted to Replicated DTensor
>>> # and the output of "w1" will return :class:`torch.Tensor` that shards on the last dim.
>>>
>>> sharded_mod = parallelize_module(m, tp_mesh, {"w1": ColwiseParallel()})
>>> ...

注意

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

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

按照行的方式拆分一个兼容的 nn.Module。目前支持 nn.Linear 和 nn.Embedding。用户可以将其与 ColwiseParallel 组合以实现更复杂模块的拆分（例如 MLP、Attention）。

Keyword Arguments

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

Returns

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

Example::

>>> from torch.distributed.tensor.parallel import parallelize_module, RowwiseParallel
>>> from torch.distributed.device_mesh import init_device_mesh
>>> ...
>>> m = Model(...)  # m is a nn.Module that contains a "w2" nn.Linear submodule
>>> tp_mesh = init_device_mesh("cuda", (8,))
>>>
>>> # By default, the input of the "w2" Linear will be converted to DTensor that shards on the last dim
>>> # and the output of "w2" will return a replicated :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)[source]¶

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

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

如果传入到此 nn.Module 的输入是一个 torch.Tensor，它假设输入已经在序列维度上进行了分片，并将其转换为在序列维度上进行分片的 DTensor。如果传入到此 nn.Module 的输入已经是 DTensor，但未在序列维度上进行分片，则会重新分布输入以使其在序列维度上进行分片。

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

Keyword Arguments

sequence_dim (整数, 可选) – 输入张量在 nn.Module 的序列维度，这用于将输入张量标注为在序列维度上分片的DTensor，默认值为 1。
use_local_output (bool, 可选) – 是否使用本地 torch.Tensor 而不是 DTensor 作为模块输出，默认值为 False。

Returns

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

Example::

>>> from torch.distributed.tensor.parallel import parallelize_module, SequenceParallel
>>> from torch.distributed.device_mesh import init_device_mesh
>>> ...
>>> m = Model(...)  # m is a nn.Module that contains a "norm" nn.LayerNorm submodule
>>> tp_mesh = init_device_mesh("cuda", (8,))
>>>
>>> # By default, the input of the "norm" will be converted to DTensor that shards on the sequence dim
>>> # and the output of "norm" will return a sharded on sequence dimension :class:`DTensor`.
>>>
>>> sharded_mod = parallelize_module(m, tp_mesh, {"norm": SequenceParallel()}),
>>> ...

注意

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

要简单配置 nn.Module 的输入和输出与 DTensor 布局，并执行必要的布局重新分配，而无需将模块参数分布到 DTensor，可以在调用 parallelize_module 时使用以下 ParallelStyle 个步骤：

class torch.distributed.tensor.parallel.PrepareModuleInput(*, input_layouts=None, desired_input_layouts=None, input_kwarg_layouts=None, desired_input_kwarg_layouts=None, use_local_output=False)[source]¶

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

Keyword Arguments

输入布局 (Union[Placement, Tuple[Optional[Placement]]]) – 输入张量的DTensor布局，用于将输入张量转换为DTensors。如果某些输入不是torch.Tensor或者不需要转换为DTensors，则需要指定None作为占位符。默认值：None。
期望的输入布局 (Union[Placement, Tuple[Optional[Placement]]]) – nn.Module 的输入张量的期望 DTensor 布局，这用于确保 nn.Module 的输入具有所需的 DTensor 布局。此参数需要与 input_layouts 的长度相同。默认值：None。
input_kwarg_layouts (Dict[str, Placement]) – nn.Module 输入关键字参数的 DTensor 布局，用于将输入关键字张量转换为 DTensors。默认值: None
desired_input_kwarg_layouts – (Dict[str, Placement])： nn.Module 输入 kwargs 的期望 DTensor 布局，用于确保 nn.Module 的输入具有所需的 DTensor 布局。默认值：None。
use_local_output (bool, 可选) – 是否使用本地 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 is a nn.Module that contains an "attn" Attention submodule
>>> tp_mesh = init_device_mesh("cuda", (8,))
>>>
>>> # According to the style specified below, the first input of attn will be annotated to Sharded DTensor
>>> # and then redistributed to Replicated DTensor.
>>> parallelize_module(
>>>     block, # this can be a submodule or module
>>>     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)[source]¶

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

Keyword Arguments

输出布局 (Union[Placement, Tuple[Placement]]) – nn.Module 的输出张量的 DTensor 布局，这用于将输出张量转换为 DTensors，如果它们是 torch.Tensor。如果某些输出不是 torch.Tensor 或者不需要转换为 DTensors，则需要指定 None 作为占位符。
期望的输出布局 (Union[Placement, Tuple[Placement]]) – nn.Module 输出张量的期望 DTensor 布局，这用于确保 nn.Module 的输出具有所需的 DTensor 布局。
use_local_output (bool, 可选) – 是否使用本地 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 is a nn.Module that contains an "attn" Attention submodule
>>> tp_mesh = init_device_mesh("cuda", (8,))
>>>
>>> # According to the style specified below, the output of the TransformerBlock will be converted to Replicated DTensor
>>> # and then redistributed to Sharded DTensor.
>>> parallelize_module(
>>>     block, # this can be a submodule or module
>>>     tp_mesh,
>>>     parallelize_plan = PrepareModuleOutput(
>>>         output_layouts=Replicate(),
>>>         desired_output_layouts=Shard(0)
>>>     )
>>> )

注意

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

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

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

torch.distributed.tensor.parallel.loss_parallel()[source]¶

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

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

Parameters

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

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 是实验性的，可能会发生变化。

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

文档

教程

资源