PyTorch DTensor 功能及原理源码解读

Tue, 23 Dec 2025 22:37:00 +0800

简介

PyTorch FSDP2 对比 FSDP1的升级上一篇文章梳理了FSDP2，于是顺着梳理下DTensor相关的内容。在DTensor出现之前，Pytorch 中的分布式张量主要依赖于 torch.distributed 模块下的DDP、 FSDP及第三方库如DeepSpeed、Megatron-LM 提供的并行能力。使用这些框架往往需要在代码中不断地插入通信(如all-gather、all-reduce等)，有了DTensor，降低用户手动管理通信操作的成本，进一步提高易用性，也是与torch的理念相匹配的。

DTensor基本用法

基本属性

DTensor是torch.Tensor的子类，因此可以像普通tensor一样操作它。DTensor通过DeviceMesh建立各个rank之间的关系，并通过Placement描述当前的状态，主要有三个Placement：Shard、Replicate和Partial

Replicate：每个rank拥有相同的完整参数；
Shard：每个rank只持有部份分块后的参数；
Partial：每个rank只持有部份数据，目前是等待all-reduce状态

三个属性之间的转换关系：

torch官方罗列了几个转换关系：

Shard(dim) -> Replicate(): all_gather
Shard(src_dim) -> Shard(dst_dim): all_to_all
Replicate() -> Shard(dim): local chunking (i.e. torch.chunk)
Partial() -> Replicate(): all_reduce
Partial() -> Shard(dim): reduce_scatter

API：

一个DTensor可以由如下API创建：

distribute_tensor() 接口

local_tensor = torch.randn(8, 16, device=f"cuda:{rank}")

dtensor_shard = distribute_tensor(
    local_tensor,
    device_mesh=mesh,
    placements=[Shard(0)],
)

from local tensor

local_tensor = torch.tensor([[1, 2], [3, 4]])
dtensor_shard = DTensor(
    local_tensor,
    device_mesh=dist.DeviceMesh("cuda", list(range(world_size))),
    placements=[dist.Placement("shard", dim=0)]  # 沿第0维分片
)
或者
local_tensor = torch.tensor([1, 2, 3])
dtensor = local_tensor.to_dtensor(
    device_mesh=dist.DeviceMesh("cuda", [0,1]),
    placements=[dist.Placement("replicate")]
)

工厂函数

dtensor_ones = DTensor.ones(
    (2, 4),
    device_mesh=dist.DeviceMesh("cuda", list(range(world_size))),
    placements=[dist.Placement("shard", dim=1)]
)

创建好DTensor之后，就可以把DTensor当作普通的Tensor对象去操作和运算，DTensor会为你在必要时插入所需要的通信操作。

DTensor on YuZhouye's Blog

PyTorch DTensor 功能及原理源码解读

简介

DTensor基本用法

基本属性

API：