目录
目录

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

Tensor Parallelism（TP）构建在 PyTorch DistributedTensor 之上 （DTensor） 并提供了几种并行样式：Rowwise、Colwise 和 Pairwise Parallelism。

警告

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

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

torch.distributed.tensor.parallel 的parallelize_module（module， device_mesh， parallelize_plan， tp_mesh_dim=0）[来源]

在 PyTorch 中应用张量并行 （TP） 的 API。我们并行化模块 或基于parallelize_plan sub_modules。parallelize_plan包含 ，它指示用户对模块或sub_module的需要 进行并行化。ParallelStyle

用户还可以为每个模块指定不同的并行样式 完全限定名称 （FQN）。 API 通过接受 n 维 device_mesh 来原生支持 2D 并行性 用户只需要指定我们执行张量并行的维度。

参数

  • module （） – 要并行化的模块。nn.Module

  • device_mesh （） – 描述网格拓扑的对象 的设备数量。DeviceMesh

  • parallelize_plan （Union[， Dict[str， ]]） – 用于并行化模块的计划。它可以是一个对象，其中包含 我们为 Tensor Parallelism 准备 input/output，或者它可以是一个 dict 及其对应的对象。ParallelStyleParallelStyleParallelStyleParallelStyle

  • tp_mesh_dim （int） – 我们执行位置的维度 Tensor Parallelism 打开。device_mesh

返回

并行化的对象。nn.Module

返回类型

模块

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

警告

PairwiseParallel现在带有约束。如果您需要更精细 granularity 时，您需要传入模块 FQN 和 parallel 样式的 dict。

Tensor Parallelism 支持以下并行样式：

类 torch.distributed.tensor.parallel.style 中。RowwiseParallel（_prepare_input=<函数make_input_shard_1d_last_dim>， _prepare_output=<函数make_output_tensor>）[来源]

对模块的行进行分区。 我们假设 input 是分片的，output 是 .DTensor

类 torch.distributed.tensor.parallel.style 中。ColwiseParallel（_prepare_input=<function make_input_replicate_1d>， _prepare_output=<function make_sharded_output_tensor>）[来源]

对张量或模块的列进行分区。 我们假设 input 是 replicad，output 是 sharded 。DTensor

类 torch.distributed.tensor.parallel.style 中。PairwiseParallel（_prepare_input=无，_prepare_output=无）[来源]

PairwiseParallel 将 colwise 和 rowwise 样式连接为固定样式 就像 Megatron-LM（https://arxiv.org/abs/1909.08053） 正在做的事情一样。 我们假设输入和输出都需要复制 DTensor。

警告

PairwiseParallel 目前不支持 。一种解决方法是将 和 应用于 变压器。我们建议仅使用 对于偶数层 MLP。nn.MultiheadAttentionnn.TransformerColwiseParallelRowwiseParallelPairwiseParallel

警告

序列并行仍处于实验阶段，尚未进行评估。

类 torch.distributed.tensor.parallel.style 中。SequenceParallel[来源]

SequenceParallel 将 colwise 和 rowwise 样式连接为固定的 与 Sequence Parallel 配对在一起，就像 Megatron-LM Sequence Parallel 一样 （https://arxiv.org/pdf/2205.05198.pdf） 正在执行。 我们假设输入和输出都需要分片 DTensor。

警告

SequenceParallel 目前不支持 。一种解决方法是将 和 应用于 变压器。我们建议仅使用 对于偶数层 MLP。nn.MultiheadAttentionnn.TransformerColwiseParallelRowwiseParallelSequenceParallel

由于 Tensor Parallelism 构建在 DTensor 之上，因此我们需要指定 使用 DTensor 放置模块的输入和输出位置，以便它可以预期地 与 Before， and After 的 Module 交互。以下是函数 用于输入/输出准备：

torch.distributed.tensor.parallel.style 的make_input_replicate_1d（输入，device_mesh=无）[来源]

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

参数

  • input （Union[， ]） – 此输入张量将在 1-D 上复制。DTensorDeviceMesh

  • device_mesh （ 可选 ） – 将复制的一维设备网格。 如果传递 no 并且 是 ，则将使用 。 如果不是 1-D，则将引发异常。 违约：DeviceMeshinputDeviceMeshinputDTensorinput.device_meshDeviceMeshNone

返回

复制的 .DTensordevice_mesh

返回类型

DTensor

torch.distributed.tensor.parallel.style 的make_input_reshard_replicate（输入，device_mesh）[来源]

从不同秩的张量构造分片 DTensor ，然后转换为复制的 DTensor。

参数

  • input （） - 每个 rank 的输入张量，由一个全局 DTensor 组成 在 1-D 上的维度上分片，然后将分片的 DTensor 转换为复制的 DTensor。0DeviceMesh

  • device_mesh （可选） – 将分片的一维设备网格。 如果不是 1-D，则将引发异常。 违约：DeviceMeshinputDeviceMeshNone

返回

一个 sharded on dimension overDTensor0device_mesh

，然后转换为 Replicate。

返回类型

DTensor

torch.distributed.tensor.parallel.style 的make_input_shard_1d（input， device_mesh=None， dim=0）[来源]

一维设备网格上的分片输入张量。此函数将在 ParallelStyle 中使用。dim

参数

  • input （Union[， ]） – 单个张量将在 1-D 上的维度上分片。DTensordimDeviceMesh

  • device_mesh （可选） – 将分片的一维设备网格。 如果传递 no 并且是 ，则将使用 input.device_mesh。 如果不是 1-D，则将引发异常。 违约：DeviceMeshinputDeviceMeshinputDTensorDeviceMeshNone

  • dim （int， optional） – tensor 的分片维度。 默认值：0input

返回

维度上 的分片 。DTensordimdevice_mesh

返回类型

DTensor

torch.distributed.tensor.parallel.style 的make_input_shard_1d_last_dim（输入，device_mesh=无）[来源]

包装器 func 与 = -1。make_input_shard_1ddim

参数

  • input （Union[， ]） – 此单个张量将在最后一个维度上分片 在 1-D 上。DTensorDeviceMesh

  • device_mesh （可选） – 将分片的一维设备网格。 如果传递 no 并且是 ，则将使用 input.device_mesh。 如果不是 1-D，则将引发异常。 违约：DeviceMeshinputDeviceMeshinputDTensorDeviceMeshNone

返回

在最后一个维度上分片的 。DTensordevice_mesh

返回类型

DTensor

torch.distributed.tensor.parallel.style 的make_output_replicate_1d（输出，device_mesh=无）[来源]

将 Output DTensor 转换为复制的 DTensor。这将在 ParallelStyle 中使用。

参数

  • output （） – 要转换的模块的输出。DTensor

  • device_mesh （ 可选 ） – 对象需要复制输出，并且它必须是 1D 的，如果传入非 1D ，我们将引发异常。 如果传入 no，我们将重用 output 中的那个。 违约：DeviceMeshdevice_meshdevice_meshdevice_meshNone

返回

已复制的对象。DTensor

返回类型

DTensor

torch.distributed.tensor.parallel.style 的make_output_reshard_tensor（输出，device_mesh=无）[来源]

将 Output DTensor 转换为分片的 DTensor 并返回本地张量。

参数

  • output （） – 要转换的模块的输出。DTensor

  • device_mesh （ 可选 ） – 对象需要对输出进行分片，并且它必须是 1D 的，如果传入非 1D ，我们将引发异常。 如果传入 no，我们将重用 output 中的那个。 违约：DeviceMeshdevice_meshdevice_meshdevice_meshNone

返回

从输出 DTensor 转换而来的对象。

返回类型

张肌

torch.distributed.tensor.parallel.style 的make_output_shard_1d（输出，device_mesh=无，dim=0）[来源]

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

参数

  • output （） – 要转换的模块的输出。DTensor

  • device_mesh （ 可选 ） – 对象需要对输出进行分片，并且它必须是 1D 的，如果传入非 1D ，我们将引发异常。 如果传入 no，我们将重用 output 中的那个。 违约：DeviceMeshdevice_meshdevice_meshdevice_meshNone

  • dim （int） – 为输出分片 dim。默认值：0

返回

在给定 dim 上分片的对象。DTensor

返回类型

DTensor

torch.distributed.tensor.parallel.style 的make_output_tensor（输出，device_mesh=无）[来源]

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

参数

  • output （） – 要转换的模块的输出。DTensor

  • device_mesh （ 可选 ） – 复制输出所需的对象 ，并且它需要 a 1D，如果传入了非 1D，我们将引发异常。如果传入 no，则 我们将重用 output 中的 one。 违约：DeviceMeshdevice_meshdevice_meshdevice_meshNone

返回

从输出 DTensor 转换而来的对象。

返回类型

张肌

目前，存在一些约束，这使得该模块很难开箱即用地处理 Tensor Parallelism，因此我们建议用户尝试 and for each parameter。现在可能需要进行一些代码更改 因为我们在 Module 的 head dim 上并行化。MultiheadAttentionColwiseParallelRowwiseParallelMultiheadAttention

我们还支持 2D 并行，我们将张量并行与数据并行相结合。 要与 集成，请执行以下操作 用户只需显式调用以下 API：FullyShardedDataParallel

torch.distributed.tensor.parallel.fsdp 的enable_2d_with_fsdp（）[来源]

API 注册了张量并行 （TP） 所需的扩展 以使用 FullyShardedDataParallel （FSDP）。我们首先并行化参数 在一个模块或parallelize_plan sub_modules中，并会让 FSDP rehard 分布式参数的局部张量，本质上是一个 DTensor。

返回

bool 指示扩展注册是否成功。

返回类型

布尔

要与 集成，请执行以下操作 用户只需显式调用以下 API：DistributedDataParallel

torch.distributed.tensor.parallel.ddp 的pre_dp_module_transform（模块）[来源]

实现 Tensor Parallelism （TP） 和 Data 之间的可组合性 使用 DDP 时 PyTorch 中的 Parallelism（DP）。我们需要转换 Parameters 是本地张量的 DTensor，然后再使用数据并行 API 进行包装。 然后我们注册两个钩子，一个用于将本地张量转换回 DTensor preforward 和 1 个用于将 DTensor 转换回 Forward 之后的张量。由 以这种方式集成，我们避免了 DDP 对 DTensor 参数的任何特殊处理 并将 DTensor 的梯度传播回 DP，例如 DDP 的梯度桶。

目前，此 API 仅适用于 .它稍后将支持 其他 DP 方法，如 FSDP。DistributedDataParallel

参数

module （） – 已应用 TP 的模块。nn.Module

例：：
>>> from torch.distributed.tensor.parallel import parallelize_module, PairwiseParallel
>>> from torch.nn.parallel import DistributedDataParallel as DDP
>>> from torch.distributed.tensor.parallel.ddp import pre_dp_module_transform
>>>
>>> # Define the module.
>>> m = module(...)
>>> parallelize_module(m, PairwiseParallel())
>>> m = pre_dp_module_transform(m)
>>> m = DDP(m)
>>>

文档

访问 PyTorch 的全面开发人员文档

查看文档

教程

获取面向初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并解答您的问题

查看资源