聊聊 PyTorch 中新的Distributed API （二）

作者丨无恶不作来源丨https://zhuanlan.zhihu.com/p/681775092编辑丨GiantPandaCV下面结合 PyTorch conference (Composable Distributed PT2(D)) 和 PyTorch 2.2 release 做介绍。总体来说相对2.0时代初期的API有了较大变化，新的通信方式对于初学者总体上更加友好。DTensorDistributedTensor 又叫 DTensor，是基础分布式 tensor，可用于 tensor 级别的切分和计算，DTensor有 2 个重要概念：DeviceMeshDeviceMesh是一种抽象表示全局拓扑关系。图1 DeviceMeshDTensor placement:DTensor placement是tensor的分布式表示方式，有两种类型：shard, replicate。图2 DTensor placementDTensor是torch.Tensor的子类，可以通过from_local和to_local和tensor进行转化，也可以通过Reshard和Redistribute进行分布式转换。提供高阶API，转换tensor/module:1. distribute_tensor(tensor)2. distribute_module(module)2. TP and User API张量并行可以根据自定义module接口进行转换，这里主要介绍 ………………………………