torch.distributed.tensor

DTensor 类 APIs

DTensor 是 torch.Tensor 的子类。这意味着一旦创建了 DTensor，就可以像使用 torch.Tensor 一样使用它，包括运行不同类型的 PyTorch 操作符，就像在单个设备上运行一样，从而实现 PyTorch 操作符的分布式计算。

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

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

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

• Shard: 在DeviceMesh维度的设备上，张量在dim维度上进行分片

• Replicate: 在DeviceMesh维度的设备上复制张量

• Partial: 张量在DeviceMesh维度的设备上等待减少

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

为了确保在调用PyTorch操作符时DTensor分片计算的数值正确性，需要将操作符的每个张量参数设置为DTensor。

返回类型

DTensor

属性device_mesh:DeviceMesh

与这个 DTensor 对象相关的 DeviceMesh 属性。

注意

device_mesh 是一个只读属性，无法进行设置。

静态 from_local(local_tensor, device_mesh=None, placements=None, *, run_check=False, shape=None, stride=None) [源代码](/_modules/torch/distributed/tensor/_api.html#DTensor.from_local)

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

参数

• local_tensor (torch.Tensor) – 每个 rank 上的本地 torch.Tensor。

• device_mesh (DeviceMesh, 可选) - 用于放置张量的设备网格。如果未指定，则需要在 DeviceMesh 上下文中进行调用， 默认值为 None。

• placements (List[Placement], optional) – 描述如何将本地 torch.Tensor 放置在 DeviceMesh 上的方式，必须与 device_mesh.ndim 的元素数量相等。

关键字参数

• run_check (bool, 可选) – 以额外通信为代价，在各个 ranks 上执行完整性检查，验证每个本地张量的元信息以确保正确性。如果在 placements 中存在 Replicate，设备网格维度上的第一个 rank 的数据将被广播到其他 ranks。默认值：False

• shape (torch.Size, 可选) – 一个整数列表，指定在local_tensor基础上构建的DTensor的大小。如果不同rank上的local_tensor形状不一致，则必须提供此参数。如果没有提供，默认假设给定的分布式张量在各个rank上均匀分片来计算shape。默认值：None

• stride (元组, 可选) – 一个整数列表，用于指定 DTensor 的步长。如果没有提供，默认情况下假定给定的分布式张量在各个 ranks 上均匀分片。stride 默认值为 None。

返回值

一个DTensor对象

返回类型

DTensor

注意

当 run_check=False 时，确保传入的本地张量在各个 ranks 上是正确的（即对于 Shard(dim) 布局进行了分片或对于 Replicate() 布局进行了复制）是用户的责任。否则，创建的 DTensor 的行为将是未定义的。

注意

from_local 是可微的，创建的 DTensor 对象的 requires_grad 属性将根据 local_tensor 是否需要梯度来决定。

full_tensor(*, grad_placements=None)[源代码]

返回此 DTensor 的完整张量。它会执行必要的通信操作，从设备网格中的其他设备收集本地张量，并将它们拼接在一起。这相当于以下代码的简化形式：

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

关键字参数

grad_placements (List[Placement]，可选) – 描述从该函数返回的完整张量的任何未来梯度布局。full_tensor 将 DTensor 转换为完整的 torch.Tensor，并且返回的 torch.tensor 可能在后续代码中不会被用作原始复制的 DTensor 布局。此参数是用户可以提供给 autograd 的提示，告知如果返回张量的梯度布局与原始复制的 DTensor 布局不匹配的情况。如果没有指定，默认假设完整张量的梯度布局为复制。

返回值

一个表示此 DTensor 完整张量的 torch.Tensor 对象。

返回类型

Tensor

注意

full_tensor 是可微的。

propertyplacements:Tuple[Placement,...]

此DTensor的placements属性描述了它在DeviceMesh上的布局。

注意

placements 是一个只读属性，无法进行设置。

redistribute(device_mesh=None, placements=None, *, async_op=False)[源代码]

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

当在设备网格的一个维度上从当前位置重新分布到新位置时，我们将执行以下操作，包括通信集合和本地操作：

1. Shard(dim) -> Replicate(): all_gather

2. Shard(src_dim) -> Shard(dst_dim): all_to_all

3. Replicate() -> Shard(dim): 当前切分（即 torch.chunk）

4. Partial() -> Replicate(): all_reduce

5. Partial() -> Shard(dim): reduce_scatter

redistribute 能正确计算无论是在一维还是多维设备网格上创建的DTensor所需的所有重新分布步骤。

参数

• device_mesh (DeviceMesh, 可选) - 用于放置 DTensor 的设备网格。如果未指定，默认使用当前 DTensor 的设备网格。默认值：None

• placements (List[Placement], optional) – 描述如何将 DTensor 放入 DeviceMesh 中的新放置方式，必须与 device_mesh.ndim 的元素数量相同。默认情况下，在所有网格维度上进行复制。

关键字参数

async_op (bool, optional) – 是否异步执行 DTensor 重分布操作。默认为 False。

返回值

一个DTensor对象

返回类型

DTensor

注意

redistribute 是可微的，因此用户无需担心重分配操作的反向计算公式。

注意

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

to_local(*, grad_placements=None)[源代码]

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

关键字参数

grad_placements (List[Placement]，可选) – 描述从该函数返回的张量的未来梯度布局。to_local 将 DTensor 转换为本地张量，并且返回的本地张量可能不会在后续代码中继续使用原始 DTensor 布局。此参数是用户可以提供给自动微分器的提示，告知其如果返回张量的梯度布局与原始 DTensor 不匹配的情况。如果没有指定，默认假设梯度布局保持不变，并以此进行计算。

返回值

A torch.Tensor 或 AsyncCollectiveTensor 对象。它表示当前 rank 上的本地张量。当返回一个 AsyncCollectiveTensor 对象时，说明本地张量还未准备好（即通信尚未完成）。此时，用户需要调用 wait 方法来等待本地张量准备就绪。

返回类型

Tensor

注意

to_local 是可微的，返回的本地张量的 requires_grad 属性取决于 DTensor 是否设置了 requires_grad。

DeviceMesh 作为分布式通讯工具

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

DTensor放置类型

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

类torch.distributed.tensor.placement_types.Shard(dim)[源代码]

The Shard(dim) placement describes how the DTensor is sharded along tensor dimension dim over a corresponding DeviceMesh dimension, where each rank on the DeviceMesh dimension holds only one shard of the global Tensor. The Shard(dim) placement follows the torch.chunk(dim) semantics: when the tensor dimension is not evenly divisible by the DeviceMesh dimension, some ranks at the end might be empty. The Shard placement can be used with all DTensor APIs (e.g., distribute_tensor, from_local).

参数

dim (int) – 表示 DTensor 在其相应 DeviceMesh 维度上进行分区的张量维度。

警告

当张量维度的大小不能被设备网格维度整除时，在设备网格维度上对该张量维度进行分片操作目前还处于实验阶段，未来可能有所变化。

dim: int

类torch.distributed.tensor.placement_types.Replicate[源代码]

The Replicate() placement describes how a DTensor is replicated across a corresponding DeviceMesh dimension, where each rank in the DeviceMesh holds a replica of the global Tensor. The Replicate placement can be used by all DTensor APIs (e.g., distribute_tensor, DTensor.from_local, etc.).

类torch.distributed.tensor.placement_types.Partial(reduce_op='sum')[源代码]

The Partial(reduce_op) placement describes the DTensor that is pending reduction on a specified DeviceMesh dimension, where each rank on the DeviceMesh dimension holds the partial value of the global Tensor. 用户可以使用redistribute操作将Partial DTensor重新分布到指定的DeviceMesh维度上的Replicate或Shard(dim)布局，这会触发底层必要的通信操作（如allreduce, reduce_scatter）。

参数

reduce_op (str, 可选) – 用于部分 DTensor 的缩减操作，以生成 Replicated 或 Sharded DTensor。仅支持元素级的缩减操作，包括： “sum”，“avg”，“product”，“max”，“min”。默认值为 “sum”。

注意

Partial 布局可以通过 DTensor 操作生成，并且只能通过 DTensor.from_local API 来使用。

reduce_op:str='sum'

类torch.distributed.tensor.placement_types.Placement[源代码]

这是 Placement 类型的基础类，用于描述 DTensor 如何放置在 DeviceMesh 上。结合 Placement 和 DeviceMesh 可以定义 DTensor 的布局。它是三种主要的 DTensor Placement 类型：Shard、Replicate 和 Partial 的基础。

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

is_partial()[源代码]

返回类型

bool

is_replicate()[源代码]

返回类型

bool

is_shard(dim=None)[源代码]

返回类型

bool

将逻辑 torch.Tensor 转换为 DTensor

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

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

torch.distributed.tensor.distribute_tensor(tensor, device_mesh=None, placements=None)

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

参数

• tensor (torch.Tensor) – 要进行分布的 torch.Tensor。如果要在某个维度上对张量进行分区，而该维度不能被该网状结构中的设备数量整除，则使用 torch.chunk 语义来分割张量并分散碎片。这种不均匀的分区行为是实验性的，并且可能会发生变化。

• device_mesh (DeviceMesh, 可选) - 用于分配张量的设备网格。如果未指定，则需要在 DeviceMesh 上下文管理器中进行调用，默认值为 None。

• placements (List[Placement], optional) – 描述如何在DeviceMesh上放置张量的位置信息，必须与device_mesh.ndim的元素数量相同。如果没有指定，默认会从device_mesh每个维度的第一个rank开始将张量在整个device_mesh上进行复制。

返回值

一个 DTensor 或 XLAShardedTensor 对象。

返回类型

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)

该函数提供了三种功能来控制模块的参数、输入和输出。

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) – 指定输入分布，可以控制模块输入的分片方式。该函数将被安装为模块的 forward_pre_hook（前向钩子）。

• output_fn (Callable) – 指定输出分布，可以控制输出的分片方式或将其转换回 torch.Tensor。output_fn 将被安装为模块 forward_hook（前向传播后的钩子）。

返回值

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

返回类型

Module

注意

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

DTensor 工厂函数

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)

返回一个填充了标量值 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, 可选) – 返回的 DTensor 的期望数据类型。默认情况下，如果为 None，则使用全局默认值（参见 torch.set_default_dtype()）。

• layout (torch.layout, 可选) – 返回的 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)

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

参数

size (int...) – 定义输出 DTensor 形状的一系列整数。这些整数可以作为可变数量的参数传递，也可以作为一个列表或元组等集合形式提供。例如：ones(1, 2, 3..)、ones([1, 2, 3..]) 或 ones((1, 2, 3..))

关键参数

• dtype (torch.dtype, 可选) – 返回的 DTensor 的期望数据类型。默认情况下，如果为 None，则使用全局默认值（参见 torch.set_default_dtype()）。

• layout (torch.layout, 可选) — 指定返回的 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)

返回一个DTensor，该张量包含未初始化的数据。DTensor 的形状由变量参数 size 定义。

参数

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

关键参数

• dtype (torch.dtype, 可选) – 返回的 DTensor 的期望数据类型。默认：如果为 None，使用全局默认值（参见 torch.set_default_dtype()）。layout (torch.layout, 可选) – 返回的 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)

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

参数

• size (int...) – 定义输出 DTensor 形状的一系列整数。这些整数可以作为可变数量的参数传递，也可以作为一个列表或元组等集合形式提供。例如：ones(1, 2, 3..)、ones([1, 2, 3..]) 或 ones((1, 2, 3..))

• fill_value (标量) – 填充输出张量的值。

关键参数

• dtype (torch.dtype, 可选) – 返回的 DTensor 的期望数据类型。默认情况下，如果为 None，则使用全局默认值（参见 torch.set_default_dtype()）。

• layout (torch.layout, 可选) — 指定返回的 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)

返回一个DTensor，该张量包含从区间 [0, 1) 上的均匀分布中随机生成的数字。张量的形状由变量参数 size 定义。

参数

size (int...) – 定义输出 DTensor 形状的一系列整数。这些整数可以作为可变数量的参数传递，也可以作为一个列表或元组等集合形式提供。例如：ones(1, 2, 3..)、ones([1, 2, 3..]) 或 ones((1, 2, 3..))

关键参数

• dtype (torch.dtype, 可选) – 返回的 DTensor 的期望数据类型。默认情况下，如果为 None，则使用全局默认值（参见 torch.set_default_dtype()）。

• layout (torch.layout, 可选) — 指定返回的 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)

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

参数

size (int...) – 定义输出 DTensor 形状的一系列整数。这些整数可以作为可变数量的参数传递，也可以作为一个列表或元组等集合形式提供。例如：ones(1, 2, 3..)、ones([1, 2, 3..]) 或 ones((1, 2, 3..))

关键参数

• dtype (torch.dtype, 可选) – 返回的 DTensor 的期望数据类型。默认情况下，如果为 None，则使用全局默认值（参见 torch.set_default_dtype()）。

• layout (torch.layout, 可选) — 指定返回的 DTensor 的布局。默认值为 torch.strided。

• requires_grad (bool, optional) – 是否应在返回的DTensor上记录自动求梯度操作。默认值：False。

• device_mesh – DeviceMesh 类型，包含排名的网格信息。

• placements — 一个包含 Placement 类型的序列：如 Shard 和 Replicate

返回值

每个进程中的 DTensor 对象

返回类型

DTensor

日志记录

在启动程序时，你可以通过设置 TORCH_LOGS 环境变量（来自torch._logging）来开启额外的日志记录：

• TORCH_LOGS=+dtensor 会显示 logging.DEBUG 及以上级别所有的日志消息。

• TORCH_LOGS=dtensor 会显示 logging.INFO 及以上级别的日志消息。

• TORCH_LOGS=-dtensor 将显示 logging.WARNING 及更高级别的日志消息。

调试工具

为了调试使用了DTensor的程序，并深入了解底层发生的集合操作细节，DTensor提供了一个CommDebugMode:

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

生成详细的表格，显示模块级别的操作和集体追踪信息。信息的多少取决于噪声级别。

0. 显示模块级别的总计数

1. 打印除简单操作外的dTensor操作以及模块信息

2. 打印不包含在简单操作中的操作

3. 显示所有操作

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 输出的替代方案，将内容写入用户指定的文件

为了可视化一个具有少于三个维度的DTensor的分区情况，DTensor提供了visualize_sharding()：

torch.distributed.tensor.debug.visualize_sharding(dtensor, header='')

在终端中可视化DTensor的一维或二维分片情况。

注意

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