评价此页

★ ★ ★ ★ ★

torch.distributed.tensor#

创建于：2025年6月13日 | 最后更新于：2025年8月23日

注意

torch.distributed.tensor 目前处于 Alpha 阶段，正在开发中。我们承诺大部分文档中列出的 API 向后兼容，但如果需要，可能会进行 API 更改。

PyTorch DTensor (分布式张量)#

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

请参阅基于 DTensor 构建的 PyTorch 原生并行解决方案的示例。

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

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

DTensor 类 API#

DTensor 是 torch.Tensor 的子类。这意味着一旦创建了 DTensor，它就可以像 torch.Tensor 一样使用，包括运行各种 PyTorch 算子，就像在单设备上运行一样，从而为 PyTorch 算子提供适当的分布式计算。

除了现有的 torch.Tensor 方法外，它还提供了一组额外的方法来与 torch.Tensor 交互，对 DTensor 进行 redistribute（重新分布）布局，获取所有设备上的完整张量内容等。

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

DTensor (分布式张量) 是 torch.Tensor 的子类，它提供了类似单设备的抽象，用于对多设备 torch.Tensor 进行编程。它通过 DeviceMesh 和以下类型的 Placement 来描述分布式张量的分片布局 (DTensor Layout)。

Shard：张量在 DeviceMesh 的 dim 维度上，按照张量的 dim 维度进行分片。
Replicate：张量在 DeviceMesh 维度上的设备上进行复制。
Partial：张量在 DeviceMesh 维度上的设备上等待归约。

调用 PyTorch 算子时，DTensor 会重写 PyTorch 算子，以便在必要时执行分片计算并发出通信。除了算子计算，DTensor 还会根据算子语义本身正确地转换或传播 placement（DTensor Layout），并生成新的 DTensor 输出。

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

注意

直接使用 Tensor 子类构造函数不是创建 DTensor 的推荐方式（即它不能正确处理 autograd，因此不是公共 API）。请参阅 create_dtensor 部分，了解如何创建 DTensor。

返回类型: DTensor

__create_chunk_list__()[source]#

返回一个 ChunkStorageMetadata 列表，这是一个描述当前 rank 上本地分片/副本大小/偏移量的 dataclass。对于 DTensor，每个 rank 将拥有一个本地分片/副本，因此返回的列表通常只有一个元素。

这个 dunder 方法主要用于分布式 checkpoint。

返回: 一个 List[ChunkStorageMetadata] 对象，表示当前 rank 上的分片大小/偏移量。

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

根据指定的 device_mesh 和 placements，从每个 rank 上的本地 torch.Tensor 创建一个 DTensor。

参数

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

关键字参数

run_check (bool, optional) – 以额外的通信为代价，跨 rank 执行健全性检查，以检查每个本地张量的元信息以确保正确性。如果在 placements 中有 Replicate，则 device mesh 维度的第一个 rank 上的数据将被广播到其他 rank。默认值：False
shape (torch.Size, optional) – 一个指定 DTensor 大小的整数列表，该 DTensor 构建在 local_tensor 之上。注意，如果 local_tensor 的大小在 rank 之间不同，则需要提供此参数。如果未提供，则假设给定的分布式张量在 rank 之间均匀分片，从而计算出 shape。默认值：None
stride (tuple, optional) – 一个指定 DTensor stride 的整数列表。如果未提供，则假设给定的分布式张量在 rank 之间均匀分片，从而计算出 stride。默认值：None

返回

一个 DTensor 对象。

返回类型

DTensor

注意

当 run_check=False 时，用户有责任确保传入的本地张量在 rank 之间是正确的（即，对于 Shard(dim) placement，张量是分片的；对于 Replicate() placement，张量是复制的）。否则，创建的 DTensor 的行为是未定义的。

注意

from_local 是可微分的，创建的 DTensor 对象上的 requires_grad 将取决于 local_tensor 是否 requires_grad。

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

返回此 DTensor 的完整张量。它将执行必要的集合操作，以收集其 DeviceMesh 上的其他 rank 的本地张量并将它们连接起来。这是以下代码的语法糖：

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

关键字参数: grad_placements (List[Placement], optional) – 描述从此函数返回的完整张量的任何梯度布局的未来布局的 placement。 full_tensor 将 DTensor 转换为一个完整的 torch.Tensor，并且返回的 torch.tensor 在代码的后续部分可能无法用作原始复制的 DTensor 布局。此参数是用户可以提供给 autograd 的提示，以防返回张量的梯度布局与原始复制的 DTensor 布局不匹配。如果未指定，我们将假设完整张量的梯度布局是复制的。
返回: 一个表示此 DTensor 完整张量的 torch.Tensor 对象。
返回类型: 张量

注意

full_tensor 是可微分的。

redistribute(device_mesh=None, placements=None, *, async_op=False, forward_dtype=None, backward_dtype=None)[source]#

redistribute 执行必要的集体操作，将当前 DTensor 从其当前 placement 重新分布到新的 placement，或从其当前 DeviceMesh 重新分布到新的 DeviceMesh。也就是说，通过为 DeviceMesh 的每个维度指定 Replicate placement，我们可以将 Sharded DTensor 转换为 Replicated DTensor。

当在单个设备 mesh 维度上从当前 placement 重新分布到新的 placement 时，我们将执行以下操作，包括通信集体或本地操作：

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 会正确地为在 1D 或 N-D DeviceMesh 上创建的 DTensor 确定必要的重新分布步骤。

参数

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

关键字参数

async_op (bool, optional) – 是否异步执行 DTensor 重新分布操作。默认值：False
forward_dtype (torch.dtype, optional) – 在其 forward 传播中重新分布本地张量之前，可以将本地张量的数据类型转换为 forward_dtype。结果 DTensor 的数据类型将为 forward_dtype。默认值：None。
backward_dtype (torch.dtype, optional) – 在其 backward 传播中重新分布本地张量之前，可以将本地张量的数据类型转换为 backward_dtype。结果 DTensor 的梯度将被转换回当前 DTensor 的数据类型。默认值：None

返回

一个 DTensor 对象。

返回类型

DTensor

注意

redistribute 是可微分的，这意味着用户无需担心 redistribute 操作的 backward 公式。

注意

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

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

获取 DTensor 在其当前 rank 上的本地张量。对于分片，它返回逻辑张量视图的本地分片；对于复制，它返回当前 rank 上的副本。

关键字参数: grad_placements (List[Placement], optional) – 描述从此函数返回的张量的任何梯度布局的未来布局的 placement。 to_local 将 DTensor 转换为本地张量，并且返回的本地张量在代码的后续部分可能无法用作原始 DTensor 布局。此参数是用户可以提供给 autograd 的提示，以防返回张量的梯度布局与原始 DTensor 布局不匹配。如果未指定，我们将假设梯度布局与原始 DTensor 保持不变，并将其用于梯度计算。
返回: 一个 torch.Tensor 或 AsyncCollectiveTensor 对象。它表示当前 rank 上的本地张量。当返回 AsyncCollectiveTensor 对象时，表示本地张量尚未就绪（即通信尚未完成）。在这种情况下，用户需要调用 wait 来等待本地张量就绪。
返回类型: 张量

注意

to_local 是可微分的，返回的本地张量的 requires_grad 将取决于 DTensor 是否 requires_grad。

property device_mesh: DeviceMesh#: 与此 DTensor 对象关联的 DeviceMesh 属性。

注意

device_mesh 是只读属性，无法设置。

property placements: tuple[torch.distributed.tensor.placement_types.Placement, ...]#: 此 DTensor 的 placements 属性，描述了该 DTensor 在其 DeviceMesh 上的布局。

注意

placements 是只读属性，无法设置。

DeviceMesh 作为分布式通信器#

DeviceMesh 是从 DTensor 构建的，作为描述集群设备拓扑并表示多维通信器（基于 ProcessGroup）的抽象。有关如何创建/使用 DeviceMesh 的详细信息，请参阅 DeviceMesh 教程。

DTensor Placement 类型#

DTensor 支持以下类型的 Placement 在每个 DeviceMesh 维度上。

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

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

参数: dim (int) – 张量维度，描述 DTensor 沿着其对应的 DeviceMesh 维度进行分片。

警告

沿张量维度进行分片，而张量维度大小不能被 DeviceMesh 维度整除，目前处于实验阶段，可能会发生变化。

dim: int#

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

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

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

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

参数: reduce_op (str, optional) – 用于部分 DTensor 生成 Replicated/Sharded DTensor 的归约 op。仅支持逐元素归约操作，包括：“sum”、“avg”、“product”、“max”、“min”，默认值：“sum”。

注意

Partial placement 可以作为 DTensor 算子的结果生成，并且只能由 DTensor.from_local API 使用。

reduce_op: str = 'sum'#

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

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

此类不应直接使用，主要作为类型存根。

is_partial(reduce_op=None)[source]#

返回类型: 布尔值

is_replicate()[source]#

返回类型: 布尔值

is_shard(dim=None)[source]#

返回类型: 布尔值

创建 DTensor 的不同方式#

有三种方法可以构造一个 DTensor：

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

从逻辑 torch.Tensor 创建 DTensor#

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

DTensor 提供了一个 distribute_tensor() API，它可以分片模型权重或张量到 DTensor，从而使创建的 DTensor 符合单设备语义，这对于 **数值正确性** 至关重要。

torch.distributed.tensor.distribute_tensor(tensor, device_mesh=None, placements=None, *, src_data_rank=0)[source]#

根据指定的 placements，将一个叶子 torch.Tensor（即 nn.Parameter/buffers）分发到 device_mesh。 device_mesh 的 rank 和 placements 的数量必须相同。要分发的 tensor 是逻辑/全局张量，API 将使用 DeviceMesh 第一个 rank 上的 tensor 作为事实来源以保留单设备语义。如果您想在 Autograd 计算的中间构建 DTensor，请改用 DTensor.from_local()。

参数

tensor (torch.Tensor) – 要分发的 torch.Tensor。请注意，如果您想在设备 mesh 维度的设备数量不能整除的维度上分片张量，我们将使用 torch.chunk 语义来分片张量并分散分片。不均匀分片行为是实验性的，可能会发生变化。
device_mesh (DeviceMesh, optional) – 用于分发张量的 DeviceMesh，如果未指定，则必须在 DeviceMesh 上下文管理器下调用，默认值：None
placements (List[Placement], optional) – 描述如何将张量放置在 DeviceMesh 上的 placement，必须与 device_mesh.ndim 的元素数量相同。如果未指定，我们将默认将张量从 device_mesh 的每个维度的第一个 rank 复制到该 device_mesh。

关键字参数

src_data_rank (int, optional) – 逻辑/全局张量源数据的 rank，distribute_tensor() 使用它来将分片/副本分散/广播到其他 rank。默认情况下，我们在每个 DeviceMesh 维度的 group_rank=0 作为源数据，以保留单设备语义。如果显式传递 None，distribute_tensor() 将直接使用其本地数据，而不是尝试通过 scatter/broadcast 来保留单设备语义。默认值：0

返回

一个每个 rank 上的 DTensor 或 XLAShardedTensor 对象。

返回类型

DTensor

注意

当使用 xla device_type 初始化 DeviceMesh 时，distribute_tensor 返回 XLAShardedTensor。有关更多详细信息，请参阅此 issue。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 将 Module 参数转换为 DTensor 参数) 来在运行时执行之前对模块进行分片。2. 通过指定 input_fn 和 output_fn 来控制模块在运行时期间的输入或输出。(例如，将输入转换为 DTensor，将输出转换回 torch.Tensor)。

参数

module (nn.Module) – 用户要分区的模块。
device_mesh (DeviceMesh) – 用于放置模块的设备 mesh。
partition_fn (Callable) – 用于分区参数的函数（例如，在 device_mesh 上分片某些参数）。如果未指定 partition_fn，则默认情况下我们将 module 的所有模块参数复制到 mesh 上。
input_fn (Callable) – 指定输入分布，例如，可以控制模块的输入如何分片。input_fn 将作为模块的 forward_pre_hook (forward 前置钩子) 安装。
output_fn (Callable) – 指定输出分布，例如，可以控制输出如何分片，或将其转换回 torch.Tensor。output_fn 将作为模块的 forward_hook (forward 后置钩子) 安装。

返回

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

返回类型

模块

注意

当使用 xla device_type 初始化 DeviceMesh 时，distribute_module 返回带有 PyTorch/XLA SPMD 注释的参数的 nn.Module。有关更多详细信息，请参阅此 issue。XLA 集成处于实验阶段，可能会发生变化。

DTensor 工厂函数#

DTensor 还提供了专用的张量工厂函数，允许使用类似 torch.Tensor 的工厂函数 API（例如 torch.ones, torch.empty, 等）直接创建 DTensor，此外还可以为创建的 DTensor 指定 DeviceMesh 和 Placement。

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

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

参数

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

关键字参数

requires_grad (bool, optional) – 如果 autograd 应该记录返回的 DTensor 上的操作。默认值：False。
dtype (torch.dtype, optional) – 所需返回 DTensor 的数据类型。默认值：如果 None，则使用全局默认值（请参阅 torch.set_default_dtype()）。
layout (torch.layout, optional) – 所需返回 DTensor 的布局。默认值：torch.strided。
device_mesh – DeviceMesh 类型，包含 rank 的 mesh 信息。
placements – 一个 Placement 类型的序列：Shard、Replicate。

返回

每个 rank 上的一个 DTensor 对象。

返回类型

DTensor

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

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

参数

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

关键字参数

dtype (torch.dtype, optional) – 所需返回 DTensor 的数据类型。默认值：如果 None，则使用全局默认值（请参阅 torch.set_default_dtype()）。
layout (torch.layout, optional) – 所需返回 DTensor 的布局。默认值：torch.strided。
requires_grad (bool, optional) – 如果 autograd 应该记录返回的 DTensor 上的操作。默认值：False。
device_mesh – DeviceMesh 类型，包含 rank 的 mesh 信息。
placements – 一个 Placement 类型的序列：Shard、Replicate。

返回

每个 rank 上的一个 DTensor 对象。

返回类型

DTensor

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

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

参数

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

关键字参数

dtype (torch.dtype, optional) – 所需返回 DTensor 的数据类型。默认值：如果 None，则使用全局默认值（请参阅 torch.set_default_dtype()）。 layout (torch.layout, optional): 所需返回 DTensor 的布局。默认值：torch.strided。
requires_grad (bool, optional) – 如果 autograd 应该记录返回的 DTensor 上的操作。默认值：False。
device_mesh – DeviceMesh 类型，包含 rank 的 mesh 信息。
placements – 一个 Placement 类型的序列：Shard、Replicate。

返回

每个 rank 上的一个 DTensor 对象。

返回类型

DTensor

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

根据 device_mesh 和 placements，使用 fill_value 填充，形状由参数 size 定义，返回一个 DTensor。

参数

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

关键字参数

dtype (torch.dtype, optional) – 所需返回 DTensor 的数据类型。默认值：如果 None，则使用全局默认值（请参阅 torch.set_default_dtype()）。
layout (torch.layout, optional) – 所需返回 DTensor 的布局。默认值：torch.strided。
requires_grad (bool, optional) – 如果 autograd 应该记录返回的 DTensor 上的操作。默认值：False。
device_mesh – DeviceMesh 类型，包含 rank 的 mesh 信息。
placements – 一个 Placement 类型的序列：Shard、Replicate。

返回

每个 rank 上的一个 DTensor 对象。

返回类型

DTensor

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

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

参数

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

关键字参数

dtype (torch.dtype, optional) – 所需返回 DTensor 的数据类型。默认值：如果 None，则使用全局默认值（请参阅 torch.set_default_dtype()）。
layout (torch.layout, optional) – 所需返回 DTensor 的布局。默认值：torch.strided。
requires_grad (bool, optional) – 如果 autograd 应该记录返回的 DTensor 上的操作。默认值：False。
device_mesh – DeviceMesh 类型，包含 rank 的 mesh 信息。
placements – 一个 Placement 类型的序列：Shard、Replicate。

返回

每个 rank 上的一个 DTensor 对象。

返回类型

DTensor

torch.distributed.tensor.randn(*size, requires_grad=False, dtype=None, layout=torch.strided, device_mesh=None, placements=None)[源代码]#

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

参数

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

关键字参数

dtype (torch.dtype, optional) – 所需返回 DTensor 的数据类型。默认值：如果 None，则使用全局默认值（请参阅 torch.set_default_dtype()）。
layout (torch.layout, optional) – 所需返回 DTensor 的布局。默认值：torch.strided。
requires_grad (bool, optional) – 如果 autograd 应该记录返回的 DTensor 上的操作。默认值：False。
device_mesh – DeviceMesh 类型，包含 rank 的 mesh 信息。
placements – 一个 Placement 类型的序列：Shard、Replicate。

返回

每个 rank 上的一个 DTensor 对象。

返回类型

DTensor

随机操作#

DTensor 提供了分布式 RNG 功能，以确保分片张量上的随机操作获得唯一值，并且复制张量上的随机操作获得相同值。该系统要求所有参与的 rank（例如 SPMD rank）在执行每个 dtensor 随机操作之前都使用相同的生成器状态开始，如果这是真的，它确保在每个 dtensor 随机操作完成后它们都处于相同的状态。随机操作期间不执行通信来同步 RNG 状态。

接受 generator 关键字参数的操作将利用用户传入的生成器，如果传入了，否则使用设备上的默认生成器。无论使用哪个生成器，在 DTensor 操作之后它都会被推进。将同一个生成器用于 DTensor 和非 DTensor 操作是有效的，但必须小心确保非 DTensor 操作在所有 rank 上平等地推进生成器状态。

当结合使用 DTensor 和流水线并行时，每个流水线阶段的 rank 应使用不同的种子，而流水线阶段内的 rank 应使用相同的种子。

DTensor 的 RNG 基础架构基于 philox 算法，并支持任何基于 philox 的后端（cuda 和其他类似 cuda 的设备），但不幸的是，尚不支持 CPU 后端。

调试#

日志记录#

启动程序时，可以使用 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)[源代码]#

生成详细表格，显示模块级别的操作和通信跟踪信息。信息的数量取决于 noise_level。

打印模块级别的通信计数。
打印未包含在平凡操作中的 dTensor 操作，以及模块信息。
打印未包含在平凡操作中的操作。
打印所有操作。

generate_json_dump(file_name='comm_mode_log.json', noise_level=3)[源代码]#: 创建用于构建浏览器可视化的 json 文件。0. 打印模块级别的通信计数；1. 打印未包含在平凡操作中的 dTensor 操作；2. 打印未包含在平凡操作中的操作；3. 打印所有操作。

get_comm_counts()[源代码]#

以字典形式返回通信计数。

返回: 通信计数以字典形式返回。
返回类型: Dict[Any, int]

get_parameter_info()[源代码]#

返回类型: dict[str, dict[str, Any]]

get_sharding_info()[源代码]#

返回类型: dict[str, dict[str, Any]]

get_total_counts()[源代码]#

返回类型: int

log_comm_debug_tracing_table_to_file(file_name='comm_mode_log.txt', noise_level=3)[源代码]#: 与控制台 CommDebugMode 输出的替代方法，写入用户指定的文件的内容。

为了可视化少于 3 个维度的 DTensor 的分片，DTensor 提供了 visualize_sharding()。

torch.distributed.tensor.debug.visualize_sharding(dtensor, header='', use_rich=False)[源代码]#

在终端中可视化 1D 或 2D DTensor 的分片。

注意

这需要 tabulate 包，或者 rich 和 matplotlib。对于空张量，将不打印任何分片信息。

实验性功能#

DTensor 还提供了一系列实验性功能。这些功能处于原型阶段，或者基本功能已完成但正在征求用户反馈。如果您对这些功能有反馈，请在 PyTorch 上提交一个 issue。

torch.distributed.tensor.experimental.context_parallel(mesh, *, buffers=None, buffer_seq_dims=None, no_restore_buffers=None)[源代码]#

context_parallel 是一个实验性 API，用于启用上下文并行 (CP)。该 API 执行两项操作：1) 使用支持 CP 的 SDPA（torch.nn.functional.scaled_dot_product_attention）进行补丁，2) 沿序列维度分片 buffers，每个 rank 根据 mesh 保留相应的分片。

参数

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

返回类型

Generator[None, None, None]

警告

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

torch.distributed.tensor.experimental.local_map(func=None, out_placements=None, in_placements=None, in_grad_placements=None, device_mesh=None, *, redistribute_inputs=False)[源代码]#

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

参数

func (Callable) – 要应用于 DTensor s 的每个局部分片的函数。
out_placements (Union[PlacementType, Tuple[PlacementType, …]]) – func 的展平输出中 DTensor s 的期望放置。如果展平的 output 是单个值，则 out_placements 应为 PlacementType 类型。否则，如果展平的 output 有多个值，则 out_placements 应为 PlacementType 值的元组，与展平的 output 是一对一映射。此外，对于 Tensor 输出，我们使用 PlacementType 作为其放置（一个 Tuple[Placement] 值）。对于非 Tensor 输出，PlacementType 应为 None。请注意，唯一的例外是当没有传入 DTensor 参数时。在这种情况下，即使 out_placements 不为 None，结果函数也应忽略期望的放置，因为函数不是用 DTensor s 运行的。
in_placements (Tuple[PlacementType, …], optional) – func 的展平输入中 DTensor s 的必需放置。如果指定了 in_placements，local_map() 将检查每个 DTensor 参数的放置是否与必需的放置相同。如果放置不相同且 redistribute_inputs 为 False，则会引发异常。否则，如果 redistribute_inputs 为 True，则参数将首先重新分发到必需的分片放置，然后才将其局部张量传递给 func。唯一的例外是当必需的放置不为 None 且参数为 torch.Tensor 时。在这种情况下，将跳过放置检查，并将参数直接传递给 func。如果 in_placements 为 None，则不执行放置检查。默认值：None
in_grad_placements (Tuple[PlacementType, …], optional) – 与展平输入 DTensor 对应的 DTensor s 梯度的放置提示。此参数是用户可以提供给 to_local() 的提示，以防局部张量输入的梯度布局与其 DTensor 输入布局不匹配。如果未指定，我们将假定局部张量输入的梯度布局与原始 DTensor 输入保持相同，并使用该布局进行梯度计算。默认值：None。
device_mesh (DeviceMesh, optional) – 输出 DTensor s 放置在其上的设备网格。如果未指定，将从第一个输入 DTensor 的设备网格推断。默认值：None。

关键字参数

redistribute_inputs (bool, optional) – 布尔值，指示当输入 DTensor s 的放置与必需的输入放置不同时，是否重新分片这些输入 DTensor s。如果此值为 False 且某些 DTensor 输入具有不同的放置，则会引发异常。默认值：False。

返回

一个 Callable，它将 func 应用于输入 DTensor 的每个局部分片，并返回一个从 func 的返回值构造的 DTensor。

引发

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 conversion
>>> 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)[源代码]#

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

参数: op (Union[OpOverload, List[OpOverload]]) – 要注册自定义分片函数的运算符或运算符列表。
返回: 一个函数装饰器，可用于包装一个定义指定运算符 op 的分片策略的函数。定义的 P分片策略将注册到 DTensor，如果 DTensor 已经实现了该运算符，则会覆盖默认的分片策略。自定义分片函数接受与原始 op 相同的输入（除了如果一个参数是 torch.Tensor，它将被 DTensor 内部使用的类似张量的对象替换）。该函数应返回一个 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 (分布式张量)#

DTensor 类 API#

DeviceMesh 作为分布式通信器#

DTensor Placement 类型#

创建 DTensor 的不同方式#

从逻辑 torch.Tensor 创建 DTensor#

DTensor 工厂函数#

随机操作#

调试#

日志记录#

调试工具#

实验性功能#

文档

教程

资源