torch.distributed.tensor¶

注意

torch.distributed.tensor 目前处于开发中的 alpha 阶段，我们正在为文档中列出的大多数 API 保证向后兼容性，但如果有必要，可能会进行 API 更改。

PyTorch DTensor（分布式张量）¶

PyTorch DTensor 提供了简单且灵活的张量分片原语，透明地处理分布式逻辑，包括分片存储、跨设备/主机的操作计算和集体通信。 DTensor 可用于构建不同的并行解决方案，并在使用多维分片时支持分片状态字典表示。

请查看基于 DTensor 构建的 PyTorch 原生并行解决方案中的示例：

DTensor 遵循 SPMD（单程序，多数据）编程模型，使用户能够编写分布式程序，就像它是具有相同收敛特性的 单设备程序 一样。它通过指定 DeviceMesh 和 Placement 提供统一的张量分片布局（DTensor 布局）：

DeviceMesh 使用一个 n 维数组来表示设备拓扑和集群的通信器。
Placement 描述了逻辑张量在 DeviceMesh 上的分片布局。 DTensor 支持三种类型的放置：Shard、Replicate 和 Partial。

PyTorch深度张量类API¶

DTensor 是一个 torch.Tensor 的子类。这意味着一旦创建了一个 DTensor，它就可以像使用 torch.Tensor 一样被使用，包括在单个设备上运行各种类型的 PyTorch 操作符，允许为 PyTorch 操作符进行适当的分布式计算。

除了现有的 torch.Tensor 种方法外，它还提供了一组额外的方法来与 torch.Tensor、redistribute DTensor 布局交互，例如创建一个新的 DTensor、获取所有设备上的完整张量内容等。

class torch.distributed.tensor.DTensor(local_tensor, spec, *, requires_grad)¶

DTensor（分布式张量）是 torch.Tensor 的子类，它为使用多设备 torch.Tensor 提供了类似单设备的抽象。它通过 DeviceMesh 和以下类型的 Placement 描述分布式张量的切分布局（DTensor 布局）：

Shard: 张量在张量维度 dim 上分片，并分布在 DeviceMesh 维度的设备上
Replicate: 在 DeviceMesh 维度设备上复制的张量
Partial: 张量在 DeviceMesh 维度的设备上处于待规约状态

在调用 PyTorch 操作时，DTensor 会覆盖 PyTorch 操作以执行分片计算并在必要时发出通信。除了操作计算之外，DTensor 还会根据操作本身的语义正确地转换或传播放置（DTensor 布局），并生成新的 DTensor 输出。

为了确保调用 PyTorch 操作符时 DTensor 个分片计算的数值正确性，DTensor 要求操作符的每个张量参数都必须是 DTensor。

Return type: DTensor

property device_mesh: DeviceMesh¶: 该DTensor对象关联的DeviceMesh属性。

注意

device_mesh 是一个只读属性，不能被设置。

static from_local(local_tensor, device_mesh=None, placements=None, *, run_check=False, shape=None, stride=None)[source][source]¶

根据每个进程上本地的torch.Tensor，按照指定的device_mesh和placements创建一个DTensor。

Parameters

local_tensor (torch.Tensor) – 每个进程上的本地 torch.Tensor。
device_mesh (DeviceMesh, 可选) – 用于放置张量的 DeviceMesh，如果未指定，则必须在 DeviceMesh 上下文管理器中调用，默认值：None
placements (List[Placement]，可选) – 描述如何将本地的 torch.Tensor 放置在 DeviceMesh 上的放置方式，必须与 device_mesh.ndim 具有相同数量的元素。

Keyword Arguments

run_check (bool, 可选) – 以额外通信为代价，在各个进程之间执行健全性检查，以验证每个本地张量的元信息是否正确。如果在 placements 中有 Replicate，设备网格维度上的第一个进程的数据将广播到其他进程。默认值：False
shape (torch.Size, 可选) – 一个整数列表，用于指定构建在 local_tensor 之上的 DTensor 的大小。如果 local_tensor 在不同进程中的形状不同，则需要提供此参数。如果没有提供，则假设给定的分布式张量在各个进程中均匀切分，并计算 shape。默认值：None
stride (元组, 可选) – 一个整数列表，用于指定 DTensor 的步长。如果未提供，则假设给定的分布式张量在各个进程之间均匀切分，并计算 stride。默认值：None

Returns

一个 DTensor 对象

Return type

DTensor

注意

当 run_check=False 时，用户有责任确保在不同进程间传递的本地张量是正确的（即对于 Shard(dim) 的放置方式，张量是分片的；对于 Replicate() 的放置方式，张量是复制的）。如果不是这样，创建的 DTensor 的行为将是未定义的。

注意

from_local 是可微的，创建的 requires_grad 对象的 DTensor 将取决于 local_tensor 是否需要梯度。

full_tensor(*, grad_placements=None)[source][source]¶

返回此 DTensor 的完整张量。该操作将执行必要的集合通信操作，以从其 DeviceMesh 中其他秩（rank）上收集本地张量，并将它们拼接在一起。它是以下代码的语法糖：

dtensor.redistribute(placements=[Replicate()] * mesh.ndim).to_local()

Keyword Arguments: grad_placements (List[Placement], 可选) – placements 描述了从该函数返回的完整张量的未来布局。如果将 DTensor 转换为完整的 torch.Tensor，返回的 torch.tensor 可能无法在代码后续部分作为原始的复制 DTensor 布局使用。此参数是用户可以提供给自动微分（autograd）的提示，以应对返回张量的梯度布局与原始复制 DTensor 布局不匹配的情况。如果没有指定，默认假设完整张量的梯度布局为复制模式。
Returns: 一个 torch.Tensor 对象，表示此 DTensor 的完整张量。
Return type: 张量

注意

full_tensor 是可微的。

property placements: Tuple[Placement, ...]¶: 此 DTensor 的 placements 属性，用于描述该 DTensor 在其 DeviceMesh 上的布局。

注意

placements 是一个只读属性，不能被设置。

redistribute(device_mesh=None, placements=None, *, async_op=False)[source][source]¶

redistribute 执行必要的集体操作，将当前的 DTensor 从其当前的放置位置重新分布到新的放置位置，或者从当前的 DeviceMesh 到一个新的 DeviceMesh。例如，我们可以通过为 DeviceMesh 的每个维度指定 Replicate 放置，将一个分片的 DTensor 转换为复制的 DTensor。

在将数据从当前布局重新分配到单个设备网格维度上的新布局时，我们将执行以下操作，包括通信集合操作或本地操作：

Shard(dim) -> Replicate(): all_gather
Shard(src_dim) -> Shard(dst_dim): all_to_all
Replicate() -> Shard(dim): 本地分块（即 torch.chunk）
Partial() -> Replicate(): all_reduce
Partial() -> Shard(dim): reduce_scatter

redistribute 将正确地确定在创建 DTensors 时所需的重新分配步骤，这些 DTensors 可以是在 1-D 或 N-D DeviceMesh 上创建的。

Parameters

device_mesh (DeviceMesh, 可选) – 用于放置 DTensor 的 DeviceMesh。如果未指定，则会使用当前 DTensor 的 DeviceMesh。默认值：None
placements (List[Placement]，可选) – 描述如何将 DTensor 放置到 DeviceMesh 中的新放置方式，必须与 device_mesh.ndim 具有相同数量的元素。默认：在所有网格维度上复制

Keyword Arguments

async_op (bool, 可选) – 是否异步执行 DTensor 重新分布操作。默认值：False

Returns

一个 DTensor 对象

Return type

DTensor

注意

redistribute 是可微的，这意味着用户不需要担心重新分配操作的反向公式。

注意

redistribute 目前仅支持在相同的 DeviceMesh 上重新分配 DTensor，如果您需要将 DTensor 重新分配到不同的 DeviceMesh，请提交一个问题。

to_local(*, grad_placements=None)[source][source]¶

获取此 DTensor 在当前进程（rank）上的本地张量。对于分片（sharding）情形，它返回逻辑张量视图的一个本地分片；对于复制（replication）情形，它返回当前进程（rank）上的副本。

Keyword Arguments: grad_placements (List[Placement]，可选) – placements 描述了从该函数返回的任何梯度布局的未来布局。 to_local 将 DTensor 转换为本地张量，返回的本地张量可能在代码后续部分不再使用原始 DTensor 布局。此参数是用户可以提供给 autograd 的提示，以防返回张量的梯度布局与原始 DTensor 布局不匹配。如果未指定，我们将假设梯度布局与原始 DTensor 保持一致，并将其用于梯度计算。
Returns: 一个 torch.Tensor 或 AsyncCollectiveTensor 对象。它表示当前进程上的本地张量。当返回一个 AsyncCollectiveTensor 对象时，表示本地张量尚未准备好（即通信尚未完成）。在这种情况下，用户需要调用 wait 来等待本地张量就绪。
Return type: 张量

注意

to_local 是可微的，局部张量返回的 requires_grad 将取决于 DTensor 是否需要梯度。

DeviceMesh 作为分布式通信器¶

DeviceMesh 是从 DTensor 构建的，DTensor 是用于描述集群设备拓扑结构并表示多维通信器（基于 ProcessGroup）的抽象。要了解如何创建和使用 DeviceMesh 的详细信息，请参阅 DeviceMesh 示例。

DTensor 位置类型¶

DTensor 支持以下类型的 Placement 每个 DeviceMesh 维度：

class torch.distributed.tensor.placement_types.Shard(dim)[source][source]¶

Shard(dim) 的放置描述了在张量维度 dim 上的 DTensor 分片，该分片分布在对应的 DeviceMesh 维度上。DeviceMesh 维度上的每个进程只持有全局张量的一个片段/切片。Shard(dim) 的放置遵循 torch.chunk(dim) 的语义，当张量维度在 DeviceMesh 维度上不能被整除时，DeviceMesh 维度上的最后几个切片可能是空的。Shard 的放置可以被所有 DTensor API 使用（例如 distribute_tensor、from_local 等）。

Parameters: dim (int) – 描述 DTensor 在其对应的 DeviceMesh 维度上进行切分的张量维度。

警告

在张量维度上进行分片（sharding）时，若该张量维度的大小无法被设备网格（DeviceMesh）维度整除，则此功能目前处于实验阶段，后续可能发生变化。

dim: int¶

class torch.distributed.tensor.placement_types.Replicate[source][source]¶

Replicate() 的放置描述了 DTensor 在对应的 DeviceMesh 维度上进行复制，其中 DeviceMesh 维度上的每个进程都持有全局张量的一个副本。Replicate 的放置可以被所有 DTensor API 使用（例如 distribute_tensor、DTensor.from_local 等）。

class torch.distributed.tensor.placement_types.Partial(reduce_op='sum')[source][source]¶

Partial(reduce_op) 位置描述了在指定 DeviceMesh 维度上等待归约的 DTensor，其中 DeviceMesh 维度上的每个进程持有全局 Tensor 的部分值。用户可以使用 redistribute 将 Partial DTensor 重新分布到指定 DeviceMesh 维度上的 Replicate 或 Shard(dim) 位置，这将触发必要的后台通信操作（即 allreduce、reduce_scatter）。

Parameters: reduce_op (str, 可选) – 用于将部分 DTensor 转换为 Replicated/Sharded DTensor 的归约操作。仅支持逐元素的归约操作，包括：“sum”、“avg”、“product”、“max”、“min”，默认值为“sum”。

注意

Partial 的放置可以作为 DTensor 操作符的结果生成，并且只能由 DTensor.from_local API 使用。

reduce_op: str = 'sum'¶

class torch.distributed.tensor.placement_types.Placement[source][source]¶

Placement 类的基础类，它描述了 DTensor 如何放置到 DeviceMesh。 Placement 和 DeviceMesh 一起可以描述 DTensor 的布局。它是三种主要 DTensor Placement 类型的基础类：Shard、Replicate，和 Partial。

此类不直接使用，主要用作类型占位符。

is_partial()[source][source]¶

Return type: 布尔

is_replicate()[source][source]¶

Return type: 布尔

is_shard(dim=None)[source][source]¶

Return type: 布尔

创建 DTensor 的不同方法¶

There’re three ways to construct a DTensor:

distribute_tensor() 从每个进程上的逻辑或“全局”torch.Tensor 创建一个 DTensor。这可以用来对叶节点 torch.Tensor 进行分片（即模型参数/缓冲区和输入）。
DTensor.from_local() 从每个进程上的本地 torch.Tensor 创建一个 DTensor，该对象可用于从非叶子 torch.Tensor（即前向/反向传播过程中的中间激活张量）创建 DTensor。
DTensor 提供了专用的张量工厂函数（例如 empty()、ones()、randn() 等），以允许通过直接指定 DeviceMesh 和 Placement 来创建不同的 DTensor。与 distribute_tensor() 相比，这可以直接在设备上实现分片内存，而不是在初始化逻辑张量内存后进行分片。

从逻辑 torch.Tensor 创建 DTensor¶

SPMD（单程序，多数据）编程模型在 torch.distributed 中通过 torchrun 启动多个进程来执行相同的程序，这意味着程序内部的模型会首先在不同的进程中初始化（例如，模型可能在 CPU 上、元设备上或直接在 GPU 上初始化，如果内存足够的话）。

DTensor 提供了一个 distribute_tensor() API，可以将模型权重或张量分片到 DTensor 上，在每个进程中从“逻辑”张量创建一个 DTensor。这将使生成的 DTensor 能够遵循单设备语义，这对于 **数值正确性** 至关重要。

torch.distributed.tensor.distribute_tensor(tensor, device_mesh=None, placements=None)[source]¶

将叶节点 torch.Tensor（即 nn.Parameter/buffers）根据指定的 device_mesh 分配到 placements。device_mesh 和 placements 的秩必须相同。tensor 是要分配的逻辑或“全局”张量，API 将使用 DeviceMesh 维度的第一个秩中的 tensor 作为真实来源，以保留单设备语义。如果您想在 Autograd 计算过程中构造一个 DTensor，请改用 DTensor.from_local()。

Parameters

张量 (torch.Tensor) – 要分发的 torch.Tensor。请注意，如果您希望在网格维度中不可被设备数量整除的维度上对张量进行切片，我们将使用 torch.chunk 语义来切片张量并分散切片。不均匀切片行为是实验性的，可能会发生变化。
device_mesh (DeviceMesh, 可选) – 用于分布张量的 DeviceMesh，如果未指定，则必须在 DeviceMesh 上下文管理器中调用，默认值：None
placements (List[Placement], 可选) – 描述如何将张量放置在 DeviceMesh 上的排列方式，必须与 device_mesh.ndim 具有相同数量的元素。如果未指定，默认将在 device_mesh 的第一个维度上复制张量，从 device_mesh 的每个维度的第一个排名开始。

Returns

A DTensor or XLAShardedTensor object.

Return type

DTensor

注意

当使用 xla 个设备类型初始化 DeviceMesh 时，distribute_tensor 返回 XLAShardedTensor。有关更多详细信息，请参阅此问题。XLA 集成是实验性的，并且可能会发生变化。

除了 distribute_tensor() 之外，DTensor 还提供了一个 distribute_module() API，以允许在 nn.Module 级别上更轻松地进行分片。

torch.distributed.tensor.distribute_module(module, device_mesh=None, partition_fn=None, input_fn=None, output_fn=None)[source]¶

此功能提供了三个用于控制模块的参数/输入/输出的函数：

1. 在运行时执行之前通过指定partition_fn（即允许用户根据指定的partition_fn将模块参数转换为DTensor参数）对模块进行分片。 2. 在运行时执行期间通过指定input_fn和output_fn来控制模块的输入或输出。（即，将输入转换为DTensor，将输出转换回torch.Tensor）

Parameters

module (nn.Module) – 用户模块，用于分区。
device_mesh (DeviceMesh) – 用于放置模块的设备网格。
partition_fn (可调用对象) – 用于划分参数的函数（即在 device_mesh 上拆分某些参数）。如果未指定 partition_fn，则默认将 module 的所有模块参数复制到网格中。
input_fn (可调用对象) – 指定输入分布，即可以控制模块的输入如何分片。input_fn 将作为模块安装forward_pre_hook（预前向钩子）。
output_fn (可调用对象) – 指定输出分布，即可以控制如何对输出进行分片，或者将其转换回torch.Tensor。output_fn 将作为模块 forward_hook 安装（后向钩子）。

Returns

一个包含所有参数/缓冲区均为 DTensor 的模块。

Return type

模块

注意

当使用 xla 设备类型初始化 DeviceMesh 时，distribute_module 返回带有 PyTorch/XLA SPMD 注释参数的 nn.Module。有关更多详细信息，请参阅此问题。XLA 集成是实验性的，可能会发生变化。

PyTorch分布式张量工厂函数¶

DTensor 还提供了专用的张量工厂函数，允许通过使用类似 torch.Tensor 的工厂函数 API（例如 torch.ones、torch.empty 等），直接创建 DTensor，并通过额外指定 DeviceMesh 和 Placement 来为创建的 DTensor 设置：

torch.distributed.tensor.zeros(*size, requires_grad=False, dtype=None, layout=torch.strided, device_mesh=None, placements=None)[source]¶

返回一个用标量值 0 填充的 DTensor。

Parameters

size (int...) – 定义输出形状的整数序列。可以是任意数量的参数，也可以是列表或元组等集合。例如：zeros(1,2,3..) 或 zeros([1,2,3..]) 或 zeros((1,2,3..))

Keyword Arguments

requires_grad (bool, 可选) – 如果自动求梯度应记录返回的DTensor上的操作。默认值：False。
数据类型 (torch.dtype，可选) – 返回的数据类型。默认值：如果为 None，使用全局默认设置（参见 torch.set_default_dtype()）。
布局 (torch.layout，可选) – 返回的DTensor的期望布局。默认值：torch.strided。
设备网格 – DeviceMesh 类型，包含排名的网格信息
位置 – 一个由 Placement 类型组成的序列： Shard, Replicate

Returns

一个 DTensor 对象在每个排名上

Return type

DTensor

torch.distributed.tensor.ones(*size, dtype=None, layout=torch.strided, requires_grad=False, device_mesh=None, placements=None)[source]¶

返回一个用标量值 1 填充的 DTensor，其形状由变量参数 size 定义。

Parameters

size (int...) – 定义输出形状的整数序列。可以是任意数量的参数，也可以是列表或元组等集合。例如：ones(1,2,3..) 或 ones([1,2,3..]) 或 ones((1,2,3..))

Keyword Arguments

数据类型 (torch.dtype，可选) – 返回的数据类型。默认值：如果为 None，使用全局默认设置（参见 torch.set_default_dtype()）。
布局 (torch.layout，可选) – 返回的DTensor所需的布局。默认：torch.strided。
requires_grad (bool, 可选) – 如果自动求梯度应记录返回的DTensor上的操作。默认值：False。
设备网格 – DeviceMesh 类型，包含排名的网格信息
位置 – 一个由 Placement 类型组成的序列： Shard, Replicate

Returns

一个 DTensor 对象在每个排名上

Return type

DTensor

torch.distributed.tensor.empty(*size, dtype=None, layout=torch.strided, requires_grad=False, device_mesh=None, placements=None)[source]¶

返回一个用未初始化数据填充的DTensor。其形状由变量参数size定义。

Parameters

size (int...) – 定义输出形状的整数序列。可以是任意数量的参数，也可以是列表或元组等集合。例如：empty(1,2,3..) 或 empty([1,2,3..]) 或 empty((1,2,3..))

Keyword Arguments

数据类型 (torch.dtype，可选) – 返回的DTensor的数据类型。默认值：如果为None，使用全局默认值（参见torch.set_default_dtype()）。布局 (torch.layout，可选)：返回的DTensor的布局。默认值：torch.strided。
requires_grad (bool, 可选) – 如果自动求梯度应记录返回的DTensor上的操作。默认值：False。
设备网格 – DeviceMesh 类型，包含排名的网格信息
位置 – 一个由 Placement 类型组成的序列： Shard, Replicate

Returns

一个 DTensor 对象在每个排名上

Return type

DTensor

torch.distributed.tensor.full(size, fill_value, *, dtype=None, layout=torch.strided, requires_grad=False, device_mesh=None, placements=None)[source]¶

返回一个用DTensor填充的张量，根据device_mesh和placements，其形状由参数size定义。

Parameters

size (int...) – 定义输出形状的整数序列。可以是任意数量的参数，也可以是列表或元组等集合。例如：ones(1,2,3..) 或 ones([1,2,3..]) 或 ones((1,2,3..))
fill_value (标量) – 用于填充输出张量的值。

Keyword Arguments

数据类型 (torch.dtype，可选) – 返回的数据类型。默认值：如果为 None，使用全局默认设置（参见 torch.set_default_dtype()）。
布局 (torch.layout，可选) – 返回的DTensor所需的布局。默认：torch.strided。
requires_grad (bool, 可选) – 如果自动求梯度应记录返回的DTensor上的操作。默认值：False。
设备网格 – DeviceMesh 类型，包含排名的网格信息。
位置 – 一个由 Placement 类型组成的序列： Shard, Replicate

Returns

一个 DTensor 对象在每个排名上

Return type

DTensor

torch.distributed.tensor.rand(*size, requires_grad=False, dtype=None, layout=torch.strided, device_mesh=None, placements=None)[source]¶

返回一个由均匀分布在区间 [0, 1) 上的随机数填充的 DTensor。张量的形状由可变参数 size 定义。

Parameters

size (int...) – 定义输出形状的整数序列。可以是任意数量的参数，也可以是列表或元组等集合。例如：ones(1,2,3..) 或 ones([1,2,3..]) 或 ones((1,2,3..))

Keyword Arguments

数据类型 (torch.dtype，可选) – 返回的数据类型。默认值：如果为 None，使用全局默认设置（参见 torch.set_default_dtype()）。
布局 (torch.layout，可选) – 返回的DTensor所需的布局。默认：torch.strided。
requires_grad (bool, 可选) – 如果自动求梯度应记录返回的DTensor上的操作。默认值：False。
设备网格 – DeviceMesh 类型，包含排名的网格信息。
位置 – 一个由 Placement 类型组成的序列： Shard, Replicate

Returns

一个 DTensor 对象在每个排名上

Return type

DTensor

torch.distributed.tensor.randn(*size, requires_grad=False, dtype=None, layout=torch.strided, device_mesh=None, placements=None)[source]¶

返回一个由均值为0、方差为1的正态分布随机数填充的DTensor。张量的形状由变量参数size定义。

Parameters

size (int...) – 定义输出形状的整数序列。可以是任意数量的参数，也可以是列表或元组等集合。例如：ones(1,2,3..) 或 ones([1,2,3..]) 或 ones((1,2,3..))

Keyword Arguments

数据类型 (torch.dtype，可选) – 返回的数据类型。默认值：如果为 None，使用全局默认设置（参见 torch.set_default_dtype()）。
布局 (torch.layout，可选) – 返回的DTensor所需的布局。默认：torch.strided。
requires_grad (bool, 可选) – 如果自动求梯度应记录返回的DTensor上的操作。默认值：False。
设备网格 – DeviceMesh 类型，包含排名的网格信息。
位置 – 一个由 Placement 类型组成的序列： Shard, Replicate

Returns

一个 DTensor 对象在每个排名上

Return type

DTensor

调试¶

日志记录¶

启动程序时，您可以使用TORCH_LOGS环境变量从torch._logging启用额外的日志记录：

TORCH_LOGS=+dtensor 将显示 logging.DEBUG 条消息及其以上的所有级别。
TORCH_LOGS=dtensor 将显示 logging.INFO 条及以上消息。
TORCH_LOGS=-dtensor 将显示 logging.WARNING 条及以上消息。

调试工具¶

为了调试应用了DTensor的程序，并更详细地了解底层发生了哪些集体操作，DTensor提供了一个CommDebugMode：

class torch.distributed.tensor.debug.CommDebugMode¶

CommDebugMode 是一个上下文管理器，用于计算其上下文中功能集体的数量。它使用一个 TorchDispatchMode 来实现。

示例用法

mod = ...
comm_mode = CommDebugMode()
with comm_mode:
    mod.sum().backward()
print(comm_mode.get_comm_counts())

generate_comm_debug_tracing_table(noise_level=3)[source][source]¶

生成详细表格，展示模块级别上的操作及集体追踪信息。信息量取决于 noise_level

打印模块级汇总计数
打印非简单操作的 dTensor 操作以及模块信息
打印操作不包括在简单操作中
打印所有操作

generate_json_dump(file_name='comm_mode_log.json', noise_level=3)[source][source]¶: 创建用于构建浏览器可视化的 JSON 文件 0. 打印模块级集合操作计数 1. 打印未包含在平凡操作中的分布式张量（dTensor）操作 2. 打印未包含在平凡操作中的操作 3. 打印所有操作

get_comm_counts()[source][source]¶

返回通信计数作为字典。

Returns: 通信被视为字典。
Return type: 字典[Any, 整数]

get_parameter_info()[source][source]¶

Return type: 字典[字符串, 字典[字符串, 任意类型]]

get_sharding_info()[source][source]¶

Return type: 字典[字符串, 字典[字符串, 任意类型]]

get_total_counts()[source][source]¶

Return type: 整数

log_comm_debug_tracing_table_to_file(file_name='comm_mode_log.txt', noise_level=3)[source][source]¶: 替代控制台 CommDebugMode 输出，写入用户指定的文件中

为了可视化维度小于3的DTensor的分片，DTensor提供了visualize_sharding()：

torch.distributed.tensor.debug.visualize_sharding(dtensor, header='')[source]¶: 在终端中可视化 DTensor 的分片，这些分片是一维或二维的。

注意

这需要tabulate包。对于空张量不会打印分片信息。

实验功能¶

DTensor 还提供了一组实验性功能。这些功能要么处于原型阶段，要么基本功能已完成但正在寻求用户反馈。如果您对这些功能有任何意见，请提交给 PyTorch。

torch.distributed.tensor.experimental.context_parallel(mesh, *, buffers=None, buffer_seq_dims=None, no_restore_buffers=None)[source]¶

context_parallel 是一个实验性 API，用于启用上下文并行（CP）。此 API 执行两个操作：1）用支持 CP 的 SDPA (torch.nn.functional.scaled_dot_product_attention) 替换原有的 SDPA，2）在序列维度上对 buffers 进行分片，并且每个排名将根据 mesh 保留相应的分片。

Parameters

网格 (DeviceMesh) – 上下文并行的设备网格。
缓冲区 (可选[List[torch.Tensor]]) – 缓冲区的使用取决于序列维度。例如，输入批次、标签和位置嵌入缓冲区。这些缓冲区必须沿序列维度进行分片以确保准确性。分片将在原地发生，缓冲区的形状会在上下文中发生变化。在上下文结束后，缓冲区将会被恢复。 no_restore_buffers 可以用于指定哪些缓冲区不需要恢复。注意，buffers 不应包含任何 nn.Parameter。
buffer_seq_dims (可选[List[int]]) – buffers 的序列维度。
no_restore_buffers (可选[集合[torch.Tensor]]) – 在上下文退出后，这些集合中的缓冲区不会被恢复。该集合必须是buffers的一个子集。如果在上下文退出后不再使用这些缓冲区，可以将它们放入此列表中以避免额外的恢复时间。

Return type

生成器[无, 无, 无]

警告

torch.distributed._tensor.experimental.attention.context_parallel 是 PyTorch 中的一个原型功能。API 可能会发生变化。

torch.distributed.tensor.experimental.local_map(func, out_placements, in_placements=None, device_mesh=None, *, redistribute_inputs=False)[source]¶

local_map() 是一个实验性 API，允许用户将 DTensor 传递给一个函数，该函数被编写为应用于 torch.Tensor。这是通过提取 DTensor 的本地组件，调用函数，并根据 out_placements 将输出包装成 DTensor 来实现的。

Parameters

函数 (可调用对象) – 要应用于每个本地分片的函数。 DTensor s.
out_placements （PlacementType 或 PlacementType 元组）– 指定在展平后的 DTensor 中的位置。如果展平后的 output 是单个值，那么 out_placements 应该是类型 PlacementType。否则，如果展平后的 output 有多个值，那么 out_placements 应该是一个与展平后的 output 一一对应的 PlacementType 值元组。此外，对于 Tensor 输出，我们使用 PlacementType 作为其位置（一个 Tuple[Placement] 值）。对于非张量输出，PlacementType 应该是 None。需要注意的是，唯一的例外是在没有传递 DTensor 参数的情况下。在这种情况下，即使 out_placements 不是 None，结果函数也应该忽略指定的位置，因为函数不是在 DTensor 上运行的。
in_placements (Tuple[PlacementType, …], 可选) – 指定输入张量中DTensor的位置。如果指定了in_placements，local_map()会检查每个DTensor参数的位置是否与所需位置一致。如果不一致且redistribute_inputs为False，则会抛出异常。否则，如果redistribute_inputs为True，参数将在传递其本地张量到func之前重新分布到所需的位置。唯一例外是当所需位置不是None且参数是一个torch.Tensor时，在这种情况下，将跳过位置检查，并直接将参数传递给func。如果in_placements为None，则不会进行位置检查。默认值：无
设备网格 (DeviceMesh，可选) – 所有DTensor放置在其上的设备网格。如果没有指定，则根据输入的DTensor的设备网格进行推断。local_map要求每个DTensor都放置在相同的设备网格上。默认值：None。
重新分配输入 (bool, 可选) – 表示是否在输入放置与所需输入放置不同时重新划分输入 DTensor 的布尔值。如果该值为 False，且某些 DTensor 输入的放置不同，则会引发异常。默认值：False。

Returns

一个 Callable，它将 func 应用于输入 DTensor 的每个本地分片，并返回由 func 的返回值构建的 DTensor。

Raises

AssertionError – 如果输入 DTensor 没有放置在相同的设备网格上，或者它们被放置在与传入的 device_mesh 参数不同的设备网格上。
AssertionError – 对于任何非DTensor输出，我们要求其对应的输出位置在out_placements为None。如果这不是这种情况，则会引发AssertionError。
ValueError – 如果 redistribute_inputs=False，但输入 DTensor 需要根据 in_placements 进行重新分配。

示例

>>> def mm_allreduce_forward(device_mesh, W, X):
>>>     partial_sum_tensor = torch.mm(W, X)
>>>     reduced_tensor = funcol.all_reduce(partial_sum_tensor, "sum", device_mesh)
>>>     return reduced_tensor
>>>
>>> W = torch.randn(12, 8, requires_grad=False)
>>> X = torch.randn(8, 16, requires_grad=False)
>>> Y = torch.mm(W, X)
>>> row_wise = [Shard(0)]  # row-wise sharding placements on 1-d mesh
>>> col_wise = [Shard(1)]  # col-wise sharding placements on 1-d mesh
>>>
>>> # local_mm_allreduce_forward is the function wrapped with DTensor/Tensor convertion
>>> local_mm_allreduce_forward = local_map(
>>>     mm_allreduce_forward,
>>>     out_placements=[Replicate()],
>>>     in_placements=[col_wise, row_wise],
>>>     device_mesh=device_mesh,
>>> )
>>>
>>> W_dt = distribute_tensor(W, device_mesh, (col_wise))  # col-wisely sharded W tensor
>>> X_dt = distribute_tensor(X, device_mesh, (row_wise))  # row-wisely sharded X tensor
>>> Y_dt = local_mm_allreduce_forward(device_mesh, W_dt, X_dt)  # apply local_mm_allreduce_forward to DTensors

注意

该 API 目前处于实验阶段，可能会发生变化。

torch.distributed.tensor.experimental.register_sharding(op)[source]¶

register_sharding() 是一个实验性的 API，允许用户在张量输入和输出为 DTensor 时为操作注册分片策略。它在以下情况下非常有用：(1) 对于 op 没有默认的分片策略，例如当 op 是一个不受 DTensor 支持的自定义操作符；(2) 当用户希望覆盖现有操作符的默认分片策略时。

Parameters: 操作 (Union[OpOverload, List[OpOverload]]) – 指定要注册自定义分片函数的一个操作或操作列表。
Returns: 一个函数装饰器，可用于包装定义算子分片策略的函数，该算子在op中指定。所定义的分片策略将被注册到DTensor，并且如果DTensor已经实现了该算子，则会覆盖默认的分片策略。自定义的分片函数接受与原始操作相同的输入（除了如果某个参数是一个torch.Tensor，它将被DTensor内部使用的张量对象替换）。该函数应返回一个2元组序列，每个2元组指定了可接受的输出位置及其对应的输入位置。

示例

>>> @register_sharding(aten._softmax.default)
>>> def custom_softmax_sharding(x, dim, half_to_float):
>>>     softmax_dim = dim if dim >= 0 else dim + x.ndim
>>>     acceptable_shardings = []
>>>
>>>     all_replicate = ([Replicate()], [Replicate(), None, None])
>>>     acceptable_shardings.append(all_replicate)
>>>
>>>     for sharding_dim in range(x.ndim):
>>>         if sharding_dim != softmax_dim:
>>>             all_sharded = (
>>>                 [Shard(sharding_dim)],
>>>                 [Shard(sharding_dim), None, None],
>>>             )
>>>             acceptable_shardings.append(all_sharded)
>>>
>>>     return acceptable_shardings

注意

该 API 目前处于实验阶段，可能会发生变化。

torch.distributed.tensor¶

PyTorch DTensor（分布式张量）¶

PyTorch深度张量类API¶

DeviceMesh 作为分布式通信器¶

DTensor 位置类型¶

创建 DTensor 的不同方法¶

从逻辑 torch.Tensor 创建 DTensor¶

PyTorch分布式张量工厂函数¶

调试¶

日志记录¶

调试工具¶

实验功能¶

文档

教程

资源