评价此页

★ ★ ★ ★ ★

torch.distributed.tensor#

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

注意

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 支持三种类型的放置：Shard、Replicate 和 Partial。

DTensor 类 API#

DTensor 是 torch.Tensor 的子类。这意味着一旦创建了 DTensor，它就可以以与 torch.Tensor 非常相似的方式使用，包括运行不同类型的 PyTorch 操作符，就像在单个设备上运行它们一样，从而允许 PyTorch 操作符进行适当的分布式计算。

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

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

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

Shard: 在 DeviceMesh 维度的设备上，张量在张量维度 dim 上分片
Replicate: 在 DeviceMesh 维度的设备上复制张量
Partial: 张量在 DeviceMesh 维度的设备上等待归约

当调用 PyTorch 运算符时，DTensor 会覆盖 PyTorch 运算符以执行分片计算并在必要时发出通信。除了运算符计算，DTensor 将根据运算符语义本身正确地转换或传播布局 (DTensor Layout) 并生成新的 DTensor 输出。

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

注意

这里不建议直接使用 Tensor 子类构造函数来创建 DTensor（即，它不能正确处理自动求导，因此不是公共 API）。请参阅 create_dtensor 部分以了解如何创建 DTensor。

返回类型: DTensor

__create_chunk_list__()[source]#

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

此双下划线方法主要用于分布式检查点目的。

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

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

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

参数

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

关键字参数

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

返回

一个 DTensor 对象

返回类型

DTensor

注意

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

注意

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

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

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

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

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

注意

full_tensor 是可微分的。

redistribute(device_mesh=None, placements=None, *, async_op=False, forward_dtype=None, backward_dtype=None)[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 将正确找出为在 1-D 或 N-D DeviceMesh 上创建的 DTensor 重新分布的必要步骤。

参数

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

关键字参数

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

返回

一个 DTensor 对象

返回类型

DTensor

注意

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

注意

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

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

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

关键字参数: grad_placements (List[Placement], optional) – 描述此函数返回的张量的任何梯度布局的未来布局的布局。to_local 将 DTensor 转换为本地张量，并且返回的本地张量稍后在代码中可能不会用作原始 DTensor 布局。此参数是用户可以给自动求导的提示，以防返回张量的梯度布局与原始 DTensor 布局不匹配。如果未指定，我们将假定梯度布局与原始 DTensor 相同，并将其用于梯度计算。
返回: 一个 torch.Tensor 或 AsyncCollectiveTensor 对象。它表示其当前排名上的本地张量。当返回 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 的布局属性，描述了此 DTensor 在其 DeviceMesh 上的布局。

注意

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

DeviceMesh 作为分布式通信器#

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

DTensor 布局类型#

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

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

Shard(dim) 布局描述了 DTensor 在张量维度 dim 上与其对应的 DeviceMesh 维度上的分片，其中 DeviceMesh 维度上的每个排名只持有全局张量的一个分片/部分。 Shard(dim) 布局遵循 torch.chunk(dim) 语义，当张量维度不能在 DeviceMesh 维度上均匀划分时，DeviceMesh 维度上的最后几个分片可能为空。 Shard 布局可以用于所有 DTensor API（即 distribute_tensor、from_local 等）

参数: dim (int) – 描述 DTensor 在其对应的 DeviceMesh 维度上分片的张量维度。

警告

张量维度不能在 DeviceMesh 维度上均匀划分的分片目前处于实验阶段，可能会有所更改。

dim: int#

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

Replicate() 布局描述了 DTensor 在相应 DeviceMesh 维度上的复制，其中 DeviceMesh 维度上的每个排名都持有全局张量的一个副本。 Replicate 布局可以用于所有 DTensor API（即 distribute_tensor、DTensor.from_local 等）

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

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

参数: reduce_op (str, optional) – 用于部分 DTensor 以生成复制/分片 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]#

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

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

is_partial()[source]#

返回类型: 布尔值

is_replicate()[source]#

返回类型: 布尔值

is_shard(dim=None)[source]#

返回类型: 布尔值

创建 DTensor 的不同方法#

有三种方法可以构造一个 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#

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

DTensor 提供了一个 distribute_tensor() API，可以将模型权重或张量分片到 DTensor，它将从每个进程上的“逻辑”张量创建一个 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 和 placements 的排名必须相同。要分发的 tensor 是逻辑或“全局”张量，API 将使用 DeviceMesh 维度的第一个排名上的 tensor 作为事实来源，以保留单设备语义。如果您想在自动求导计算中间构造 DTensor，请改用 DTensor.from_local()。

参数

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

关键字参数

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

返回

一个 DTensor 或 XLAShardedTensor 对象。

返回类型

DTensor

注意

当使用 xla device_type 初始化 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）。

参数

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

返回

一个包含参数/缓冲区都是 DTensor 的模块。

返回类型

模块

注意

当使用 xla device_type 初始化 DeviceMesh 时，distribute_module 返回带有 PyTorch/XLA SPMD 注释参数的 nn.Module。有关更多详细信息，请参见此问题。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) – 如果自动求导应记录返回的 DTensor 上的操作。默认值：False。
dtype (torch.dtype, optional) – 返回 DTensor 所需的数据类型。默认值：如果为 None，则使用全局默认值（参见 torch.set_default_dtype()）。
layout (torch.layout, optional) – 返回 DTensor 所需的布局。默认值：torch.strided。
device_mesh – DeviceMesh 类型，包含排名的网格信息
placements – Placement 类型的序列：Shard、Replicate

返回

每个排名上的 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) – 如果自动求导应记录返回的 DTensor 上的操作。默认值：False。
device_mesh – DeviceMesh 类型，包含排名的网格信息
placements – Placement 类型的序列：Shard、Replicate

返回

每个排名上的 DTensor 对象

返回类型

DTensor

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

返回一个填充了未初始化数据的 DTensor。 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()）。布局 (torch.layout, optional)：返回 DTensor 所需的布局。默认值：torch.strided。
requires_grad (bool, optional) – 如果自动求导应记录返回的 DTensor 上的操作。默认值：False。
device_mesh – DeviceMesh 类型，包含排名的网格信息
placements – Placement 类型的序列：Shard、Replicate

返回

每个排名上的 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 的 DTensor，其形状由参数 size 定义。

参数

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) – 如果自动求导应记录返回的 DTensor 上的操作。默认值：False。
device_mesh – DeviceMesh 类型，包含排名的网格信息。
placements – Placement 类型的序列：Shard、Replicate

返回

每个排名上的 DTensor 对象

返回类型

DTensor

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

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

参数

关键字参数

dtype (torch.dtype, optional) – 返回 DTensor 所需的数据类型。默认值：如果为 None，则使用全局默认值（参见 torch.set_default_dtype()）。
layout (torch.layout, optional) – 返回 DTensor 所需的布局。默认值：torch.strided。
requires_grad (bool, optional) – 如果自动求导应记录返回的 DTensor 上的操作。默认值：False。
device_mesh – DeviceMesh 类型，包含排名的网格信息。
placements – Placement 类型的序列：Shard、Replicate

返回

每个排名上的 DTensor 对象

返回类型

DTensor

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

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

参数

关键字参数

dtype (torch.dtype, optional) – 返回 DTensor 所需的数据类型。默认值：如果为 None，则使用全局默认值（参见 torch.set_default_dtype()）。
layout (torch.layout, optional) – 返回 DTensor 所需的布局。默认值：torch.strided。
requires_grad (bool, optional) – 如果自动求导应记录返回的 DTensor 上的操作。默认值：False。
device_mesh – DeviceMesh 类型，包含排名的网格信息。
placements – Placement 类型的序列：Shard、Replicate

返回

每个排名上的 DTensor 对象

返回类型

DTensor

调试#

日志记录#

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

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]#

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

打印模块级集体计数
打印不包含在琐碎操作中的 DTensor 操作，模块信息
打印不包含在琐碎操作中的操作
打印所有操作

generate_json_dump(file_name='comm_mode_log.json', noise_level=3)[source]#: 创建用于构建浏览器可视化的 json 文件。0. 打印模块级集体计数 1. 打印不包含在琐碎操作中的 DTensor 操作 2. 打印不包含在琐碎操作中的操作 3. 打印所有操作

get_comm_counts()[source]#

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

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

get_parameter_info()[source]#

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

get_sharding_info()[source]#

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

get_total_counts()[source]#

返回类型: int

log_comm_debug_tracing_table_to_file(file_name='comm_mode_log.txt', noise_level=3)[source]#: 除了控制台 CommDebugMode 输出之外，还写入用户指定的文件

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

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

可视化 1D 或 2D DTensor 的分片。

注意

这需要 tabulate 包，或 rich 和 matplotlib。空张量不会打印分片信息。

实验性功能#

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) 进行修补，2) 沿序列维度分片 buffers，并且每个排名将根据 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, out_placements, in_placements=None, in_grad_placements=None, device_mesh=None, *, redistribute_inputs=False)[source]#

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

参数

func (Callable) – 要应用于 DTensor 的每个本地分片的函数。
out_placements (Union[PlacementType, Tuple[PlacementType, …]]) – func 的扁平化输出中 DTensor 的所需放置。如果扁平化 output 是单个值，则 out_placements 应为 PlacementType 类型。否则，如果扁平化 output 具有多个值，则 out_placements 应为 PlacementType 值的元组，与扁平化 output 一一对应。此外，对于 Tensor 输出，我们使用 PlacementType 作为其放置（Tuple[Placement] 值）。对于非 Tensor 输出，PlacementType 应为 None。请注意，唯一的例外是当没有传入 DTensor 参数时。在这种情况下，即使 out_placements 不是 None，结果函数也应忽略所需的放置，因为该函数未与 DTensor 一起运行。
in_placements (Tuple[PlacementType, …], optional) – func 的扁平化输入中 DTensor 的所需放置。如果指定了 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 梯度的放置提示。此参数是用户可以提供给 to_local() 的提示，以防本地张量输入的梯度布局与其 DTensor 输入布局不匹配。如果未指定，我们将假定本地张量输入的梯度布局与原始 DTensor 输入保持相同，并将其用于梯度计算。默认值：None。
device_mesh (DeviceMesh, optional) – 所有 DTensor 放置所在的设备网格。如果未指定，则将从输入 DTensor 的设备网格推断。 local_map 要求所有 DTensor 放置在相同的设备网格上。默认值：None。
redistribute_inputs (bool, optional) – 布尔值，指示当输入 DTensor 的放置与所需的输入放置不同时，是否重新切分。如果此值为 False 且某些 DTensor 输入具有不同的放置，则会引发异常。默认值：False。

返回

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

引发

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) 当用户希望覆盖现有操作符的默认切分策略时。

参数: op (Union[OpOverload, List[OpOverload]]) – 要注册自定义切分函数的一个操作符或操作符列表。
返回: 一个函数装饰器，可用于包装一个函数，该函数定义了 op 中指定的操作符的切分策略。定义的切分策略将注册到 DTensor，如果 DTensor 已经实现了该操作符，则将覆盖默认切分策略。自定义切分函数接受与原始操作符相同的输入（除了如果参数是 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 布局类型#

创建 DTensor 的不同方法#

从逻辑 torch.Tensor 创建 DTensor#

DTensor 工厂函数#

调试#

日志记录#

调试工具#

实验性功能#

文档

教程

资源