PyTorch DTensor 分布式张量编程指南SPMD 抽象、Placement 体系与模块级分发实战【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorchPyTorch DTensor即torch.distributed.tensor原型阶段曾被称为 DistributedTensor是 PyTorch 官方为统一分布式训练而设计的张量级抽象它用Shard/Replicate/Partial三种 Placement 描述张量在多维设备网格DeviceMesh上的分布方式让开发者可以在 SPMD单程序多设备范式下像编写单机代码一样编写分布式程序。本文以仓库内 torch/distributed/tensor/README.md 为核心结合torch/distributed/tensor/目录下的源码实现系统讲解 DTensor 的设计动机、Placement 语义、三大核心 APIinit_device_mesh/distribute_tensor/distribute_module的完整用法与底层原理。读完本文你将能够用几行代码完成张量切分、实现模块级 Tensor Parallelism并理解 DTensor 与 DDP/FSDP、编译器式分布式训练之间的关系。DTensor 是什么三行代码实现张量切分DTensor 提出了分布式张量原语distributed tensor primitives目标是让开发者以 SPMD 范式轻松编写分布式计算只需描述张量的分布方式sharding 与 replication 并行策略即可在任意设备网格上表达张量分布从而支撑原生 Tensor Parallelism 及更高级的并行探索。官方 README 给出了一个最直观的例子——把一个 10 万行的大张量切分到 4 张卡上仅需三行核心代码# to run this file (i.e. dtensor_example.py): # torchrun --standalone --nnodes1 --nproc-per-node4 dtensor_example.py import os import torch from torch.distributed.tensor import init_device_mesh, Shard, distribute_tensor # Create a mesh topology with the available devices: # 1. We can directly create the mesh using elastic launcher, (recommended) # 2. If using mp.spawn, one needs to initialize the world process_group first and set device # i.e. torch.distributed.init_process_group(backendnccl, world_sizeworld_size) mesh init_device_mesh(cuda, (int(os.environ[WORLD_SIZE]),)) big_tensor torch.randn(100000, 88) # Shard this tensor over the mesh by sharding big_tensors 0th dimension over the 0th dimension of mesh. my_dtensor distribute_tensor(big_tensor, mesh, [Shard(dim0)])其中torchrun --standalone --nnodes1 --nproc-per-node4 dtensor_example.py是官方推荐的启动方式elastic launcher如果使用mp.spawn则需先调用torch.distributed.init_process_group(backendnccl, world_sizeworld_size)初始化进程组。启动后每个进程运行同一份代码SPMD但各自持有big_tensor第 0 维的 1/4 分片。init_device_mesh构建多维设备网格init_device_mesh定义于 torch/distributed/device_mesh.py它根据device_type和mesh_shape创建一个 n 维数组布局的DeviceMeshn 为mesh_shape的长度。其完整签名如下def init_device_mesh( device_type: str, mesh_shape: tuple[int, ...], *, mesh_dim_names: tuple[str, ...] | None None, backend_override: dict[ int | str, str | C10dBackend.Options | tuple[str, C10dBackend.Options] ] | None None, ) - DeviceMesh从源码的 docstring 可以确认以下使用前提与限制SPMD 一致性init_device_mesh遵循 SPMD 编程模型集群中所有进程运行相同的 Python 程序因此mesh_shape必须在所有 rank 上保持一致否则可能导致挂起。自动初始化进程组如果当前没有找到进程组init_device_mesh会在后台自动初始化分布式通信所需的进程组。device_type 支持范围目前支持cpu、cuda/cuda-like、xpu不允许传入带 GPU 索引的设备类型如cuda:0。mesh_dim_names为每个网格维度命名长度必须与mesh_shape一致且每个名字必须唯一源码中会做重复名检查例如init_device_mesh(cuda, (2, 8), mesh_dim_names(dp, tp))可分别以dp、tp标识数据并行与张量并行维度。backend_override可对部分或全部网格维度的 ProcessGroup 进行后端覆盖key 可以是维度索引或维度名提供了mesh_dim_names时value 可以是后端名、C10dBackend.Options或两者的元组。动机为什么需要一个统一的张量级抽象当前扩展大规模分布式训练主要有三种方式数据并行Data Parallel、张量并行Tensor Parallel和流水线并行Pipeline Parallel。它们各自在独立维度上工作PyTorch 生态为此分别构建了 DDP、FSDP、ShardedTensor、PiPPy 等独立方案。当用户训练超大模型时往往需要将这三种并行组合使用即三维并行 3-D Parallelism但现有方案之间的互操作性不佳、使用门槛高——用户很难随意组合数据并行、张量并行与流水线并行。根本原因在于缺少一个能在不同并行策略之间搭建桥梁的公共抽象。理想场景是用户构建分布式程序就像在单机/单设备上编写程序一样无需关心集群内的分布式训练细节由框架自动决定如何拆分模型、如何在节点间跑流水线并行、如何在节点内跑数据并行和张量并行。要实现这一点就需要公共抽象来统一描述张量值的分布方式与分布式计算。DTensor 正是在这一背景下被引入作为下一代 ShardedTensor为分布式的存储与计算提供基础抽象并作为分布式程序转换的基本构建块可无缝构建张量并行、DDP、FSDP 等并行策略。价值主张DTensor 带来什么README 明确了 PyTorch DTensor 的三点核心价值统一的state_dict存取在进行 checkpoint 时提供统一的保存/加载方式即使存在复杂的张量存储分布策略如将张量并行与 FSDP 参数分片结合也能正确处理。eager 模式下的张量并行相比 ShardedTensorDTensor 允许灵活混用 sharding 与 replication从而在 eager 模式下直接表达张量并行。SPMD 编程模型的入口作为 SPMD 编程模型的入口点同时也是基于编译器的分布式训练的基础构建块。从目录结构看torch/distributed/tensor/内部还提供了_sharding_prop.py分片传播、_redistribute.py重分布、_dispatch.py算子分发以及_ops/内建 DTensor 算子实现等模块共同支撑上述能力。Placement 体系Shard / Replicate / Partial 的语义与限制DTensor 的核心是三种 Placement放置策略它们定义在 torch/distributed/tensor/placement_types.py 中分别是Shard、Replicate和PartialShard(dim)按维度切分Shard(dim)描述 DTensor 在张量的第dim维上、沿 DeviceMesh 的对应维度切分网格维度上的每个 rank 只持有全局张量的一份分片。源码明确说明其遵循torch.chunk(dim)语义——当张量维度不能被网格维度整除时网格维度上最后几个分片可能为空对应不均匀切分目前为实验性行为可能发生变化。因此Shard用于表达 row-wiseShard(0)或 col-wiseShard(1)切分等场景。Replicate()整张复制Replicate()描述 DTensor 在对应 DeviceMesh 维度上进行复制网格维度上的每个 rank 都持有全局张量的完整副本。源码中Replicate._make_replicate_tensor的实现表明复制通过从网格维度的首个坐标作为数据源进行广播broadcast完成若 rank 不属于该网格则返回空张量。Replicate常用于数据并行式的复制语义。Partial(reduce_op)待归约的局部结果Partial(reduce_op)描述的是一个待归约的 DTensor它在对应网格维度上并不持有完整值或单一分片而是持有部分和/局部结果需要通过指定的reduce_op归约后才能得到全局值。源码显示支持的reduce_op包括sum、avg、min、max其__repr__为Partial(sum)这样的形式字符串表示缩写为P(sum)。Partial常用于表达 matmul 等算子中间结果的分布状态。混用限制Partial 归约类型必须一致三种 Placement 一般可以自由组合但存在一个重要限制一个 DTensor 不能同时混用不同的Partial归约类型例如Partial(sum)与Partial(max)混用。原因是非线性归约如max与线性归约如sum不可交换不满足交换律在重分布redistribution过程中相对顺序在语义上至关重要。因此单个 DTensor 中的所有PartialPlacement 必须使用相同的归约操作# Valid: same partial type placements [Partial(sum), Partial(sum), Shard(0)] # OK # Invalid: mixed partial types - will raise ValueError placements [Partial(sum), Partial(max), Shard(0)] # Error!该限制同样体现在实现层面在 torch/distributed/tensor/_api.py 中distribute_tensor在分发前会调用assert_no_mixed_partial_types(placements)对 placements 做校验混用不同归约类型会直接抛出ValueError。DTensor 基础 API 实战构造、from_local 与 redistributetorch.distributed.tensor对外导出了DTensor、Shard、Replicate、Partial、distribute_tensor、distribute_module、init_device_mesh、DeviceMesh等公共 API见 torch/distributed/tensor/init.py。以下示例覆盖三类典型用法直接构造 DTensor表达纯切分、纯复制、切分复制等不同策略从本地torch.Tensor构造 DTensorDTensor.from_local对已有 DTensor 做重分片reshard迁移到新的 DTensor Layoutredistribute。# torchrun --standalone --nnodes1 --nproc-per-node4 dtensor_example.py import torch from torch.distributed.tensor import DTensor, Shard, Replicate, distribute_tensor, distribute_module, init_device_mesh # construct a device mesh with available devices (multi-host or single host) device_mesh init_device_mesh(cuda, (4,)) # if we want to do row-wise sharding rowwise_placement[Shard(0)] # if we want to do col-wise sharding colwise_placement[Shard(1)] big_tensor torch.randn(888, 12) # distributed tensor returned will be sharded across the dimension specified in placements rowwise_tensor distribute_tensor(big_tensor, device_meshdevice_mesh, placementsrowwise_placement) # if we want to do replication across a certain device list replica_placement [Replicate()] # distributed tensor will be replicated to all four GPUs. replica_tensor distribute_tensor(big_tensor, device_meshdevice_mesh, placementsreplica_placement) # if we want to distribute a tensor with both replication and sharding device_mesh init_device_mesh(cuda, (2, 2)) # replicate across the first dimension of device mesh, then sharding on the second dimension of device mesh spec[Replicate(), Shard(0)] partial_replica distribute_tensor(big_tensor, device_meshdevice_mesh, placementsspec) # create a DistributedTensor that shards on dim 0, from a local torch.Tensor local_tensor torch.randn((8, 8), requires_gradTrue) rowwise_tensor DTensor.from_local(local_tensor, device_mesh, rowwise_placement) # reshard the current row-wise tensor to a colwise tensor or replicate tensor colwise_tensor rowwise_tensor.redistribute(device_mesh, colwise_placement) replica_tensor colwise_tensor.redistribute(device_mesh, replica_placement)要点说明placements 与 mesh 维数必须一致源码中distribute_tensor会校验len(placements) device_mesh.ndim不一致时抛出ValueErrortorch/distributed/tensor/_api.py。placements 缺省行为若不指定placements默认将张量在 DeviceMesh 每个维度上复制[Replicate() for _ in range(device_mesh.ndim)]。leaf tensor 限制distribute_tensor用于分发 leaf 张量如nn.Parameter/buffer若传入非叶子张量会抛出RuntimeError。若需在 Autograd 计算中途构造 DTensor应改用DTensor.from_local。src_data_rankdistribute_tensor默认以每个网格维度的group_rank0作为逻辑/全局张量的数据源通过 scatter/broadcast 保持单设备语义显式传入None时则直接使用本地数据。redistributeredistribute(device_mesh, placements)将当前 DTensor 迁移到新的 Layout其底层实现位于 torch/distributed/tensor/_redistribute.py负责在三种 Placement 之间做转换如从 row-wise 切分转到列切分、或转成全复制。高层 APIdistribute_tensor 与 distribute_module用户可以直接使用 DTensor 张量构造器如distributed.ones/empty创建分布式张量但对于nn.Linear这类已持有torch.Tensor参数的现有模块如何让参数变为分布式DTensor 提供了两个高层 APIdef distribute_tensor(tensor: torch.Tensor, device_mesh: DeviceMeshNone, placements: List[Placement]None): distribute the tensor according to device_mesh and placements, tensor could be a meta tensor. def distribute_module( module: nn.Module, device_mesh: DeviceMeshNone, partition_fn: Callable[[str, nn.Module, DeviceMesh], ...]None, input_fn: Callable[...., None]None, output_fn: Callable[...., None]None, ): This function converts all module parameters to distributed tensor parameters according to the partition_fn specified. It could also control the input/output of the module by specifying the input_fn and output_fn. distribute_module 的三个回调机制从 torch/distributed/tensor/_api.py 的实现细节可以确认三个回调的分工与执行时机partition_fn(name, module, device_mesh)在运行时执行之前被调用用于切分模块的参数/buffer——把普通torch.Tensor参数替换为DTensor参数。它对每个子模块含模块自身各调用一次name是子模块的完整限定名顶层模块为即nn.Module.named_modules的返回。未被partition_fn转换的参数/buffer 之后会被复制到整个 DeviceMesh若partition_fn为None则所有参数/buffer 全部复制。input_fn(module, inputs, device_mesh)在运行时执行期间通过forward_pre_hook安装将模块输入转换为 DTensorinputs是forward位置参数构成的元组关键字参数不会传给input_fn。output_fn(module, outputs, device_mesh)在运行时执行期间通过forward_hook安装将模块输出通常是 DTensor可能为任意嵌套结构转换回普通torch.Tensor典型实现是调用DTensor.to_local或DTensor.full_tensor或重分布到其他 placements。注意input_fn/output_fn早期版本采用双参数形式(inputs, device_mesh)该形式已废弃请使用上面文档化的三参数形式。模块级张量并行完整示例下面的例子将自定义模块的nn.Linear参数按列切分Shard(0)实现模块级张量并行import torch.nn as nn from torch.distributed.tensor import Shard, distribute_tensor, distribute_module, init_device_mesh class MyModule(nn.Module): def __init__(self) - None: super().__init__() self.fc1 nn.Linear(8, 8) self.fc2 nn.Linear(8, 8) self.relu nn.ReLU() def forward(self, input): return self.relu(self.fc1(input) self.fc2(input)) mesh init_device_mesh(cuda, (4,)) def shard_params(mod_name, mod, mesh): col_linear_placement [Shard(0)] # shard fc1 and fc2 if isinstance(mod, nn.Linear): for name, param in mod.named_parameters(): dist_param nn.Parameter( distribute_tensor(param, mesh, col_linear_placement) ) mod.register_parameter(name, dist_param) sharded_module distribute_module(MyModule(), mesh, partition_fnshard_params)结合input_fn/output_fn还可以进一步控制运行时输入输出例如用to_replicated把普通输入张量分发为复制 DTensor 再进入forward用to_local内部调用outputs.full_tensor()把 DTensor 输出还原为普通张量。torch/distributed/tensor/examples/目录下的torchrec_sharding_example.py、convnext_example.py等示例脚本可作为继续深入的参考。编译器与 DTensor从 eager 优化到编译式分布式训练DTensor 为张量并行等场景提供了高效方案但以数据并行的方式使用 DTensor 的复制Replication时性能可能明显低于 DDP/FSDP 等现有方案。原因在于DDP/FSDP 拥有整个模型架构的全局视图可以针对数据并行做专门的优化如集合通信融合、计算与通信重叠等而 DTensor 作为类张量对象只能在单个张量操作内部做优化。为此官方正在 DTensor 之上探索基于编译器的方案从用户程序中提取图信息以暴露更多性能优化机会从而提升基于 DTensor 的数据并行训练效率。这也是 DTensor 被定位为编译器式分布式训练基础构建块的原因。相关技术脉络与后续演进DTensor 的设计主要受三个使用单一分布式张量概念统一表达复制与切分的工作启发这些工作都支持用统一的 SPMD 编程模型构建分布式训练程序GSPMDJAX/TensorFlow 分布式训练的基础组件借助 XLA 编译器实现大规模高效训练。它在一个张量内定义了三种切分策略——tiled分片、replicated复制、partially tiled部分分片来表达切分与复制GSPMD Partitioner 的核心是借助 XLA 编译器做分片传播与基于编译器的融合等高级优化。OneFlow GlobalTensorOneFlow 构建了自己的 GlobalTensor 概念是 GSPMD 切分的一种变体。它也定义了三类张量但与 GSPMD 略有不同split切分、broadcast广播、partial sum部分和——不使用 partially tiled而是用 partial sum 概念来表达值的归约。TensorFlow DTensorTensorFlow 同步分布式训练的扩展支持在多维 mesh如设备网络上同时做切分与复制并实现了 MLIR passes 来完成传播与算子实现。此外还有多个将张量切分嵌入系统的高阶研究方向例如 Megatron-LM 面向 Transformer 模型的张量并行、DeepSpeed 在张量切分之上结合多种优化技术训练大规模模型。本项目中的 DTensor 原型Prototype同样处于早期反馈收集阶段官方已发布相关 RFC 并在 dev-discuss 论坛同步征求反馈完整设计文档可结合torch/distributed/tensor/源码目录继续研读——尤其是 placement_types.pyPlacement 定义与张量切分/广播/归约实现、_api.pydistribute_tensor/distribute_module实现、_sharding_prop.py分片传播规则以及_ops/目录内建算子的分片规则与策略它们共同构成了 DTensor 从声明式 Placement 到实际分布式计算的核心链路。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
PyTorch DTensor 分布式张量编程指南：SPMD 抽象、Placement 体系与模块级分发实战
PyTorch DTensor 分布式张量编程指南SPMD 抽象、Placement 体系与模块级分发实战【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorchPyTorch DTensor即torch.distributed.tensor原型阶段曾被称为 DistributedTensor是 PyTorch 官方为统一分布式训练而设计的张量级抽象它用Shard/Replicate/Partial三种 Placement 描述张量在多维设备网格DeviceMesh上的分布方式让开发者可以在 SPMD单程序多设备范式下像编写单机代码一样编写分布式程序。本文以仓库内 torch/distributed/tensor/README.md 为核心结合torch/distributed/tensor/目录下的源码实现系统讲解 DTensor 的设计动机、Placement 语义、三大核心 APIinit_device_mesh/distribute_tensor/distribute_module的完整用法与底层原理。读完本文你将能够用几行代码完成张量切分、实现模块级 Tensor Parallelism并理解 DTensor 与 DDP/FSDP、编译器式分布式训练之间的关系。DTensor 是什么三行代码实现张量切分DTensor 提出了分布式张量原语distributed tensor primitives目标是让开发者以 SPMD 范式轻松编写分布式计算只需描述张量的分布方式sharding 与 replication 并行策略即可在任意设备网格上表达张量分布从而支撑原生 Tensor Parallelism 及更高级的并行探索。官方 README 给出了一个最直观的例子——把一个 10 万行的大张量切分到 4 张卡上仅需三行核心代码# to run this file (i.e. dtensor_example.py): # torchrun --standalone --nnodes1 --nproc-per-node4 dtensor_example.py import os import torch from torch.distributed.tensor import init_device_mesh, Shard, distribute_tensor # Create a mesh topology with the available devices: # 1. We can directly create the mesh using elastic launcher, (recommended) # 2. If using mp.spawn, one needs to initialize the world process_group first and set device # i.e. torch.distributed.init_process_group(backendnccl, world_sizeworld_size) mesh init_device_mesh(cuda, (int(os.environ[WORLD_SIZE]),)) big_tensor torch.randn(100000, 88) # Shard this tensor over the mesh by sharding big_tensors 0th dimension over the 0th dimension of mesh. my_dtensor distribute_tensor(big_tensor, mesh, [Shard(dim0)])其中torchrun --standalone --nnodes1 --nproc-per-node4 dtensor_example.py是官方推荐的启动方式elastic launcher如果使用mp.spawn则需先调用torch.distributed.init_process_group(backendnccl, world_sizeworld_size)初始化进程组。启动后每个进程运行同一份代码SPMD但各自持有big_tensor第 0 维的 1/4 分片。init_device_mesh构建多维设备网格init_device_mesh定义于 torch/distributed/device_mesh.py它根据device_type和mesh_shape创建一个 n 维数组布局的DeviceMeshn 为mesh_shape的长度。其完整签名如下def init_device_mesh( device_type: str, mesh_shape: tuple[int, ...], *, mesh_dim_names: tuple[str, ...] | None None, backend_override: dict[ int | str, str | C10dBackend.Options | tuple[str, C10dBackend.Options] ] | None None, ) - DeviceMesh从源码的 docstring 可以确认以下使用前提与限制SPMD 一致性init_device_mesh遵循 SPMD 编程模型集群中所有进程运行相同的 Python 程序因此mesh_shape必须在所有 rank 上保持一致否则可能导致挂起。自动初始化进程组如果当前没有找到进程组init_device_mesh会在后台自动初始化分布式通信所需的进程组。device_type 支持范围目前支持cpu、cuda/cuda-like、xpu不允许传入带 GPU 索引的设备类型如cuda:0。mesh_dim_names为每个网格维度命名长度必须与mesh_shape一致且每个名字必须唯一源码中会做重复名检查例如init_device_mesh(cuda, (2, 8), mesh_dim_names(dp, tp))可分别以dp、tp标识数据并行与张量并行维度。backend_override可对部分或全部网格维度的 ProcessGroup 进行后端覆盖key 可以是维度索引或维度名提供了mesh_dim_names时value 可以是后端名、C10dBackend.Options或两者的元组。动机为什么需要一个统一的张量级抽象当前扩展大规模分布式训练主要有三种方式数据并行Data Parallel、张量并行Tensor Parallel和流水线并行Pipeline Parallel。它们各自在独立维度上工作PyTorch 生态为此分别构建了 DDP、FSDP、ShardedTensor、PiPPy 等独立方案。当用户训练超大模型时往往需要将这三种并行组合使用即三维并行 3-D Parallelism但现有方案之间的互操作性不佳、使用门槛高——用户很难随意组合数据并行、张量并行与流水线并行。根本原因在于缺少一个能在不同并行策略之间搭建桥梁的公共抽象。理想场景是用户构建分布式程序就像在单机/单设备上编写程序一样无需关心集群内的分布式训练细节由框架自动决定如何拆分模型、如何在节点间跑流水线并行、如何在节点内跑数据并行和张量并行。要实现这一点就需要公共抽象来统一描述张量值的分布方式与分布式计算。DTensor 正是在这一背景下被引入作为下一代 ShardedTensor为分布式的存储与计算提供基础抽象并作为分布式程序转换的基本构建块可无缝构建张量并行、DDP、FSDP 等并行策略。价值主张DTensor 带来什么README 明确了 PyTorch DTensor 的三点核心价值统一的state_dict存取在进行 checkpoint 时提供统一的保存/加载方式即使存在复杂的张量存储分布策略如将张量并行与 FSDP 参数分片结合也能正确处理。eager 模式下的张量并行相比 ShardedTensorDTensor 允许灵活混用 sharding 与 replication从而在 eager 模式下直接表达张量并行。SPMD 编程模型的入口作为 SPMD 编程模型的入口点同时也是基于编译器的分布式训练的基础构建块。从目录结构看torch/distributed/tensor/内部还提供了_sharding_prop.py分片传播、_redistribute.py重分布、_dispatch.py算子分发以及_ops/内建 DTensor 算子实现等模块共同支撑上述能力。Placement 体系Shard / Replicate / Partial 的语义与限制DTensor 的核心是三种 Placement放置策略它们定义在 torch/distributed/tensor/placement_types.py 中分别是Shard、Replicate和PartialShard(dim)按维度切分Shard(dim)描述 DTensor 在张量的第dim维上、沿 DeviceMesh 的对应维度切分网格维度上的每个 rank 只持有全局张量的一份分片。源码明确说明其遵循torch.chunk(dim)语义——当张量维度不能被网格维度整除时网格维度上最后几个分片可能为空对应不均匀切分目前为实验性行为可能发生变化。因此Shard用于表达 row-wiseShard(0)或 col-wiseShard(1)切分等场景。Replicate()整张复制Replicate()描述 DTensor 在对应 DeviceMesh 维度上进行复制网格维度上的每个 rank 都持有全局张量的完整副本。源码中Replicate._make_replicate_tensor的实现表明复制通过从网格维度的首个坐标作为数据源进行广播broadcast完成若 rank 不属于该网格则返回空张量。Replicate常用于数据并行式的复制语义。Partial(reduce_op)待归约的局部结果Partial(reduce_op)描述的是一个待归约的 DTensor它在对应网格维度上并不持有完整值或单一分片而是持有部分和/局部结果需要通过指定的reduce_op归约后才能得到全局值。源码显示支持的reduce_op包括sum、avg、min、max其__repr__为Partial(sum)这样的形式字符串表示缩写为P(sum)。Partial常用于表达 matmul 等算子中间结果的分布状态。混用限制Partial 归约类型必须一致三种 Placement 一般可以自由组合但存在一个重要限制一个 DTensor 不能同时混用不同的Partial归约类型例如Partial(sum)与Partial(max)混用。原因是非线性归约如max与线性归约如sum不可交换不满足交换律在重分布redistribution过程中相对顺序在语义上至关重要。因此单个 DTensor 中的所有PartialPlacement 必须使用相同的归约操作# Valid: same partial type placements [Partial(sum), Partial(sum), Shard(0)] # OK # Invalid: mixed partial types - will raise ValueError placements [Partial(sum), Partial(max), Shard(0)] # Error!该限制同样体现在实现层面在 torch/distributed/tensor/_api.py 中distribute_tensor在分发前会调用assert_no_mixed_partial_types(placements)对 placements 做校验混用不同归约类型会直接抛出ValueError。DTensor 基础 API 实战构造、from_local 与 redistributetorch.distributed.tensor对外导出了DTensor、Shard、Replicate、Partial、distribute_tensor、distribute_module、init_device_mesh、DeviceMesh等公共 API见 torch/distributed/tensor/init.py。以下示例覆盖三类典型用法直接构造 DTensor表达纯切分、纯复制、切分复制等不同策略从本地torch.Tensor构造 DTensorDTensor.from_local对已有 DTensor 做重分片reshard迁移到新的 DTensor Layoutredistribute。# torchrun --standalone --nnodes1 --nproc-per-node4 dtensor_example.py import torch from torch.distributed.tensor import DTensor, Shard, Replicate, distribute_tensor, distribute_module, init_device_mesh # construct a device mesh with available devices (multi-host or single host) device_mesh init_device_mesh(cuda, (4,)) # if we want to do row-wise sharding rowwise_placement[Shard(0)] # if we want to do col-wise sharding colwise_placement[Shard(1)] big_tensor torch.randn(888, 12) # distributed tensor returned will be sharded across the dimension specified in placements rowwise_tensor distribute_tensor(big_tensor, device_meshdevice_mesh, placementsrowwise_placement) # if we want to do replication across a certain device list replica_placement [Replicate()] # distributed tensor will be replicated to all four GPUs. replica_tensor distribute_tensor(big_tensor, device_meshdevice_mesh, placementsreplica_placement) # if we want to distribute a tensor with both replication and sharding device_mesh init_device_mesh(cuda, (2, 2)) # replicate across the first dimension of device mesh, then sharding on the second dimension of device mesh spec[Replicate(), Shard(0)] partial_replica distribute_tensor(big_tensor, device_meshdevice_mesh, placementsspec) # create a DistributedTensor that shards on dim 0, from a local torch.Tensor local_tensor torch.randn((8, 8), requires_gradTrue) rowwise_tensor DTensor.from_local(local_tensor, device_mesh, rowwise_placement) # reshard the current row-wise tensor to a colwise tensor or replicate tensor colwise_tensor rowwise_tensor.redistribute(device_mesh, colwise_placement) replica_tensor colwise_tensor.redistribute(device_mesh, replica_placement)要点说明placements 与 mesh 维数必须一致源码中distribute_tensor会校验len(placements) device_mesh.ndim不一致时抛出ValueErrortorch/distributed/tensor/_api.py。placements 缺省行为若不指定placements默认将张量在 DeviceMesh 每个维度上复制[Replicate() for _ in range(device_mesh.ndim)]。leaf tensor 限制distribute_tensor用于分发 leaf 张量如nn.Parameter/buffer若传入非叶子张量会抛出RuntimeError。若需在 Autograd 计算中途构造 DTensor应改用DTensor.from_local。src_data_rankdistribute_tensor默认以每个网格维度的group_rank0作为逻辑/全局张量的数据源通过 scatter/broadcast 保持单设备语义显式传入None时则直接使用本地数据。redistributeredistribute(device_mesh, placements)将当前 DTensor 迁移到新的 Layout其底层实现位于 torch/distributed/tensor/_redistribute.py负责在三种 Placement 之间做转换如从 row-wise 切分转到列切分、或转成全复制。高层 APIdistribute_tensor 与 distribute_module用户可以直接使用 DTensor 张量构造器如distributed.ones/empty创建分布式张量但对于nn.Linear这类已持有torch.Tensor参数的现有模块如何让参数变为分布式DTensor 提供了两个高层 APIdef distribute_tensor(tensor: torch.Tensor, device_mesh: DeviceMeshNone, placements: List[Placement]None): distribute the tensor according to device_mesh and placements, tensor could be a meta tensor. def distribute_module( module: nn.Module, device_mesh: DeviceMeshNone, partition_fn: Callable[[str, nn.Module, DeviceMesh], ...]None, input_fn: Callable[...., None]None, output_fn: Callable[...., None]None, ): This function converts all module parameters to distributed tensor parameters according to the partition_fn specified. It could also control the input/output of the module by specifying the input_fn and output_fn. distribute_module 的三个回调机制从 torch/distributed/tensor/_api.py 的实现细节可以确认三个回调的分工与执行时机partition_fn(name, module, device_mesh)在运行时执行之前被调用用于切分模块的参数/buffer——把普通torch.Tensor参数替换为DTensor参数。它对每个子模块含模块自身各调用一次name是子模块的完整限定名顶层模块为即nn.Module.named_modules的返回。未被partition_fn转换的参数/buffer 之后会被复制到整个 DeviceMesh若partition_fn为None则所有参数/buffer 全部复制。input_fn(module, inputs, device_mesh)在运行时执行期间通过forward_pre_hook安装将模块输入转换为 DTensorinputs是forward位置参数构成的元组关键字参数不会传给input_fn。output_fn(module, outputs, device_mesh)在运行时执行期间通过forward_hook安装将模块输出通常是 DTensor可能为任意嵌套结构转换回普通torch.Tensor典型实现是调用DTensor.to_local或DTensor.full_tensor或重分布到其他 placements。注意input_fn/output_fn早期版本采用双参数形式(inputs, device_mesh)该形式已废弃请使用上面文档化的三参数形式。模块级张量并行完整示例下面的例子将自定义模块的nn.Linear参数按列切分Shard(0)实现模块级张量并行import torch.nn as nn from torch.distributed.tensor import Shard, distribute_tensor, distribute_module, init_device_mesh class MyModule(nn.Module): def __init__(self) - None: super().__init__() self.fc1 nn.Linear(8, 8) self.fc2 nn.Linear(8, 8) self.relu nn.ReLU() def forward(self, input): return self.relu(self.fc1(input) self.fc2(input)) mesh init_device_mesh(cuda, (4,)) def shard_params(mod_name, mod, mesh): col_linear_placement [Shard(0)] # shard fc1 and fc2 if isinstance(mod, nn.Linear): for name, param in mod.named_parameters(): dist_param nn.Parameter( distribute_tensor(param, mesh, col_linear_placement) ) mod.register_parameter(name, dist_param) sharded_module distribute_module(MyModule(), mesh, partition_fnshard_params)结合input_fn/output_fn还可以进一步控制运行时输入输出例如用to_replicated把普通输入张量分发为复制 DTensor 再进入forward用to_local内部调用outputs.full_tensor()把 DTensor 输出还原为普通张量。torch/distributed/tensor/examples/目录下的torchrec_sharding_example.py、convnext_example.py等示例脚本可作为继续深入的参考。编译器与 DTensor从 eager 优化到编译式分布式训练DTensor 为张量并行等场景提供了高效方案但以数据并行的方式使用 DTensor 的复制Replication时性能可能明显低于 DDP/FSDP 等现有方案。原因在于DDP/FSDP 拥有整个模型架构的全局视图可以针对数据并行做专门的优化如集合通信融合、计算与通信重叠等而 DTensor 作为类张量对象只能在单个张量操作内部做优化。为此官方正在 DTensor 之上探索基于编译器的方案从用户程序中提取图信息以暴露更多性能优化机会从而提升基于 DTensor 的数据并行训练效率。这也是 DTensor 被定位为编译器式分布式训练基础构建块的原因。相关技术脉络与后续演进DTensor 的设计主要受三个使用单一分布式张量概念统一表达复制与切分的工作启发这些工作都支持用统一的 SPMD 编程模型构建分布式训练程序GSPMDJAX/TensorFlow 分布式训练的基础组件借助 XLA 编译器实现大规模高效训练。它在一个张量内定义了三种切分策略——tiled分片、replicated复制、partially tiled部分分片来表达切分与复制GSPMD Partitioner 的核心是借助 XLA 编译器做分片传播与基于编译器的融合等高级优化。OneFlow GlobalTensorOneFlow 构建了自己的 GlobalTensor 概念是 GSPMD 切分的一种变体。它也定义了三类张量但与 GSPMD 略有不同split切分、broadcast广播、partial sum部分和——不使用 partially tiled而是用 partial sum 概念来表达值的归约。TensorFlow DTensorTensorFlow 同步分布式训练的扩展支持在多维 mesh如设备网络上同时做切分与复制并实现了 MLIR passes 来完成传播与算子实现。此外还有多个将张量切分嵌入系统的高阶研究方向例如 Megatron-LM 面向 Transformer 模型的张量并行、DeepSpeed 在张量切分之上结合多种优化技术训练大规模模型。本项目中的 DTensor 原型Prototype同样处于早期反馈收集阶段官方已发布相关 RFC 并在 dev-discuss 论坛同步征求反馈完整设计文档可结合torch/distributed/tensor/源码目录继续研读——尤其是 placement_types.pyPlacement 定义与张量切分/广播/归约实现、_api.pydistribute_tensor/distribute_module实现、_sharding_prop.py分片传播规则以及_ops/目录内建算子的分片规则与策略它们共同构成了 DTensor 从声明式 Placement 到实际分布式计算的核心链路。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
