PyTorch DTensor 功能及原理源码解读
简介 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会为你在必要时插入所需要的通信操作。 ...